package auth

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

type Permission string

const (
	OIDCSessionCookieName       = "easyai_gateway_oidc_session"
	localBreakGlassTokenPurpose = "local_break_glass_manager"
	legacyAccessTokenPurpose    = "legacy_access"

	PermissionPublic  Permission = "public"
	PermissionBasic   Permission = "basic"
	PermissionCreat   Permission = "creat"
	PermissionPower   Permission = "power"
	PermissionManager Permission = "manager"
)

type User struct {
	ID                string    `json:"sub"`
	Username          string    `json:"username"`
	DisplayName       string    `json:"-"`
	Email             string    `json:"-"`
	Phone             string    `json:"-"`
	AvatarURL         string    `json:"-"`
	Roles             []string  `json:"role,omitempty"`
	ContextType       string    `json:"contextType,omitempty"`
	TenantID          string    `json:"tenantId,omitempty"`
	TenantName        string    `json:"tenantName,omitempty"`
	GatewayTenantID   string    `json:"gatewayTenantId,omitempty"`
	TenantKey         string    `json:"tenantKey,omitempty"`
	SSOID             string    `json:"sso_id,omitempty"`
	Source            string    `json:"source,omitempty"`
	GatewayUserID     string    `json:"gatewayUserId,omitempty"`
	UserGroupID       string    `json:"userGroupId,omitempty"`
	UserGroupKey      string    `json:"userGroupKey,omitempty"`
	UserGroupKeys     []string  `json:"userGroupKeys,omitempty"`
	APIKeyID          string    `json:"apiKeyId,omitempty"`
	APIKeySecret      string    `json:"apiKeySecret,omitempty"`
	APIKeyName        string    `json:"apiKeyName,omitempty"`
	APIKeyPrefix      string    `json:"apiKeyPrefix,omitempty"`
	APIKeyScopes      []string  `json:"apiKeyScopes,omitempty"`
	TokenExpiresAt    time.Time `json:"-"`
	TokenIssuedAt     time.Time `json:"-"`
	Issuer            string    `json:"-"`
	ApplicationID     string    `json:"-"`
	OIDCClientID      string    `json:"-"`
	OIDCUserBindingID string    `json:"-"`
	TokenPurpose      string    `json:"-"`
}

type contextKey string

const userContextKey contextKey = "easyai-auth-user"

var ErrUnauthorized = errors.New("unauthorized")

type RequestAuthError struct {
	Status            int
	Code              string
	Message           string
	RetryAfterSeconds int
}

func (e *RequestAuthError) Error() string { return e.Code }

func NewRequestAuthError(status int, code, message string) error {
	return &RequestAuthError{Status: status, Code: code, Message: message}
}

func NewRetryableRequestAuthError(status int, code, message string, retryAfterSeconds int) error {
	return &RequestAuthError{
		Status: status, Code: code, Message: message, RetryAfterSeconds: retryAfterSeconds,
	}
}

type Authenticator struct {
	JWTSecret                   string
	ServerMainBaseURL           string
	ServerMainInternalToken     string
	ServerMainInternalKey       string
	ServerMainInternalSecret    string
	HTTPClient                  *http.Client
	LocalAPIKeyVerifier         func(ctx context.Context, apiKey string) (*User, error)
	OIDCVerifier                *OIDCVerifier
	OIDCVerifierProvider        func() *OIDCVerifier
	OIDCSessionResolver         func(ctx context.Context, sessionID string) (*User, error)
	OIDCSessionResolverProvider func(ctx context.Context, sessionID string) (*User, error)
	LegacyJWTEnabled            bool
	LegacyJWTEnabledProvider    func() bool
}

func New(jwtSecret string, serverMainBaseURL string, internalToken string) *Authenticator {
	return &Authenticator{
		JWTSecret:               jwtSecret,
		ServerMainBaseURL:       strings.TrimRight(serverMainBaseURL, "/"),
		ServerMainInternalToken: internalToken,
		HTTPClient: &http.Client{
			Timeout: 10 * time.Second,
		},
		LegacyJWTEnabled: true,
	}
}

