package crypto

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"strings"
	"time"

	"golang.org/x/crypto/bcrypt"
)

const (
	// TokenLength defines the length of generated tokens in bytes
	TokenLength = 32
	// TokenPrefix is prepended to all tokens for identification
	TokenPrefix = "kms_"
	// BcryptCost defines the bcrypt cost for 2025 security standards (minimum 14)
	BcryptCost = 14
)

// TokenGenerator provides secure token generation and validation
type TokenGenerator struct {
	hmacKey    []byte
	bcryptCost int
}

// NewTokenGenerator creates a new token generator with the provided HMAC key
func NewTokenGenerator(hmacKey string) *TokenGenerator {
	return &TokenGenerator{
		hmacKey:    []byte(hmacKey),
		bcryptCost: BcryptCost,
	}
}

// NewTokenGeneratorWithCost creates a new token generator with custom bcrypt cost
func NewTokenGeneratorWithCost(hmacKey string, bcryptCost int) *TokenGenerator {
	// Validate bcrypt cost (must be between 4 and 31)
	if bcryptCost < 4 {
		bcryptCost = 4
	} else if bcryptCost > 31 {
		bcryptCost = 31
	}
	
	// Warn if cost is too low for production
	if bcryptCost < 12 {
		// This should log a warning, but we don't have logger here
		// In a real implementation, you'd pass a logger or use a global one
	}
	
	return &TokenGenerator{
		hmacKey:    []byte(hmacKey),
		bcryptCost: bcryptCost,
	}
}

// GenerateSecureToken generates a cryptographically secure random token
func (tg *TokenGenerator) GenerateSecureToken() (string, error) {
	return tg.GenerateSecureTokenWithPrefix("", "")
}

// GenerateSecureTokenWithPrefix generates a cryptographically secure random token with custom prefix
func (tg *TokenGenerator) GenerateSecureTokenWithPrefix(appPrefix string, tokenType string) (string, error) {
	// Generate random bytes
	tokenBytes := make([]byte, TokenLength)
	if _, err := rand.Read(tokenBytes); err != nil {
		return "", fmt.Errorf("failed to generate random token: %w", err)
	}

	// Encode to base64 for safe transmission
	tokenData := base64.URLEncoding.EncodeToString(tokenBytes)
	
	// Build prefix based on application and token type
	var prefix string
	if appPrefix != "" {
		// Use custom application prefix
		if tokenType == "user" {
			prefix = appPrefix + "UT-" // User Token
		} else {
			prefix = appPrefix + "T-"  // Static Token
		}
	} else {
		// Use default prefix
		prefix = TokenPrefix
	}
	
	token := prefix + tokenData

	return token, nil
}

// HashToken creates a secure hash of the token for storage
func (tg *TokenGenerator) HashToken(token string) (string, error) {
	// Use bcrypt with configured cost
	hash, err := bcrypt.GenerateFromPassword([]byte(token), tg.bcryptCost)
	if err != nil {
		return "", fmt.Errorf("failed to hash token with bcrypt cost %d: %w", tg.bcryptCost, err)
	}

	return string(hash), nil
}

// VerifyToken verifies a token against its stored hash
func (tg *TokenGenerator) VerifyToken(token, hash string) bool {
	err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(token))
	return err == nil
}

// GenerateHMACKey generates a new HMAC key for token signing
func GenerateHMACKey() (string, error) {
	key := make([]byte, 32) // 256-bit key
	if _, err := rand.Read(key); err != nil {
		return "", fmt.Errorf("failed to generate HMAC key: %w", err)
	}

	return hex.EncodeToString(key), nil
}

// SignToken creates an HMAC signature for a token
func (tg *TokenGenerator) SignToken(token string, timestamp time.Time) string {
	h := hmac.New(sha256.New, tg.hmacKey)
	h.Write([]byte(token))
	h.Write([]byte(timestamp.Format(time.RFC3339)))
	
	signature := h.Sum(nil)
	return hex.EncodeToString(signature)
}

// VerifyTokenSignature verifies an HMAC signature for a token
func (tg *TokenGenerator) VerifyTokenSignature(token, signature string, timestamp time.Time) bool {
	expectedSignature := tg.SignToken(token, timestamp)
	return hmac.Equal([]byte(signature), []byte(expectedSignature))
}

