package middleware

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
	"github.com/tingly-dev/tingly-box/internal/data/db"
	"github.com/tingly-dev/tingly-box/internal/server/config"
	"github.com/tingly-dev/tingly-box/pkg/auth"
)

// AuthMiddleware provides authentication middleware for different types of authentication
type AuthMiddleware struct {
	config          *config.Config
	jwtManager      *auth.JWTManager
	apiTokenManager *auth.APITokenManager
	apiTokenStore   APITokenStore
}

// APITokenStore interface for token validation (abstracted for testability)
type APITokenStore interface {
	ValidateToken(tokenID string) (*db.APITokenRecord, error)
	UpdateLastUsed(tokenID string) error
}

// ErrorResponse represents an error response
type ErrorResponse struct {
	Error ErrorDetail `json:"error"`
}

// abortWithError sends an error response and adds the error to gin context for logging
func abortWithError(c *gin.Context, statusCode int, message string, errorType string) {
	// Add error to context so logging middleware can capture it
	c.Error(fmt.Errorf("%s: %s", errorType, message)).SetType(gin.ErrorTypePublic)

	// Send JSON response
	c.JSON(statusCode, ErrorResponse{
		Error: ErrorDetail{
			Message: message,
			Type:    errorType,
		},
	})
	c.Abort()
}

// ErrorDetail represents error details
type ErrorDetail struct {
	Message string `json:"message"`
	Type    string `json:"type"`
	Code    string `json:"code,omitempty"`
}

// NewAuthMiddleware creates a new authentication middleware
func NewAuthMiddleware(cfg *config.Config, jwtManager *auth.JWTManager, apiTokenManager *auth.APITokenManager, apiTokenStore APITokenStore) *AuthMiddleware {
	return &AuthMiddleware{
		config:          cfg,
		jwtManager:      jwtManager,
		apiTokenManager: apiTokenManager,
		apiTokenStore:   apiTokenStore,
	}
}

type enterpriseContextClaims struct {
	UserID       string
	DepartmentID string
	KeyPrefix    string
	Tier         string
	JTI          string
}

func containsFold(items []string, target string) bool {
	for _, item := range items {
		if strings.EqualFold(strings.TrimSpace(item), strings.TrimSpace(target)) {
			return true
		}
	}
	return false
}

func resolveSecretRef(raw string) (string, error) {
	ref := strings.TrimSpace(raw)
	if ref == "" {
		return "", nil
	}
	if strings.HasPrefix(ref, "env:") {
		return strings.TrimSpace(os.Getenv(strings.TrimSpace(strings.TrimPrefix(ref, "env:")))), nil
	}
	if strings.HasPrefix(ref, "file:") {
		path := strings.TrimSpace(strings.TrimPrefix(ref, "file:"))
		content, err := os.ReadFile(path)
		if err != nil {
			return "", err
		}
		return strings.TrimSpace(string(content)), nil
	}
	if st, err := os.Stat(ref); err == nil && !st.IsDir() {
		content, readErr := os.ReadFile(ref)
		if readErr != nil {
			return "", readErr
		}
		return strings.TrimSpace(string(content)), nil
	}
	return ref, nil
}

func parseRSAPublicKeyFromPEM(raw string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(raw))
	if block == nil {
		return nil, errors.New("invalid pem block")
	}
	if key, err := x509.ParsePKIXPublicKey(block.Bytes); err == nil {
		rsaKey, ok := key.(*rsa.PublicKey)
		if !ok {
			return nil, errors.New("public key is not RSA")
		}
		return rsaKey, nil
	}
	cert, certErr := x509.ParseCertificate(block.Bytes)
	if certErr == nil {
		if rsaKey, ok := cert.PublicKey.(*rsa.PublicKey); ok {
			return rsaKey, nil
		}
	}
	return nil, errors.New("failed to parse rsa public key")
}

