package auth

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/archivmail/internal/ldapauth"
	ldapcfg "github.com/archivmail/internal/ldapconfig"
	"github.com/archivmail/internal/userstore"
)

// Session holds the claims extracted from a validated JWT.
type Session struct {
	UserID     int64
	Username   string
	Role       string
	JTI        string // unique JWT ID
	TenantID   *int64
	TenantSlug string
}

// TenantDomainLookup is the interface used by Manager to resolve an email domain
// to a tenant ID. Implemented by tenantstore.Store.
type TenantDomainLookup interface {
	// GetTenantIDByDomain returns the tenant_id for a given domain, or nil if not found.
	GetTenantIDByDomain(ctx context.Context, domain string) (*int64, error)
}

// Manager handles login, token issuance, validation, and logout.
type Manager struct {
	store           *userstore.Store
	ldapStore       *ldapcfg.Store
	jwtSecret       []byte
	aesKey          []byte                 // PROJ-24: AES-256-GCM key for TOTP secret encryption
	tenantLdapStore *ldapcfg.TenantStore   // PROJ-23: per-tenant LDAP config
	tenantLookup    TenantDomainLookup     // PROJ-23: domain -> tenant_id resolution
}

// New creates a new auth Manager.
// ldapStore may be nil; in that case LDAP fallback is disabled.
// aesKey is the hex-encoded AES-256 key for encrypting TOTP secrets.
func New(store *userstore.Store, ldapStore *ldapcfg.Store, jwtSecret string, aesKey string) *Manager {
	aesKeyBytes, _ := hex.DecodeString(aesKey)
	return &Manager{
		store:     store,
		ldapStore: ldapStore,
		jwtSecret: []byte(jwtSecret),
		aesKey:    aesKeyBytes,
	}
}

// SetTenantLDAP wires the per-tenant LDAP config store and tenant domain lookup
// into the auth manager. Both may be nil to disable per-tenant LDAP.
func (m *Manager) SetTenantLDAP(tenantLdapStore *ldapcfg.TenantStore, tenantLookup TenantDomainLookup) {
	m.tenantLdapStore = tenantLdapStore
	m.tenantLookup = tenantLookup
}

