package handlers

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"golang.org/x/crypto/bcrypt"

	"switchboard-core/auth"
	"switchboard-core/config"
	"switchboard-core/crypto"
	"switchboard-core/models"
	"switchboard-core/store"
)

// Claims represents the JWT payload.
//
// bcryptCost is shared across auth.go, settings.go, admin.go
const bcryptCost = 12

type Claims struct {
	UserID string `json:"user_id"`
	Email  string `json:"email"`
	Role   string `json:"role"`
	jwt.RegisteredClaims
}

type AuthHandler struct {
	cfg      *config.Config
	stores   store.Stores
	uekCache *crypto.UEKCache
	provider auth.Provider
}

func NewAuthHandler(cfg *config.Config, s store.Stores, uekCache *crypto.UEKCache, provider auth.Provider) *AuthHandler {
	return &AuthHandler{cfg: cfg, stores: s, uekCache: uekCache, provider: provider}
}

func (h *AuthHandler) Register(c *gin.Context) {
	if !h.provider.SupportsRegistration() {
		c.JSON(http.StatusMethodNotAllowed, gin.H{"error": "registration not supported for this auth mode"})
		return
	}

	result, err := h.provider.Register(c, h.stores)
	if err != nil {
		c.JSON(auth.ErrorToHTTPStatus(err), gin.H{"error": err.Error()})
		return
	}

	if result.VaultHint != "" {
		if err := h.initVault(c.Request.Context(), result.User.ID, result.VaultHint); err != nil {
			log.Printf("⚠ Failed to init vault for user %s: %v", result.User.ID, err)
		}
	}

	if !result.User.IsActive {
		c.JSON(http.StatusCreated, gin.H{
			"message": "Account created but requires admin approval",
			"user_id": result.User.ID,
		})
		return
	}

	tokens, err := h.generateTokens(result.User)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate tokens"})
		return
	}

	c.JSON(http.StatusCreated, tokens)
}

func (h *AuthHandler) Login(c *gin.Context) {
	result, err := h.provider.Authenticate(c, h.stores)
	if err != nil {
		c.JSON(auth.ErrorToHTTPStatus(err), gin.H{"error": err.Error()})
		return
	}

	if result.VaultHint != "" {
		h.unlockVault(c.Request.Context(), result.User, result.VaultHint)
	}

	h.stores.Users.UpdateLastLogin(c.Request.Context(), result.User.ID)

	tokens, err := h.generateTokens(result.User)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate tokens"})
		return
	}

	c.JSON(http.StatusOK, tokens)
}