func verifyEnterpriseContextJWT(cfg *config.Config, rawToken string) (*enterpriseContextClaims, error) {
	tokenString := strings.TrimSpace(rawToken)
	if tokenString == "" || cfg == nil {
		return nil, nil
	}
	jwtCfg := cfg.EnterpriseContextJWT
	if !jwtCfg.Enabled {
		return nil, nil
	}

	clockSkew := time.Duration(jwtCfg.ClockSkewSeconds) * time.Second
	parsed, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		alg := strings.ToUpper(strings.TrimSpace(token.Method.Alg()))
		if len(jwtCfg.AlgAllowlist) > 0 && !containsFold(jwtCfg.AlgAllowlist, alg) {
			return nil, fmt.Errorf("algorithm %s is not allowed", alg)
		}
		switch alg {
		case "HS256":
			secret, secretErr := resolveSecretRef(jwtCfg.HS256SecretRef)
			if secretErr != nil {
				return nil, secretErr
			}
			if strings.TrimSpace(secret) == "" {
				return nil, errors.New("hs256 secret is empty")
			}
			return []byte(secret), nil
		case "RS256":
			// If no public keys were provided, fall back to RS256PublicKeyRef.
			if len(jwtCfg.PublicKeys) == 0 && strings.TrimSpace(jwtCfg.RS256PublicKeyRef) != "" {
				ref, refErr := resolveSecretRef(jwtCfg.RS256PublicKeyRef)
				if refErr == nil && strings.TrimSpace(ref) != "" {
					if pub, parseErr := parseRSAPublicKeyFromPEM(ref); parseErr == nil {
						return pub, nil
					}
				}
			}
			kid, _ := token.Header["kid"].(string)
			kid = strings.TrimSpace(kid)
			for _, item := range jwtCfg.PublicKeys {
				if kid != "" && strings.TrimSpace(item.KID) != kid {
					continue
				}
				pub, parseErr := parseRSAPublicKeyFromPEM(item.PEM)
				if parseErr != nil {
					continue
				}
				return pub, nil
			}
			if strings.TrimSpace(jwtCfg.JWKSURL) != "" {
				return nil, errors.New("jwks_url is configured but dynamic jwks fetch is not implemented yet")
			}
			return nil, errors.New("no rs256 public key matched")
		default:
			return nil, fmt.Errorf("unsupported jwt alg: %s", alg)
		}
	}, jwt.WithLeeway(clockSkew))
	if err != nil || !parsed.Valid {
		if err != nil {
			return nil, err
		}
		return nil, errors.New("invalid context jwt")
	}

	mapClaims, ok := parsed.Claims.(jwt.MapClaims)
	if !ok {
		return nil, errors.New("invalid jwt claims")
	}

	if len(jwtCfg.AllowedIssuers) > 0 {
		issuer, _ := mapClaims.GetIssuer()
		if !containsFold(jwtCfg.AllowedIssuers, issuer) {
			return nil, fmt.Errorf("issuer %s is not allowed", issuer)
		}
	}

	if len(jwtCfg.AllowedAudiences) > 0 {
		aud, audErr := mapClaims.GetAudience()
		if audErr != nil {
			return nil, audErr
		}
		found := false
		for _, candidate := range aud {
			if containsFold(jwtCfg.AllowedAudiences, candidate) {
				found = true
				break
			}
		}
		if !found {
			return nil, errors.New("audience is not allowed")
		}
	}

	userID, _ := mapClaims["sub"].(string)
	departmentID, _ := mapClaims["dept_id"].(string)
	keyPrefix, _ := mapClaims["key_prefix"].(string)
	tier, _ := mapClaims["tier"].(string)
	jti, _ := mapClaims["jti"].(string)

	if strings.TrimSpace(userID) == "" {
		return nil, errors.New("missing sub claim")
	}
	if jwtCfg.RequireJTI && strings.TrimSpace(jti) == "" {
		return nil, errors.New("missing jti claim")
	}

	return &enterpriseContextClaims{
		UserID:       strings.TrimSpace(userID),
		DepartmentID: strings.TrimSpace(departmentID),
		KeyPrefix:    strings.TrimSpace(keyPrefix),
		Tier:         strings.TrimSpace(tier),
		JTI:          strings.TrimSpace(jti),
	}, nil
}

// UserAuthMiddleware middleware for UI and control API authentication
func (am *AuthMiddleware) UserAuthMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		authHeader := c.GetHeader("Authorization")
		if authHeader == "" {
			abortWithError(c, http.StatusUnauthorized, "User authorization header required", "invalid_request_error")
			return
		}

		// Extract token from "Bearer <token>" format
		tokenParts := strings.Split(authHeader, " ")
		if len(tokenParts) != 2 || tokenParts[0] != "Bearer" {
			abortWithError(c, http.StatusUnauthorized, "Invalid user authorization header format. Expected: 'Bearer <token>'", "invalid_request_error")
			return
		}

		token := tokenParts[1]

		// Check against global config user token first
		cfg := am.config
		if cfg != nil && cfg.HasUserToken() {
			configToken := cfg.GetUserToken()

			// Remove "Bearer " prefix if present in the token
			if strings.HasPrefix(token, "Bearer ") {
				token = token[7:]
			}

			// Direct token comparison
			if token == configToken || strings.TrimPrefix(token, "Bearer ") == configToken {
				// Token matches the one in global config, allow access
				c.Set("client_id", "user_authenticated")
				// For UI user authentication, set user_id to default admin
				// This matches the migrated user_id in usage_records
				c.Set("user_id", db.DefaultAdminUserID)
				c.Next()
				return
			}
		}

		abortWithError(c, http.StatusUnauthorized, "Invalid user authorization token.", "invalid_request_error")
	}
}