func UserFromContext(ctx context.Context) (*User, bool) {
	user, ok := ctx.Value(userContextKey).(*User)
	return user, ok
}

func WithUser(ctx context.Context, user *User) context.Context {
	return context.WithValue(ctx, userContextKey, user)
}

func (a *Authenticator) Require(permission Permission, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		user, err := a.Authenticate(r)
		if err != nil {
			if permission == PermissionPublic {
				next.ServeHTTP(w, r)
				return
			}
			var requestError *RequestAuthError
			if errors.As(err, &requestError) {
				if requestError.RetryAfterSeconds > 0 {
					w.Header().Set("Retry-After", strconv.Itoa(requestError.RetryAfterSeconds))
				}
				w.Header().Set("Content-Type", "application/json; charset=utf-8")
				w.Header().Set("Cache-Control", "no-store")
				w.WriteHeader(requestError.Status)
				_ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{
					"message": requestError.Message, "status": requestError.Status, "code": requestError.Code,
				}})
				return
			}
			if strings.HasPrefix(err.Error(), ErrUnauthorized.Error()+":") {
				slog.WarnContext(r.Context(), "OIDC authentication rejected", "reason", err.Error())
			}
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		if !hasPermission(user.Roles, permission) {
			http.Error(w, "forbidden", http.StatusForbidden)
			return
		}
		next.ServeHTTP(w, r.WithContext(WithUser(r.Context(), user)))
	})
}

func (a *Authenticator) Authenticate(r *http.Request) (*User, error) {
	var token string
	if authorization := strings.TrimSpace(r.Header.Get("Authorization")); authorization != "" {
		token = extractBearer(authorization)
		if token == "" {
			return nil, ErrUnauthorized
		}
	} else if value := strings.TrimSpace(r.Header.Get("x-comfy-api-key")); value != "" {
		token = value
	} else if value := strings.TrimSpace(r.Header.Get("x-goog-api-key")); value != "" {
		token = value
	} else if value := strings.TrimSpace(r.URL.Query().Get("key")); value != "" {
		// Query credentials are retained only for API compatibility with
		// providers that use `?key=sk-*`. Bearer/OIDC tokens in URLs would leak
		// through browser history, reverse-proxy logs, and diagnostics.
		if !strings.HasPrefix(value, "sk-") {
			return nil, ErrUnauthorized
		}
		token = value
	}
	if token == "" {
		if cookie, err := r.Cookie(OIDCSessionCookieName); err == nil {
			sessionID := strings.TrimSpace(cookie.Value)
			resolver := a.OIDCSessionResolver
			if a.OIDCSessionResolverProvider != nil {
				resolver = a.OIDCSessionResolverProvider
			}
			if sessionID == "" || resolver == nil {
				return nil, ErrUnauthorized
			}
			return resolver(r.Context(), sessionID)
		}
	}
	if token == "" {
		return nil, ErrUnauthorized
	}
	if strings.HasPrefix(token, "sk-") {
		return a.verifyAPIKey(r.Context(), token)
	}
	algorithm := jwtAlgorithm(token)
	if algorithm == "RS256" || algorithm == "ES256" {
		return a.AuthenticateOIDCAccessToken(r.Context(), token)
	}
	user, err := a.verifyJWT(token)
	if err != nil {
		return nil, err
	}
	if a.legacyJWTEnabled() || isLocalBreakGlassManager(user) {
		return user, nil
	}
	return nil, ErrUnauthorized
}

func (a *Authenticator) AuthenticateOIDCAccessToken(ctx context.Context, token string) (*User, error) {
	algorithm := jwtAlgorithm(token)
	verifier := a.OIDCVerifier
	if a.OIDCVerifierProvider != nil {
		verifier = a.OIDCVerifierProvider()
	}
	if verifier == nil || algorithm != "RS256" && algorithm != "ES256" {
		return nil, ErrUnauthorized
	}
	return verifier.Verify(ctx, token)
}

func (a *Authenticator) legacyJWTEnabled() bool {
	if a.LegacyJWTEnabledProvider != nil {
		return a.LegacyJWTEnabledProvider()
	}
	return a.LegacyJWTEnabled
}

