package handlers

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"

	"git.elijahkuntz.com/Elijah/drive/config"
	"git.elijahkuntz.com/Elijah/drive/database"
	"git.elijahkuntz.com/Elijah/drive/middleware"
	"github.com/gofiber/fiber/v2"
	"github.com/pquerna/otp/totp"
	"golang.org/x/crypto/argon2"
)

type AuthHandler struct {
	DB     *database.DB
	Config *config.Config
	Guard  *middleware.BruteForceGuard
}

type SetupRequest struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

type LoginRequest struct {
	Username string `json:"username"`
	Password string `json:"password"`
	TOTPCode string `json:"totp_code,omitempty"`
}

// Setup handles the first-run account creation wizard.
// POST /api/setup
func (h *AuthHandler) Setup(c *fiber.Ctx) error {
	// Check if setup is already complete
	complete, err := h.DB.IsSetupComplete()
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
	}
	if complete {
		return c.Status(fiber.StatusConflict).JSON(fiber.Map{"error": "setup already completed"})
	}

	var req SetupRequest
	if err := c.BodyParser(&req); err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request body"})
	}

	if len(req.Username) < 3 {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "username must be at least 3 characters"})
	}
	if len(req.Password) < 8 {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "password must be at least 8 characters"})
	}

	hash, err := hashPassword(req.Password)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to hash password"})
	}

	_, err = h.DB.Exec(
		"INSERT INTO users (username, password_hash) VALUES (?, ?)",
		req.Username, hash,
	)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to create user"})
	}

	h.DB.AddAuditLog("setup_complete", fmt.Sprintf("User '%s' created", req.Username), c.IP())

	return c.Status(fiber.StatusCreated).JSON(fiber.Map{
		"message": "setup complete, please log in",
	})
}

// CheckSetup returns whether initial setup has been completed.
// GET /api/setup/status
func (h *AuthHandler) CheckSetup(c *fiber.Ctx) error {
	complete, err := h.DB.IsSetupComplete()
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
	}
	return c.JSON(fiber.Map{"setup_complete": complete})
}

// Login authenticates a user and returns JWT tokens.
// POST /api/auth/login
func (h *AuthHandler) Login(c *fiber.Ctx) error {
	var req LoginRequest
	if err := c.BodyParser(&req); err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request body"})
	}

	var id int
	var passwordHash, totpSecret string
	var totpEnabled bool

	err := h.DB.QueryRow(
		"SELECT id, password_hash, totp_secret, totp_enabled FROM users WHERE username = ?",
		req.Username,
	).Scan(&id, &passwordHash, &totpSecret, &totpEnabled)

	if err != nil {
		h.Guard.RecordFailure(c.IP())
		h.DB.AddAuditLog("login_failed", fmt.Sprintf("Unknown user '%s'", req.Username), c.IP())
		return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid credentials"})
	}

	if !verifyPassword(req.Password, passwordHash) {
		h.Guard.RecordFailure(c.IP())
		h.DB.AddAuditLog("login_failed", fmt.Sprintf("Bad password for '%s'", req.Username), c.IP())
		return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid credentials"})
	}

	// Check TOTP if enabled
	if totpEnabled {
		if req.TOTPCode == "" {
			return c.Status(fiber.StatusPreconditionRequired).JSON(fiber.Map{
				"error":      "2fa_required",
				"message":    "Two-factor authentication code required",
				"totp_required": true,
			})
		}
		if !totp.Validate(req.TOTPCode, totpSecret) {
			h.Guard.RecordFailure(c.IP())
			h.DB.AddAuditLog("login_failed", fmt.Sprintf("Bad TOTP for '%s'", req.Username), c.IP())
			return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid 2FA code"})
		}
	}

	// Generate tokens
	accessToken, err := middleware.GenerateAccessToken(id, req.Username, h.Config.JWTSecret, h.Config.JWTExpirySecs)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate token"})
	}

	refreshToken, err := middleware.GenerateRefreshToken(id, req.Username, h.Config.JWTSecret)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate refresh token"})
	}

	h.Guard.RecordSuccess(c.IP())
	h.DB.AddAuditLog("login_success", fmt.Sprintf("User '%s' logged in", req.Username), c.IP())

	return c.JSON(fiber.Map{
		"access_token":  accessToken,
		"refresh_token": refreshToken,
		"expires_in":    h.Config.JWTExpirySecs,
	})
}

// Refresh issues a new access token from a valid refresh token.
// POST /api/auth/refresh
func (h *AuthHandler) Refresh(c *fiber.Ctx) error {
	// The auth middleware already validated the token
	userID := c.Locals("userID")
	username := c.Locals("username")

	accessToken, err := middleware.GenerateAccessToken(
		int(userID.(float64)),
		username.(string),
		h.Config.JWTSecret,
		h.Config.JWTExpirySecs,
	)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate token"})
	}

	return c.JSON(fiber.Map{
		"access_token": accessToken,
		"expires_in":   h.Config.JWTExpirySecs,
	})
}