// ModelAuthMiddleware middleware for OpenAI and Anthropic API authentication
// The auth will support both `Authorization` and `X-Api-Key`
// Supports three authentication methods (in order of precedence):
// 1. JWT API tokens (when multi-tenant is enabled)
// 2. Global config model token (backward compatibility)
// 3. Enterprise context JWT (X-TBE-Context-JWT header)
func (am *AuthMiddleware) ModelAuthMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		authHeader := c.GetHeader("Authorization")
		xApiKey := c.GetHeader("X-Api-Key")
		if authHeader == "" && xApiKey == "" {
			abortWithError(c, http.StatusUnauthorized, "Model authorization header required", "invalid_request_error")
			return
		}

		if authHeader == "" {
			authHeader = xApiKey
		}

		token := authHeader
		// Remove "Bearer " prefix if present in the token
		if strings.HasPrefix(token, "Bearer ") {
			token = token[7:]
		}
		token = strings.TrimSpace(token)
		xApiKey = strings.TrimSpace(xApiKey)

		cfg := am.config

		// Try API token authentication first (if multi-tenant is enabled)
		// Check if token has tb-share- prefix
		if cfg != nil && cfg.IsMultiTenantEnabled() && am.apiTokenStore != nil {
			if strings.HasPrefix(token, "tb-share-") {
				// This is an API token, validate it directly from database
				tokenRecord, validateErr := am.apiTokenStore.ValidateToken(token)
				if validateErr == nil && tokenRecord != nil {
					// Token is valid and enabled
					c.Set("user_id", tokenRecord.UserID)
					c.Set("token_id", tokenRecord.TokenID)
					c.Set("client_id", "model_authenticated")
					c.Set("auth_method", "api_token")

					// Update last used asynchronously
					go am.apiTokenStore.UpdateLastUsed(tokenRecord.TokenID)

					c.Next()
					return
				}
			}
		}

		// Check global config model token (backward compatibility)
		if cfg == nil || !cfg.HasModelToken() {
			abortWithError(c, http.StatusInternalServerError, "config or config model token missing", "invalid_request_error")
			return
		}

		// If global token is disabled and multi-tenant is enabled, reject
		if cfg.IsMultiTenantEnabled() && cfg.IsGlobalTokenDisabled() {
			abortWithError(c, http.StatusUnauthorized, "Global token disabled, use API token", "invalid_token_error")
			return
		}

		configToken := cfg.GetModelToken()

		// Direct token comparison
		if token == configToken || xApiKey == configToken {
			c.Set("client_id", "model_authenticated")
			c.Set("auth_method", "global_token")
			// Set UserID to default admin for usage tracking consistency
			// This matches the migrated user_id values in usage_records
			c.Set("user_id", db.DefaultAdminUserID)
			contextJWT := strings.TrimSpace(c.GetHeader("X-TBE-Context-JWT"))
			if contextJWT != "" {
				claims, verifyErr := verifyEnterpriseContextJWT(cfg, contextJWT)
				if verifyErr != nil {
					abortWithError(c, http.StatusUnauthorized, "Invalid enterprise context jwt", "invalid_request_error")
					return
				}
				if claims != nil {
					c.Set("enterprise_user_id", claims.UserID)
					c.Set("enterprise_department_id", claims.DepartmentID)
					c.Set("enterprise_key_prefix", claims.KeyPrefix)
					c.Set("enterprise_user_tier", claims.Tier)
					c.Set("enterprise_context_jti", claims.JTI)
					c.Set("enterprise_context_verified", true)
				}
			}
			c.Next()
			return
		}

		// Enterprise short-lived access token authentication.
		requestToken := token
		if requestToken == "" {
			requestToken = xApiKey
		}
		if strings.HasPrefix(strings.TrimSpace(requestToken), "sk-tbe-") {
			abortWithError(c, http.StatusUnauthorized, "Virtual key must be used through TBE /tbe/* endpoints", "invalid_request_error")
			return
		}

		abortWithError(c, http.StatusUnauthorized, "Invalid model authorization token.", "invalid_request_error")
	}
}