func (a *Authenticator) verifyJWT(tokenString string) (*User, error) {
	token, err := jwt.ParseWithClaims(tokenString, jwt.MapClaims{}, func(token *jwt.Token) (any, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return []byte(a.JWTSecret), nil
	})
	if err != nil || !token.Valid {
		return nil, ErrUnauthorized
	}
	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return nil, ErrUnauthorized
	}

	user := &User{
		ID:              stringClaim(claims, "sub"),
		Username:        stringClaim(claims, "username"),
		Roles:           stringSliceClaim(claims, "role"),
		TenantID:        stringClaim(claims, "tenantId"),
		GatewayTenantID: stringClaim(claims, "gatewayTenantId"),
		TenantKey:       stringClaim(claims, "tenantKey"),
		SSOID:           stringClaim(claims, "sso_id"),
		Source:          stringClaim(claims, "source"),
		GatewayUserID:   stringClaim(claims, "gatewayUserId"),
		UserGroupID:     stringClaim(claims, "userGroupId"),
		UserGroupKey:    stringClaim(claims, "userGroupKey"),
		UserGroupKeys:   stringSliceClaim(claims, "userGroupKeys"),
		APIKeyID:        stringClaim(claims, "apiKeyId"),
		APIKeySecret:    stringClaim(claims, "apiKeySecret"),
		APIKeyName:      stringClaim(claims, "apiKeyName"),
		APIKeyPrefix:    stringClaim(claims, "apiKeyPrefix"),
		APIKeyScopes:    stringSliceClaim(claims, "apiKeyScopes"),
		TokenPurpose:    stringClaim(claims, "tokenPurpose"),
	}
	if user.Source == "" {
		user.Source = "gateway"
	}
	if user.ID == "" {
		return nil, ErrUnauthorized
	}
	return user, nil
}

func (a *Authenticator) SignJWT(user *User, ttl time.Duration) (string, error) {
	if ttl <= 0 {
		ttl = time.Hour
	}
	now := time.Now()
	tokenPurpose := legacyAccessTokenPurpose
	if user.Source == "gateway" && hasManagerRole(user.Roles) {
		tokenPurpose = localBreakGlassTokenPurpose
	}
	claims := jwt.MapClaims{
		"sub":             user.ID,
		"username":        user.Username,
		"role":            user.Roles,
		"tenantId":        user.TenantID,
		"gatewayTenantId": user.GatewayTenantID,
		"tenantKey":       user.TenantKey,
		"source":          user.Source,
		"gatewayUserId":   user.GatewayUserID,
		"userGroupId":     user.UserGroupID,
		"userGroupKey":    user.UserGroupKey,
		"userGroupKeys":   user.UserGroupKeys,
		"apiKeyId":        user.APIKeyID,
		"apiKeyName":      user.APIKeyName,
		"apiKeyPrefix":    user.APIKeyPrefix,
		"apiKeyScopes":    user.APIKeyScopes,
		"tokenPurpose":    tokenPurpose,
		"iat":             now.Unix(),
		"exp":             now.Add(ttl).Unix(),
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(a.JWTSecret))
}

func isLocalBreakGlassManager(user *User) bool {
	return user != nil && user.Source == "gateway" && user.TokenPurpose == localBreakGlassTokenPurpose && hasManagerRole(user.Roles)
}

func hasManagerRole(roles []string) bool {
	for _, role := range roles {
		if role == "manager" || role == "admin" {
			return true
		}
	}
	return false
}