func (h *AuthHandler) Refresh(c *gin.Context) {
	var req struct {
		RefreshToken string `json:"refresh_token" binding:"required"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	tokenHash := hashToken(req.RefreshToken)
	userID, err := h.stores.Users.GetRefreshToken(c.Request.Context(), tokenHash)
	if err != nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid refresh token"})
		return
	}

	// Revoke the used token (rotate)
	h.stores.Users.RevokeRefreshToken(c.Request.Context(), tokenHash)

	user, err := h.stores.Users.GetByID(c.Request.Context(), userID)
	if err != nil || !user.IsActive {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "user not found or inactive"})
		return
	}

	tokens, err := h.generateTokens(user)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate tokens"})
		return
	}

	c.JSON(http.StatusOK, tokens)
}

func (h *AuthHandler) Logout(c *gin.Context) {
	var req struct {
		RefreshToken string `json:"refresh_token"`
	}
	c.ShouldBindJSON(&req)

	if req.RefreshToken != "" {
		tokenHash := hashToken(req.RefreshToken)
		h.stores.Users.RevokeRefreshToken(c.Request.Context(), tokenHash)
	}

	// Evict UEK from session cache
	if userID, exists := c.Get("user_id"); exists {
		h.uekCache.Evict(userID.(string))
	}

	// Clear the sb_token cookie so SSR middleware doesn't trust a stale token
	c.SetCookie("sb_token", "", -1, "/", "", false, false)

	c.JSON(http.StatusOK, gin.H{"message": "logged out"})
}

// ── OIDC Flow ──────────────────────────────

// OIDCLogin initiates the authorization code flow by redirecting to the IdP.
// GET /api/v1/auth/oidc/login
func (h *AuthHandler) OIDCLogin(c *gin.Context) {
	oidcProv, ok := h.provider.(*auth.OIDCProvider)
	if !ok {
		c.JSON(http.StatusBadRequest, gin.H{"error": "OIDC not configured"})
		return
	}

	state := uuid.New().String()
	nonce := uuid.New().String()

	// Determine redirect URI
	redirectURI := h.cfg.OIDCRedirectURL
	if redirectURI == "" {
		// Auto-derive from request
		scheme := "https"
		if c.Request.TLS == nil && c.GetHeader("X-Forwarded-Proto") != "https" {
			scheme = "http"
		}
		redirectURI = scheme + "://" + c.Request.Host + h.cfg.BasePath + "/api/v1/auth/oidc/callback"
	}

	// Store state for callback verification
	if err := h.stores.GlobalConfig.SaveOIDCState(c.Request.Context(), state, nonce, c.Query("redirect")); err != nil {
		log.Printf("[auth/oidc] warn: could not store state: %v", err)
	}

	authURL := oidcProv.AuthorizationURL(state, nonce, redirectURI)
	c.Redirect(http.StatusFound, authURL)
}

// OIDCCallback handles the IdP redirect after successful authentication.
// GET /api/v1/auth/oidc/callback?code=...&state=...
func (h *AuthHandler) OIDCCallback(c *gin.Context) {
	oidcProv, ok := h.provider.(*auth.OIDCProvider)
	if !ok {
		c.JSON(http.StatusBadRequest, gin.H{"error": "OIDC not configured"})
		return
	}

	code := c.Query("code")
	state := c.Query("state")

	// Check for IdP error response (Keycloak returns ?error=...&error_description=...)
	if errCode := c.Query("error"); errCode != "" {
		errDesc := c.Query("error_description")
		log.Printf("[auth/oidc] IdP returned error: %s — %s", errCode, errDesc)
		c.JSON(http.StatusUnauthorized, gin.H{"error": "IdP error: " + errCode, "description": errDesc})
		return
	}

	if code == "" || state == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "missing code or state"})
		return
	}

	// Verify state (nonce + redirectTo retrieved but not yet validated — TODO)
	_, _, err := h.stores.GlobalConfig.ConsumeOIDCState(c.Request.Context(), state)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid or expired state"})
		return
	}

	// Clean up stale states (> 10 minutes old)
	_ = h.stores.GlobalConfig.CleanupOIDCState(c.Request.Context())

	// Determine redirect URI (must match what was sent in the login request)
	redirectURI := h.cfg.OIDCRedirectURL
	if redirectURI == "" {
		scheme := "https"
		if c.Request.TLS == nil && c.GetHeader("X-Forwarded-Proto") != "https" {
			scheme = "http"
		}
		redirectURI = scheme + "://" + c.Request.Host + h.cfg.BasePath + "/api/v1/auth/oidc/callback"
	}

	// Exchange code for tokens
	tokenResp, err := oidcProv.ExchangeCode(c.Request.Context(), code, redirectURI)
	if err != nil {
		log.Printf("[auth/oidc] code exchange failed: %v", err)
		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
		return
	}

	// Use the ID token (or access token) to authenticate via the provider
	// Temporarily set the Authorization header so Authenticate() can read it
	tokenToValidate := tokenResp.IDToken
	if tokenToValidate == "" {
		tokenToValidate = tokenResp.AccessToken
	}
	c.Request.Header.Set("Authorization", "Bearer "+tokenToValidate)

	result, err := oidcProv.Authenticate(c, h.stores)
	if err != nil {
		log.Printf("[auth/oidc] authenticate after code exchange failed: %v", err)
		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
		return
	}

	// Issue internal JWT
	tokens, err := h.generateTokens(result.User)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate tokens"})
		return
	}

	h.stores.Users.UpdateLastLogin(c.Request.Context(), result.User.ID)

	// For browser flow: encode tokens into a fragment that the login page JS
	// can pick up, store in localStorage, and redirect to the app.
	// This avoids httpOnly cookie issues — JS needs the token in localStorage.
	accessToken, _ := tokens["access_token"].(string)
	refreshToken, _ := tokens["refresh_token"].(string)
	userJSON, _ := json.Marshal(tokens["user"])

	// Set page-auth cookie too (for SSR middleware)
	c.SetCookie("sb_token", accessToken, 900, "/", "", false, false)

	// Base64-encode the token payload for the fragment
	payload := fmt.Sprintf(`{"access_token":"%s","refresh_token":"%s","user":%s}`,
		accessToken, refreshToken, string(userJSON))
	encoded := base64.StdEncoding.EncodeToString([]byte(payload))

	dest := h.cfg.BasePath + "/login#oidc_result=" + encoded
	c.Redirect(http.StatusFound, dest)
}

func (h *AuthHandler) generateTokens(user *models.User) (gin.H, error) {
	// Access token (15 min)
	accessClaims := Claims{
		UserID: user.ID,
		Email:  user.Email,
		Role:   user.Role,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(15 * time.Minute)),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
			ID:        uuid.New().String(),
		},
	}
	accessToken := jwt.NewWithClaims(jwt.SigningMethodHS256, accessClaims)
	accessString, err := accessToken.SignedString([]byte(h.cfg.JWTSecret))
	if err != nil {
		return nil, err
	}

	// Refresh token (7 days)
	refreshRaw := uuid.New().String()
	refreshHash := hashToken(refreshRaw)
	expiresAt := time.Now().Add(7 * 24 * time.Hour)

	if err := h.stores.Users.CreateRefreshToken(context.Background(), user.ID, refreshHash, expiresAt); err != nil {
		log.Printf("warn: failed to store refresh token: %v", err)
	}

	return gin.H{
		"access_token":  accessString,
		"refresh_token": refreshRaw,
		"token_type":    "Bearer",
		"expires_in":    900,
		"user": gin.H{
			"id":           user.ID,
			"username":     user.Username,
			"email":        user.Email,
			"display_name": user.DisplayName,
			"role":         user.Role,
			"handle":       user.Handle,
			"auth_source":  user.AuthSource,
		},
	}, nil
}

func hashToken(token string) string {
	h := sha256.Sum256([]byte(token))
	return hex.EncodeToString(h[:])
}

// ── Vault Lifecycle ─────────────────────────

// DestroyVaultDB nullifies a user's vault columns and deletes personal
// provider configs. This is the shared DB-level operation used by:
//   - unlockVault (stale seal recovery)
//   - BootstrapAdmin / SeedUsers (password rotation at startup)
//   - AdminHandler.destroyVault (admin-initiated reset)
//
// Does NOT evict from UEK cache or write audit logs — callers handle that.
// v0.29.0: accepts stores instead of using database.DB directly.
func DestroyVaultDB(ctx context.Context, stores store.Stores, userID string) (providersDeleted int64) {
	if err := stores.Users.ClearVaultKeys(ctx, userID); err != nil {
		log.Printf("⚠ DestroyVaultDB: failed to clear vault columns for user %s: %v", userID, err)
	}

	rows, err := stores.Providers.DeletePersonalByOwner(ctx, userID)
	if err != nil {
		log.Printf("⚠ DestroyVaultDB: failed to delete personal providers for user %s: %v", userID, err)
		return 0
	}
	return rows
}

// ProbeAndRepairVault checks whether a user's vault can be unlocked with
// the given password. If the vault doesn't exist (vault_set=false), this is
// a no-op. If the vault exists but the seal is stale (encrypted_uek was
// wrapped with a different password), the vault and personal providers are
// destroyed so initVault fires cleanly on next login.
//
// When uekCache is non-nil and the vault unlocks successfully, the UEK is
// cached so that BYOK operations work immediately without requiring a fresh
// login. This is important for bootstrapped/seeded users whose browser
// sessions survive server restarts.
//
// Used by BootstrapAdmin and SeedUsers where the password is known at startup.
// v0.29.0: accepts stores instead of using database.DB directly.
// v0.30.2: accepts optional uekCache to pre-warm vault on restart.
func ProbeAndRepairVault(ctx context.Context, stores store.Stores, userID, password string, uekCache ...*crypto.UEKCache) {
	vaultSet, encryptedUEK, salt, nonce, err := stores.Users.GetVaultKeys(ctx, userID)
	if err != nil || !vaultSet {
		return // no vault to probe
	}

	pdk := crypto.DeriveKeyFromPassword(password, salt)
	uek, err := crypto.UnwrapUEK(encryptedUEK, nonce, pdk)
	if err == nil {
		// Vault seal matches current password — cache UEK if cache provided
		if len(uekCache) > 0 && uekCache[0] != nil {
			uekCache[0].Store(userID, uek)
		}
		return
	}

	// Stale seal: password has actually changed since the vault was sealed
	deleted := DestroyVaultDB(ctx, stores, userID)
	if deleted > 0 {
		log.Printf("  🔐 Vault stale-seal repair for user %s: cleared %d personal provider(s)", userID, deleted)
	} else {
		log.Printf("  🔐 Vault stale-seal repair for user %s: vault columns cleared", userID)
	}
}

// initVault generates a UEK, wraps it with the user's password, and stores
// the encrypted UEK + salt + nonce on the user record. Called on registration
// and on first login for pre-migration users.
func (h *AuthHandler) initVault(ctx context.Context, userID, password string) error {
	if h.uekCache == nil {
		return nil // Vault not configured (e.g. tests)
	}

	uek, err := crypto.GenerateUEK()
	if err != nil {
		return err
	}

	salt, err := crypto.GenerateSalt()
	if err != nil {
		return err
	}

	pdk := crypto.DeriveKeyFromPassword(password, salt)
	encryptedUEK, nonce, err := crypto.WrapUEK(uek, pdk)
	if err != nil {
		return err
	}

	if err := h.stores.Users.InitVaultKeys(ctx, userID, encryptedUEK, salt, nonce); err != nil {
		return err
	}

	// Cache the UEK for the session
	h.uekCache.Store(userID, uek)
	log.Printf("  🔐 Vault initialized for user %s", userID)
	return nil
}

// unlockVault unwraps the UEK on login and caches it. If the user doesn't
// have a vault yet (pre-migration account), initializes one.
func (h *AuthHandler) unlockVault(ctx context.Context, user *models.User, password string) {
	if h.uekCache == nil {
		return
	}

	vaultSet, encryptedUEK, salt, nonce, err := h.stores.Users.GetVaultKeys(ctx, user.ID)
	if err != nil {
		log.Printf("⚠ Vault query failed for user %s: %v", user.ID, err)
		return
	}

	if !vaultSet {
		// Pre-migration user: initialize vault on first login
		if err := h.initVault(ctx, user.ID, password); err != nil {
			log.Printf("⚠ Vault init on login failed for user %s: %v", user.ID, err)
		}
		return
	}

	// Derive PDK from password and unwrap UEK
	pdk := crypto.DeriveKeyFromPassword(password, salt)
	uek, err := crypto.UnwrapUEK(encryptedUEK, nonce, pdk)
	if err != nil {
		// Stale seal: encrypted_uek was wrapped with a different password
		deleted := DestroyVaultDB(ctx, h.stores, user.ID)
		if deleted > 0 {
			log.Printf("⚠ Vault stale-seal recovery for user %s: cleared %d personal provider(s)", user.ID, deleted)
		} else {
			log.Printf("⚠ Vault stale-seal recovery for user %s: no personal providers to clear", user.ID)
		}

		if err := h.initVault(ctx, user.ID, password); err != nil {
			log.Printf("⚠ Vault re-init failed for user %s: %v", user.ID, err)
		}
		return
	}

	h.uekCache.Store(user.ID, uek)
}

// BootstrapAdmin creates/updates the admin user from env vars (K8s secret).
// When uekCache is provided, the admin's vault is pre-warmed so BYOK
// operations work immediately without requiring a fresh login.
func BootstrapAdmin(cfg *config.Config, s store.Stores, uekCache ...*crypto.UEKCache) {
	if cfg.AdminUsername == "" || cfg.AdminPassword == "" {
		return
	}

	ctx := context.Background()
	hash, err := bcrypt.GenerateFromPassword([]byte(cfg.AdminPassword), bcryptCost)
	if err != nil {
		log.Printf("⚠ Failed to hash admin password: %v", err)
		return
	}

	existing, _ := s.Users.GetByUsername(ctx, cfg.AdminUsername)
	if existing != nil {
		// Update password and ensure admin role
		s.Users.Update(ctx, existing.ID, map[string]interface{}{
			"password_hash": string(hash),
			"role":          models.UserRoleAdmin,
			"is_active":     true,
		})
		// If the actual password changed, the vault seal is stale. Probe it
		// with the current password and destroy only if unwrap fails.
		// Ensure handle exists (backfill for pre-v0.24.0 users)
		if existing.Handle == "" {
			handle := auth.UniqueHandle(ctx, s.Users, models.HandleFromName(cfg.AdminUsername))
			s.Users.Update(ctx, existing.ID, map[string]interface{}{"handle": handle})
		}
		var cache *crypto.UEKCache
		if len(uekCache) > 0 {
			cache = uekCache[0]
		}
		ProbeAndRepairVault(ctx, s, existing.ID, cfg.AdminPassword, cache)
		log.Printf("  ✅ Admin user '%s' updated", cfg.AdminUsername)
		return
	}

	email := cfg.AdminEmail
	if email == "" {
		email = cfg.AdminUsername + "@switchboard.local"
	}

	handle := auth.UniqueHandle(ctx, s.Users, models.HandleFromName(cfg.AdminUsername))
	user := &models.User{
		Username:     strings.ToLower(cfg.AdminUsername),
		Email:        strings.ToLower(email),
		PasswordHash: string(hash),
		Role:         models.UserRoleAdmin,
		IsActive:     true,
		AuthSource:   "builtin",
		Handle:       handle,
	}

	if err := s.Users.Create(ctx, user); err != nil {
		log.Printf("⚠ Failed to create admin user: %v", err)
		return
	}
	log.Printf("  ✅ Admin user '%s' created", cfg.AdminUsername)
}

// SeedUsers creates or updates users from the SEED_USERS env var.
// Format: "user:pass:role,user2:pass2:role2" where role is "admin" or "user".
// Upsert: existing users get their password and role refreshed on every restart.
// Gated to non-production environments.
func SeedUsers(cfg *config.Config, s store.Stores, uekCache ...*crypto.UEKCache) {
	if cfg.SeedUsers == "" {
		log.Printf("  ℹ SEED_USERS not set, skipping")
		return
	}

	if cfg.Environment == "production" {
		log.Printf("⚠ SEED_USERS ignored in production environment")
		return
	}

	ctx := context.Background()
	entries := strings.Split(cfg.SeedUsers, ",")
	log.Printf("  🌱 SEED_USERS: %d entries to process", len(entries))

	for _, entry := range entries {
		entry = strings.TrimSpace(entry)
		if entry == "" {
			continue
		}

		parts := strings.SplitN(entry, ":", 3)
		if len(parts) < 2 {
			log.Printf("⚠ Seed user skipped (bad format, want user:pass[:role]): %q", entry)
			continue
		}

		username := strings.ToLower(strings.TrimSpace(parts[0]))
		password := strings.TrimSpace(parts[1])
		role := models.UserRoleUser
		if len(parts) >= 3 && strings.TrimSpace(parts[2]) == "admin" {
			role = models.UserRoleAdmin
		}

		if username == "" || password == "" {
			log.Printf("⚠ Seed user skipped (empty username or password)")
			continue
		}

		hash, err := bcrypt.GenerateFromPassword([]byte(password), bcryptCost)
		if err != nil {
			log.Printf("⚠ Seed user '%s' skipped (hash error): %v", username, err)
			continue
		}

		// Upsert: update password+role if user exists, create if not
		existing, _ := s.Users.GetByUsername(ctx, username)
		if existing != nil {
			s.Users.Update(ctx, existing.ID, map[string]interface{}{
				"password_hash": string(hash),
				"role":          role,
				"is_active":     true,
			})
			// Probe vault with current password — only destroys if seal is stale
			if existing.Handle == "" {
				handle := auth.UniqueHandle(ctx, s.Users, models.HandleFromName(username))
				s.Users.Update(ctx, existing.ID, map[string]interface{}{"handle": handle})
			}
			var cache *crypto.UEKCache
			if len(uekCache) > 0 {
				cache = uekCache[0]
			}
			ProbeAndRepairVault(ctx, s, existing.ID, password, cache)
			log.Printf("  🌱 Seed user '%s' updated (role=%s)", username, role)
			continue
		}

		handle := auth.UniqueHandle(ctx, s.Users, models.HandleFromName(username))
		user := &models.User{
			Username:     username,
			Email:        username + "@switchboard.local",
			PasswordHash: string(hash),
			Role:         role,
			IsActive:     true,
			AuthSource:   "builtin",
			Handle:       handle,
		}

		if err := s.Users.Create(ctx, user); err != nil {
			log.Printf("⚠ Seed user '%s' failed: %v", username, err)
			continue
		}
		log.Printf("  🌱 Seed user '%s' created (role=%s, active=true)", username, role)
	}
}

// IsRegistrationEnabled checks the platform policy.
func IsRegistrationEnabled(s store.Stores) bool {
	val, _ := s.Policies.GetBool(context.Background(), "allow_registration")
	return val
}