// Login verifies credentials and returns a signed JWT token.
// It first attempts a local password check. If that fails and LDAP is
// configured and enabled, it falls back to LDAP authentication.
// If the user has TOTP enabled, totpRequired is true and the token is a
// short-lived pending token that can only be used with ValidateTOTPLogin.
func (m *Manager) Login(username, password string) (token string, user *userstore.User, totpRequired bool, err error) {
	// 1. Try local authentication first.
	user, err = m.store.VerifyPassword(username, password)
	if err == nil {
		if user.TOTPEnabled {
			t, e := m.issuePendingTOTPToken(user)
			return t, user, true, e
		}
		t, u, e := m.issueToken(user)
		return t, u, false, e
	}

	// 2. PROJ-23: Per-tenant LDAP — checked first so tenant config takes priority
	// over global LDAP (spec: tenant_ldap > ldap_config > config.yml).
	// Extract domain from username (user@domain.tld), resolve tenant, try auth.
	if m.tenantLdapStore != nil && m.tenantLookup != nil {
		if domain := extractDomain(username); domain != "" {
			ctx := context.Background()
			tenantID, lookupErr := m.tenantLookup.GetTenantIDByDomain(ctx, domain)
			if lookupErr != nil {
				fmt.Printf("[DEBUG] tenant domain lookup failed for %q: %v\n", domain, lookupErr)
			}
			if lookupErr == nil && tenantID != nil {
				tcfg, tErr := m.tenantLdapStore.GetWithPassword(ctx, *tenantID)
				if tErr != nil {
					fmt.Printf("[DEBUG] tenant LDAP GetWithPassword failed: %v\n", tErr)
				}
				if tErr == nil && tcfg != nil && tcfg.Enabled && tcfg.URL != "" && tcfg.BindPassword != "" {
					attrs, authErr := ldapauth.Authenticate(ldapauth.Config{
						URL:           tcfg.URL,
						BindDN:        tcfg.BindDN,
						BindPassword:  tcfg.BindPassword,
						BaseDN:        tcfg.BaseDN,
						UserFilter:    tcfg.UserFilter,
						TLS:           tcfg.TLS,
						TLSSkipVerify: tcfg.TLSSkipVerify,
					}, username, password)
					if authErr != nil {
						fmt.Printf("[DEBUG] tenant LDAP auth failed for %q: %v\n", username, authErr)
					}
					if authErr == nil {
						role := tcfg.DefaultRole
						if role == "" {
							role = userstore.RoleUser
						}
						memberOf := attrs["memberOf"]
						if memberOf != "" {
							for _, gm := range tcfg.GroupMappings {
								if gm.GroupDN != "" && containsGroup(memberOf, gm.GroupDN) {
									role = gm.Role
									break
								}
							}
						}
						email := attrs["mail"]
						if email == "" {
							email = username
						}
						ldapUser, upsertErr := m.store.UpsertLDAPUser(username, email, role, tenantID)
						if upsertErr == nil {
							if ldapUser.TOTPEnabled {
								t, e := m.issuePendingTOTPToken(ldapUser)
								return t, ldapUser, true, e
							}
							t, u, e := m.issueToken(ldapUser)
							return t, u, false, e
						}
					}
				}
			}
		}
	}

	// 3. Global LDAP fallback — only reached if no matching tenant LDAP config found.
	if m.ldapStore != nil {
		cfg, ldapErr := m.ldapStore.GetWithPassword(context.Background())
		if ldapErr == nil && cfg != nil && cfg.Enabled && cfg.URL != "" && cfg.BindPassword != "" {
			attrs, authErr := ldapauth.Authenticate(ldapauth.Config{
				URL:           cfg.URL,
				BindDN:        cfg.BindDN,
				BindPassword:  cfg.BindPassword,
				BaseDN:        cfg.BaseDN,
				UserFilter:    cfg.UserFilter,
				TLS:           cfg.TLS,
				TLSSkipVerify: cfg.TLSSkipVerify,
			}, username, password)
			if authErr == nil {
				role := cfg.DefaultRole
				if role == "" {
					role = userstore.RoleUser
				}
				memberOf := attrs["memberOf"]
				if memberOf != "" {
					for _, gm := range cfg.GroupMappings {
						if gm.GroupDN != "" && containsGroup(memberOf, gm.GroupDN) {
							role = gm.Role
							break
						}
					}
				}
				email := attrs["mail"]
				if email == "" {
					email = username + "@ldap.local"
				}
				ldapUser, upsertErr := m.store.UpsertLDAPUser(username, email, role, nil)
				if upsertErr == nil {
					if ldapUser.TOTPEnabled {
						t, e := m.issuePendingTOTPToken(ldapUser)
						return t, ldapUser, true, e
					}
					t, u, e := m.issueToken(ldapUser)
					return t, u, false, e
				}
			}
		}
	}

	return "", nil, false, fmt.Errorf("auth: login: invalid credentials")
}