func (a *Authenticator) verifyAPIKey(ctx context.Context, apiKey string) (*User, error) {
	if a.LocalAPIKeyVerifier != nil {
		user, err := a.LocalAPIKeyVerifier(ctx, apiKey)
		if err == nil {
			return user, nil
		}
		if !errors.Is(err, ErrUnauthorized) {
			return nil, NewRetryableRequestAuthError(
				http.StatusServiceUnavailable,
				"AUTH_STORE_UNAVAILABLE",
				"authentication service temporarily unavailable",
				2,
			)
		}
		if strings.HasPrefix(apiKey, "sk-gw-") {
			return nil, ErrUnauthorized
		}
	}
	secret := a.ServerMainInternalSecret
	if secret == "" {
		secret = a.ServerMainInternalToken
	}
	if a.ServerMainBaseURL == "" || a.ServerMainInternalKey == "" || secret == "" {
		return nil, ErrUnauthorized
	}
	body, _ := json.Marshal(map[string]string{"apiKey": apiKey})
	const path = "/internal/auth/verify-apikey"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.ServerMainBaseURL+path, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	timestamp := strconv.FormatInt(time.Now().Unix(), 10)
	nonceBytes := make([]byte, 16)
	if _, err := rand.Read(nonceBytes); err != nil {
		return nil, err
	}
	nonce := hex.EncodeToString(nonceBytes)
	bodyHash := sha256.Sum256(body)
	signingText := strings.Join([]string{http.MethodPost, path, timestamp, nonce, hex.EncodeToString(bodyHash[:]), a.ServerMainInternalKey}, "\n")
	mac := hmac.New(sha256.New, []byte(secret))
	_, _ = mac.Write([]byte(signingText))
	req.Header.Set("X-Internal-Key", a.ServerMainInternalKey)
	req.Header.Set("X-Timestamp", timestamp)
	req.Header.Set("X-Nonce", nonce)
	req.Header.Set("X-Signature", hex.EncodeToString(mac.Sum(nil)))

	resp, err := a.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, ErrUnauthorized
	}
	var payload struct {
		Success bool `json:"success"`
		Data    struct {
			User User `json:"user"`
		} `json:"data"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
		return nil, err
	}
	user := payload.Data.User
	if user.ID == "" {
		return nil, ErrUnauthorized
	}
	if user.Source == "" {
		user.Source = "server-main"
	}
	return &user, nil
}

func extractBearer(value string) string {
	fields := strings.Fields(value)
	if len(fields) == 2 && strings.EqualFold(fields[0], "bearer") {
		return fields[1]
	}
	return ""
}

func jwtAlgorithm(token string) string {
	parsed, _, err := new(jwt.Parser).ParseUnverified(token, jwt.MapClaims{})
	if err != nil || parsed == nil {
		return ""
	}
	algorithm, _ := parsed.Header["alg"].(string)
	return algorithm
}

func hasPermission(roles []string, required Permission) bool {
	if required == PermissionPublic {
		return true
	}
	granted := map[Permission]bool{PermissionPublic: true}
	for _, role := range roles {
		for _, permission := range permissionsForRole(role) {
			granted[permission] = true
		}
	}
	return granted[required]
}

func PermissionLevel(roles []string) int {
	level := 0
	for _, role := range roles {
		switch role {
		case "admin", "manager":
			if level < 4 {
				level = 4
			}
		case "operator":
			if level < 3 {
				level = 3
			}
		case "creator":
			if level < 2 {
				level = 2
			}
		case "user":
			if level < 1 {
				level = 1
			}
		}
	}
	return level
}

func permissionsForRole(role string) []Permission {
	switch role {
	case "admin", "manager":
		return []Permission{PermissionPublic, PermissionBasic, PermissionCreat, PermissionPower, PermissionManager}
	case "operator":
		return []Permission{PermissionPublic, PermissionBasic, PermissionCreat, PermissionPower}
	case "creator":
		return []Permission{PermissionPublic, PermissionBasic, PermissionCreat}
	case "user":
		return []Permission{PermissionPublic, PermissionBasic}
	default:
		return []Permission{PermissionPublic}
	}
}

func stringClaim(claims jwt.MapClaims, key string) string {
	value, _ := claims[key].(string)
	return value
}

func stringSliceClaim(claims jwt.MapClaims, key string) []string {
	value := claims[key]
	switch typed := value.(type) {
	case []string:
		return typed
	case []any:
		out := make([]string, 0, len(typed))
		for _, item := range typed {
			if s, ok := item.(string); ok && s != "" {
				out = append(out, s)
			}
		}
		return out
	case string:
		if typed == "" {
			return nil
		}
		return []string{typed}
	default:
		return nil
	}
}