// ExtractTokenFromHeader extracts a token from an Authorization header
func ExtractTokenFromHeader(authHeader string) string {
	// Support both "Bearer token" and "token" formats
	if strings.HasPrefix(authHeader, "Bearer ") {
		return strings.TrimPrefix(authHeader, "Bearer ")
	}
	return authHeader
}

// IsValidTokenFormat checks if a token has the expected format
func IsValidTokenFormat(token string) bool {
	return IsValidTokenFormatWithPrefix(token, "")
}

// IsValidTokenFormatWithPrefix checks if a token has the expected format with custom prefix
func IsValidTokenFormatWithPrefix(token string, expectedPrefix string) bool {
	var prefix string
	if expectedPrefix != "" {
		prefix = expectedPrefix
	} else {
		prefix = TokenPrefix
	}
	
	if !strings.HasPrefix(token, prefix) {
		// If expected prefix doesn't match, check if it's a valid token with any custom prefix
		if expectedPrefix == "" {
			// Check for custom prefix pattern: 2-4 uppercase letters + "T-" or "UT-"
			if len(token) < 6 { // minimum: "ABT-" + some data
				return false
			}
			
			// Look for T- or UT- suffix in the first part
			dashIndex := strings.Index(token, "-")
			if dashIndex < 2 || dashIndex > 6 { // 2-4 chars + "T" or "UT"
				// Not a custom prefix, check default
				if !strings.HasPrefix(token, TokenPrefix) {
					return false
				}
				prefix = TokenPrefix
			} else {
				prefixPart := token[:dashIndex+1]
				if !strings.HasSuffix(prefixPart, "T-") && !strings.HasSuffix(prefixPart, "UT-") {
					if !strings.HasPrefix(token, TokenPrefix) {
						return false
					}
					prefix = TokenPrefix
				} else {
					prefix = prefixPart
				}
			}
		} else {
			return false
		}
	}
	
	// Remove prefix and check if remaining part is valid base64
	tokenData := strings.TrimPrefix(token, prefix)
	if len(tokenData) == 0 {
		return false
	}
	
	// Try to decode base64
	_, err := base64.URLEncoding.DecodeString(tokenData)
	return err == nil
}

// TokenInfo holds information about a token
type TokenInfo struct {
	Token     string
	Hash      string
	Signature string
	CreatedAt time.Time
}

// GenerateTokenWithInfo generates a complete token with hash and signature
func (tg *TokenGenerator) GenerateTokenWithInfo() (*TokenInfo, error) {
	return tg.GenerateTokenWithInfoAndPrefix("", "")
}

// GenerateTokenWithInfoAndPrefix generates a complete token with hash, signature, and custom prefix
func (tg *TokenGenerator) GenerateTokenWithInfoAndPrefix(appPrefix string, tokenType string) (*TokenInfo, error) {
	// Generate the token
	token, err := tg.GenerateSecureTokenWithPrefix(appPrefix, tokenType)
	if err != nil {
		return nil, fmt.Errorf("failed to generate token: %w", err)
	}

	// Hash the token for storage
	hash, err := tg.HashToken(token)
	if err != nil {
		return nil, fmt.Errorf("failed to hash token: %w", err)
	}

	// Create timestamp and signature
	now := time.Now()
	signature := tg.SignToken(token, now)

	return &TokenInfo{
		Token:     token,
		Hash:      hash,
		Signature: signature,
		CreatedAt: now,
	}, nil
}

// ValidateTokenInfo validates a complete token with all its components
func (tg *TokenGenerator) ValidateTokenInfo(token, hash, signature string, createdAt time.Time) error {
	// Check token format
	if !IsValidTokenFormat(token) {
		return fmt.Errorf("invalid token format")
	}

	// Verify token against hash
	if !tg.VerifyToken(token, hash) {
		return fmt.Errorf("token verification failed")
	}

	// Verify signature
	if !tg.VerifyTokenSignature(token, signature, createdAt) {
		return fmt.Errorf("token signature verification failed")
	}

	return nil
}