// issueToken signs a JWT for the given user and returns the token string.
func (m *Manager) issueToken(user *userstore.User) (string, *userstore.User, error) {
	jti := generateJTI()
	now := time.Now()

	var tenantIDVal int64
	if user.TenantID != nil {
		tenantIDVal = *user.TenantID
	}

	claims := jwt.MapClaims{
		"sub":         user.Username,
		"role":        user.Role,
		"uid":         user.ID,
		"jti":         jti,
		"iat":         now.Unix(),
		"exp":         now.Add(8 * time.Hour).Unix(),
		"tenant_id":   tenantIDVal,
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	signed, err := token.SignedString(m.jwtSecret)
	if err != nil {
		return "", nil, fmt.Errorf("auth: sign token: %w", err)
	}

	return signed, user, nil
}

// issuePendingTOTPToken issues a short-lived JWT (5 min) that signals TOTP is required.
// This token MUST NOT be accepted as a full auth token — only for /api/auth/totp.
func (m *Manager) issuePendingTOTPToken(user *userstore.User) (string, error) {
	jti := generateJTI()
	now := time.Now()

	claims := jwt.MapClaims{
		"sub":          user.Username,
		"uid":          user.ID,
		"jti":          jti,
		"iat":          now.Unix(),
		"exp":          now.Add(5 * time.Minute).Unix(),
		"totp_pending": true,
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	signed, err := token.SignedString(m.jwtSecret)
	if err != nil {
		return "", fmt.Errorf("auth: sign pending totp token: %w", err)
	}
	return signed, nil
}

// ValidateTOTPLogin validates a pending TOTP token and TOTP code, then issues a full JWT.
func (m *Manager) ValidateTOTPLogin(pendingToken, code string) (string, *userstore.User, error) {
	// Parse the pending token
	token, err := jwt.Parse(pendingToken, func(t *jwt.Token) (interface{}, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("auth: unexpected signing method: %v", t.Header["alg"])
		}
		return m.jwtSecret, nil
	})
	if err != nil {
		return "", nil, fmt.Errorf("auth: invalid pending token: %w", err)
	}
	if !token.Valid {
		return "", nil, errors.New("auth: pending token not valid")
	}

	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return "", nil, errors.New("auth: bad claims")
	}

	// Verify this is actually a pending TOTP token
	pending, _ := claims["totp_pending"].(bool)
	if !pending {
		return "", nil, errors.New("auth: not a pending TOTP token")
	}

	// Extract user ID
	var userID int64
	switch v := claims["uid"].(type) {
	case float64:
		userID = int64(v)
	case int64:
		userID = v
	default:
		return "", nil, errors.New("auth: missing uid in pending token")
	}

	// Load user from DB
	user, err := m.store.GetByID(userID)
	if err != nil {
		return "", nil, fmt.Errorf("auth: user lookup: %w", err)
	}
	if !user.TOTPEnabled {
		return "", nil, errors.New("auth: TOTP not enabled for this user")
	}

	// Load and decrypt TOTP secret
	encSecret, enabled, err := m.store.GetTOTPSecret(context.Background(), userID)
	if err != nil {
		return "", nil, fmt.Errorf("auth: get totp secret: %w", err)
	}
	if !enabled || len(encSecret) == 0 {
		return "", nil, errors.New("auth: TOTP not configured")
	}

	plainSecret, err := m.decryptAES(encSecret)
	if err != nil {
		return "", nil, fmt.Errorf("auth: decrypt totp secret: %w", err)
	}

	// Validate the TOTP code
	if !ValidateTOTP(string(plainSecret), code) {
		return "", nil, errors.New("auth: invalid TOTP code")
	}

	// Issue a full auth token
	return m.issueToken(user)
}

// ValidateToken parses and validates the token, checking the blacklist.
func (m *Manager) ValidateToken(tokenStr string) (*Session, error) {
	token, err := jwt.Parse(tokenStr, func(t *jwt.Token) (interface{}, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("auth: unexpected signing method: %v", t.Header["alg"])
		}
		return m.jwtSecret, nil
	})
	if err != nil {
		return nil, fmt.Errorf("auth: invalid token: %w", err)
	}
	if !token.Valid {
		return nil, errors.New("auth: token not valid")
	}

	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return nil, errors.New("auth: bad claims")
	}

	// PROJ-24: Reject pending TOTP tokens — they must not be used as full auth tokens.
	if pending, _ := claims["totp_pending"].(bool); pending {
		return nil, errors.New("auth: pending TOTP token cannot be used for authentication")
	}

	jti, _ := claims["jti"].(string)
	blacklisted, err := m.store.IsBlacklisted(jti)
	if err != nil {
		return nil, fmt.Errorf("auth: blacklist check: %w", err)
	}
	if blacklisted {
		return nil, errors.New("auth: token revoked")
	}

	username, _ := claims["sub"].(string)
	role, _ := claims["role"].(string)

	var userID int64
	switch v := claims["uid"].(type) {
	case float64:
		userID = int64(v)
	case int64:
		userID = v
	}

	var tenantID *int64
	switch v := claims["tenant_id"].(type) {
	case float64:
		if v != 0 {
			id := int64(v)
			tenantID = &id
		}
	case int64:
		if v != 0 {
			tenantID = &v
		}
	}

	return &Session{
		UserID:   userID,
		Username: username,
		Role:     role,
		JTI:      jti,
		TenantID: tenantID,
	}, nil
}

// Logout revokes the token by adding its JTI to the blacklist.
func (m *Manager) Logout(tokenStr string) error {
	token, err := jwt.Parse(tokenStr, func(t *jwt.Token) (interface{}, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("auth: unexpected signing method")
		}
		return m.jwtSecret, nil
	})
	if err != nil {
		return fmt.Errorf("auth: logout parse: %w", err)
	}

	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return errors.New("auth: bad claims on logout")
	}

	jti, _ := claims["jti"].(string)
	var exp time.Time
	switch v := claims["exp"].(type) {
	case float64:
		exp = time.Unix(int64(v), 0)
	case int64:
		exp = time.Unix(v, 0)
	default:
		exp = time.Now().Add(8 * time.Hour)
	}

	return m.store.BlacklistToken(jti, exp)
}

