package middleware

import (
	"crypto/rand"
	"encoding/base64"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
)

// CSRF token store (in-memory, consider Redis for production)
type csrfStore struct {
	sync.RWMutex
	tokens map[string]time.Time
}

var store = &csrfStore{
	tokens: make(map[string]time.Time),
}

// Clean expired tokens periodically
func init() {
	go func() {
		ticker := time.NewTicker(10 * time.Minute)
		defer ticker.Stop()
		for range ticker.C {
			store.Lock()
			now := time.Now()
			for token, expiry := range store.tokens {
				if now.After(expiry) {
					delete(store.tokens, token)
				}
			}
			store.Unlock()
		}
	}()
}

// generateToken creates a cryptographically secure random token
func generateToken() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(b), nil
}

// CSRFProtection middleware validates CSRF tokens for state-changing operations
func CSRFProtection() gin.HandlerFunc {
	return func(c *gin.Context) {
		// Skip CSRF for GET, HEAD, OPTIONS (safe methods)
		if c.Request.Method == "GET" || c.Request.Method == "HEAD" || c.Request.Method == "OPTIONS" {
			c.Next()
			return
		}

		// Skip CSRF for public API endpoints that use Bearer tokens
		// (Bearer token auth provides similar protection)
		authHeader := c.GetHeader("Authorization")
		if authHeader != "" && len(authHeader) > 7 && authHeader[:7] == "Bearer " {
			c.Next()
			return
		}

		// Get token from header or form
		token := c.GetHeader("X-CSRF-Token")
		if token == "" {
			token = c.PostForm("csrf_token")
		}

		if token == "" {
			c.JSON(http.StatusForbidden, gin.H{"error": "CSRF token missing"})
			c.Abort()
			return
		}

		// Validate token
		store.RLock()
		expiry, exists := store.tokens[token]
		store.RUnlock()

		if !exists || time.Now().After(expiry) {
			c.JSON(http.StatusForbidden, gin.H{"error": "Invalid or expired CSRF token"})
			c.Abort()
			return
		}

		c.Next()
	}
}

// GetCSRFToken returns a new CSRF token (GET endpoint)
func GetCSRFToken(c *gin.Context) {
	token, err := generateToken()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to generate CSRF token"})
		return
	}

	// Store token with 1 hour expiry
	store.Lock()
	store.tokens[token] = time.Now().Add(1 * time.Hour)
	store.Unlock()

	// Return token in response
	c.JSON(http.StatusOK, gin.H{"csrf_token": token})
}

// CSRFCookie sets CSRF token as a cookie (alternative approach)
func CSRFCookie() gin.HandlerFunc {
	return func(c *gin.Context) {
		// Check if token already exists in cookie
		existingToken, err := c.Cookie("csrf_token")
		if err != nil || existingToken == "" {
			// Generate new token
			token, err := generateToken()
			if err != nil {
				c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to generate CSRF token"})
				c.Abort()
				return
			}

			// Store in memory
			store.Lock()
			store.tokens[token] = time.Now().Add(24 * time.Hour)
			store.Unlock()

			// Set cookie
			c.SetCookie(
				"csrf_token",
				token,
				86400, // 24 hours
				"/",
				"",
				c.Request.TLS != nil || c.GetHeader("X-Forwarded-Proto") == "https", // Secure flag
				false, // HttpOnly = false (needs to be read by JS)
			)

			c.Set("csrf_token", token)
		} else {
			c.Set("csrf_token", existingToken)
		}

		c.Next()
	}
}