// Enable2FA generates a TOTP secret and returns the provisioning URI.
// POST /api/auth/2fa/enable
func (h *AuthHandler) Enable2FA(c *fiber.Ctx) error {
	username := c.Locals("username").(string)

	key, err := totp.Generate(totp.GenerateOpts{
		Issuer:      "Drive",
		AccountName: username,
	})
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to generate TOTP secret"})
	}

	// Store the secret but don't enable 2FA until confirmed
	_, err = h.DB.Exec("UPDATE users SET totp_secret = ? WHERE username = ?", key.Secret(), username)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to save TOTP secret"})
	}

	return c.JSON(fiber.Map{
		"secret": key.Secret(),
		"url":    key.URL(),
	})
}

// Confirm2FA verifies a TOTP code and enables 2FA.
// POST /api/auth/2fa/confirm
func (h *AuthHandler) Confirm2FA(c *fiber.Ctx) error {
	username := c.Locals("username").(string)

	var body struct {
		Code string `json:"code"`
	}
	if err := c.BodyParser(&body); err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request"})
	}

	var secret string
	err := h.DB.QueryRow("SELECT totp_secret FROM users WHERE username = ?", username).Scan(&secret)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
	}

	if !totp.Validate(body.Code, secret) {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid TOTP code"})
	}

	_, err = h.DB.Exec("UPDATE users SET totp_enabled = 1 WHERE username = ?", username)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to enable 2FA"})
	}

	h.DB.AddAuditLog("2fa_enabled", fmt.Sprintf("User '%s' enabled 2FA", username), c.IP())

	return c.JSON(fiber.Map{"message": "2FA enabled successfully"})
}

// Disable2FA turns off two-factor authentication.
// POST /api/auth/2fa/disable
func (h *AuthHandler) Disable2FA(c *fiber.Ctx) error {
	username := c.Locals("username").(string)

	_, err := h.DB.Exec("UPDATE users SET totp_enabled = 0, totp_secret = '' WHERE username = ?", username)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to disable 2FA"})
	}

	h.DB.AddAuditLog("2fa_disabled", fmt.Sprintf("User '%s' disabled 2FA", username), c.IP())

	return c.JSON(fiber.Map{"message": "2FA disabled successfully"})
}

// ChangePassword allows an authenticated user to change their password.
// PUT /api/auth/password
func (h *AuthHandler) ChangePassword(c *fiber.Ctx) error {
	username := c.Locals("username").(string)

	var req struct {
		OldPassword string `json:"old_password"`
		NewPassword string `json:"new_password"`
	}
	if err := c.BodyParser(&req); err != nil {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "invalid request body"})
	}

	if len(req.NewPassword) < 8 {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "new password must be at least 8 characters"})
	}

	var currentHash string
	if err := h.DB.QueryRow("SELECT password_hash FROM users WHERE username = ?", username).Scan(&currentHash); err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "database error"})
	}

	if !verifyPassword(req.OldPassword, currentHash) {
		h.Guard.RecordFailure(c.IP())
		return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "incorrect current password"})
	}

	newHash, err := hashPassword(req.NewPassword)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to hash new password"})
	}

	if _, err := h.DB.Exec("UPDATE users SET password_hash = ? WHERE username = ?", newHash, username); err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"error": "failed to update password"})
	}

	h.DB.AddAuditLog("password_changed", fmt.Sprintf("User '%s' changed their password", username), c.IP())
	return c.JSON(fiber.Map{"message": "password changed successfully"})
}

// --- Password Hashing with Argon2 ---

func hashPassword(password string) (string, error) {
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return "", err
	}

	hash := argon2.IDKey([]byte(password), salt, 2, 64*1024, 4, 32)
	encoded := fmt.Sprintf("%s$%s",
		base64.RawStdEncoding.EncodeToString(salt),
		base64.RawStdEncoding.EncodeToString(hash),
	)
	return encoded, nil
}

func verifyPassword(password, encoded string) bool {
	parts := splitN(encoded, "$", 2)
	if len(parts) != 2 {
		return false
	}

	salt, err := base64.RawStdEncoding.DecodeString(parts[0])
	if err != nil {
		return false
	}

	expectedHash, err := base64.RawStdEncoding.DecodeString(parts[1])
	if err != nil {
		return false
	}

	hash := argon2.IDKey([]byte(password), salt, 2, 64*1024, 4, 32)

	if len(hash) != len(expectedHash) {
		return false
	}

	// Constant-time comparison to prevent timing attacks
	var diff byte
	for i := range hash {
		diff |= hash[i] ^ expectedHash[i]
	}
	return diff == 0
}

func splitN(s, sep string, n int) []string {
	result := make([]string, 0, n)
	for i := 0; i < n-1; i++ {
		idx := indexOf(s, sep)
		if idx < 0 {
			break
		}
		result = append(result, s[:idx])
		s = s[idx+len(sep):]
	}
	result = append(result, s)
	return result
}

func indexOf(s, sub string) int {
	for i := 0; i <= len(s)-len(sub); i++ {
		if s[i:i+len(sub)] == sub {
			return i
		}
	}
	return -1
}