core/server/handlers/auth.go

package handlers

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"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"

	"git.gobha.me/xcaliber/chat-switchboard/config"
	"git.gobha.me/xcaliber/chat-switchboard/crypto"
	"git.gobha.me/xcaliber/chat-switchboard/database"
	"git.gobha.me/xcaliber/chat-switchboard/models"
	"git.gobha.me/xcaliber/chat-switchboard/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
}

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

func (h *AuthHandler) Register(c *gin.Context) {
	var req struct {
		Username string `json:"username" binding:"required"`
		Email    string `json:"email" binding:"required"`
		Password string `json:"password" binding:"required,min=8"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	// Check registration policy
	allowed, _ := h.stores.Policies.GetBool(c.Request.Context(), "allow_registration")
	if !allowed {
		c.JSON(http.StatusForbidden, gin.H{"error": "registration is disabled"})
		return
	}

	// Check duplicate
	if existing, _ := h.stores.Users.GetByUsername(c.Request.Context(), req.Username); existing != nil {
		c.JSON(http.StatusConflict, gin.H{"error": "username already taken"})
		return
	}
	if existing, _ := h.stores.Users.GetByEmail(c.Request.Context(), req.Email); existing != nil {
		c.JSON(http.StatusConflict, gin.H{"error": "email already registered"})
		return
	}

	hash, err := bcrypt.GenerateFromPassword([]byte(req.Password), bcryptCost)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to hash password"})
		return
	}

	// Check if user should be active by default
	defaultActive, _ := h.stores.Policies.GetBool(c.Request.Context(), "default_user_active")

	user := &models.User{
		Username:     strings.ToLower(req.Username),
		Email:        strings.ToLower(req.Email),
		PasswordHash: string(hash),
		Role:         models.UserRoleUser,
		IsActive:     defaultActive,
	}

	if err := h.stores.Users.Create(c.Request.Context(), user); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create user"})
		return
	}

	// Generate and store the User Encryption Key (vault)
	if err := h.initVault(c.Request.Context(), user.ID, req.Password); err != nil {
		log.Printf("⚠ Failed to init vault for user %s: %v", user.ID, err)
		// Non-fatal: user can still use the app, vault will init on next login
	}

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

	tokens, err := h.generateTokens(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) {
	var req struct {
		Login    string `json:"login" binding:"required"`    // username or email
		Password string `json:"password" binding:"required"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	user, err := h.stores.Users.GetByLogin(c.Request.Context(), req.Login)
	if err != nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid credentials"})
		return
	}

	if !user.IsActive {
		c.JSON(http.StatusForbidden, gin.H{"error": "account is inactive"})
		return
	}

	if err := bcrypt.CompareHashAndPassword([]byte(user.PasswordHash), []byte(req.Password)); err != nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid credentials"})
		return
	}

	// Vault: unwrap UEK into session cache (or init if first login post-migration)
	h.unlockVault(c.Request.Context(), user, req.Password)

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

	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) 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))
	}

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

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,
		},
	}, nil
}

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

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

// 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
	}

	_, err = database.DB.ExecContext(ctx, `
		UPDATE users
		SET encrypted_uek = $1, uek_salt = $2, uek_nonce = $3, vault_set = true
		WHERE id = $4
	`, encryptedUEK, salt, nonce, userID)
	if 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
	}

	var vaultSet bool
	var encryptedUEK, salt, nonce []byte

	err := database.DB.QueryRowContext(ctx, `
		SELECT vault_set, encrypted_uek, uek_salt, uek_nonce
		FROM users WHERE id = $1
	`, user.ID).Scan(&vaultSet, &encryptedUEK, &salt, &nonce)
	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 {
		log.Printf("⚠ Vault unlock 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).
func BootstrapAdmin(cfg *config.Config, s store.Stores) {
	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,
		})
		log.Printf("  ✅ Admin user '%s' updated", cfg.AdminUsername)
		return
	}

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

	user := &models.User{
		Username:     strings.ToLower(cfg.AdminUsername),
		Email:        strings.ToLower(email),
		PasswordHash: string(hash),
		Role:         models.UserRoleAdmin,
		IsActive:     true,
	}

	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) {
	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,
			})
			log.Printf("  🌱 Seed user '%s' updated (role=%s)", username, role)
			continue
		}

		user := &models.User{
			Username:     username,
			Email:        username + "@switchboard.local",
			PasswordHash: string(hash),
			Role:         role,
			IsActive:     true,
		}

		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
}