// HasRole returns true when userRole satisfies the required role level.
// Hierarchy: superadmin > admin > domain_admin > auditor > user
func HasRole(userRole, required string) bool {
	levels := map[string]int{
		userstore.RoleUser:        1,
		userstore.RoleAuditor:     2,
		userstore.RoleDomainAdmin: 3,
		userstore.RoleAdmin:       4,
		userstore.RoleSuperAdmin:  5,
	}
	return levels[userRole] >= levels[required]
}

// containsGroup checks whether a comma-separated memberOf string contains groupDN
// (case-insensitive substring match to handle varying DN formats).
func containsGroup(memberOf, groupDN string) bool {
	for _, dn := range splitMemberOf(memberOf) {
		if dn == groupDN {
			return true
		}
	}
	return false
}

func splitMemberOf(s string) []string {
	var out []string
	for _, part := range splitComma(s) {
		part = trimSpace(part)
		if part != "" {
			out = append(out, part)
		}
	}
	return out
}

// splitComma splits on commas that are not part of DN attribute values.
// For simplicity we split on ", " and "," and let callers match.
func splitComma(s string) []string {
	// Groups returned from LDAP memberOf are newline-separated in the map
	// because we join with "," in ldapauth. Re-split here.
	parts := []string{}
	start := 0
	for i := 0; i < len(s); i++ {
		if s[i] == ',' && (i == 0 || s[i-1] != '\\') {
			parts = append(parts, s[start:i])
			start = i + 1
		}
	}
	parts = append(parts, s[start:])
	return parts
}

func trimSpace(s string) string {
	start, end := 0, len(s)
	for start < end && (s[start] == ' ' || s[start] == '\t') {
		start++
	}
	for end > start && (s[end-1] == ' ' || s[end-1] == '\t') {
		end--
	}
	return s[start:end]
}

// extractDomain returns the domain part from a username containing '@'.
// Returns empty string if no '@' is found.
func extractDomain(username string) string {
	idx := strings.LastIndex(username, "@")
	if idx < 0 || idx == len(username)-1 {
		return ""
	}
	return strings.ToLower(username[idx+1:])
}

// generateJTI returns a cryptographically random identifier for a JWT.
func generateJTI() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		// fallback: should never happen on a healthy system
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	return hex.EncodeToString(b)
}

// ── PROJ-24: AES-256-GCM encryption helpers for TOTP secrets ─────────────

// EncryptAES encrypts plaintext using AES-256-GCM with the manager's AES key.
// The 12-byte nonce is prepended to the ciphertext.
func (m *Manager) EncryptAES(plaintext []byte) ([]byte, error) {
	return encryptAESGCM(m.aesKey, plaintext)
}

// decryptAES decrypts data encrypted with EncryptAES.
func (m *Manager) decryptAES(data []byte) ([]byte, error) {
	return decryptAESGCM(m.aesKey, data)
}

// DecryptAESForHandler exposes AES decryption for use in API handlers (e.g., TOTP setup confirmation).
func (m *Manager) DecryptAESForHandler(data []byte) ([]byte, error) {
	return decryptAESGCM(m.aesKey, data)
}

func encryptAESGCM(key, plaintext []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("auth: aes cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("auth: gcm: %w", err)
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("auth: random nonce: %w", err)
	}
	ciphertext := gcm.Seal(nonce, nonce, plaintext, nil)
	return ciphertext, nil
}

func decryptAESGCM(key, data []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("auth: aes cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("auth: gcm: %w", err)
	}
	nonceSize := gcm.NonceSize()
	if len(data) < nonceSize {
		return nil, fmt.Errorf("auth: ciphertext too short")
	}
	nonce, ciphertext := data[:nonceSize], data[nonceSize:]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("auth: decrypt: %w", err)
	}
	return plaintext, nil
}

// GetUserStore returns the underlying user store (used by TOTP handlers).
func (m *Manager) GetUserStore() *userstore.Store {
	return m.store
}