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_"
)

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

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

// GenerateSecureToken generates a cryptographically secure random token
func (tg *TokenGenerator) GenerateSecureToken() (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)
	
	// Add prefix for identification
	token := TokenPrefix + tokenData

	return token, nil
}

// HashToken creates a secure hash of the token for storage
func (tg *TokenGenerator) HashToken(token string) (string, error) {
	// Use bcrypt for secure password-like hashing
	hash, err := bcrypt.GenerateFromPassword([]byte(token), bcrypt.DefaultCost)
	if err != nil {
		return "", fmt.Errorf("failed to hash token: %w", 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 {
	if !strings.HasPrefix(token, TokenPrefix) {
		return false
	}
	
	// Remove prefix and check if remaining part is valid base64
	tokenData := strings.TrimPrefix(token, TokenPrefix)
	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) {
	// Generate the token
	token, err := tg.GenerateSecureToken()
	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
}