package main

import (
    "encoding/json"
    "fmt"
    "net/http"
    "os"
    "strings"
    "time"

    jwt "github.com/golang-jwt/jwt/v5"
    "github.com/pquerna/otp/totp"
)

// Simple in-memory storage for TOTPs per user. In production, use a DB.
var totpSecrets = map[string]string{}

// API key (default) - can be overridden by WG_API_KEY env var.
var apiKey string

// JWT secret (default) - can be overridden by JWT_SECRET env var.
var jwtSecret string

func init() {
    apiKey = os.Getenv("WG_API_KEY")
    if apiKey == "" {
        apiKey = "test-api-key" // default for testing
    }
    jwtSecret = os.Getenv("JWT_SECRET")
    if jwtSecret == "" {
        jwtSecret = "secret" // default for testing
    }
    // Preload a test user so JWTs can be generated in tests if needed
    if _, ok := totpSecrets["test"]; !ok {
        // generate a random-looking secret for test user if desired
        // but we won't force it here; user can call /auth/setup-totp?user=test
    }
}

type Claims struct {
    Username string                 `json:"username"`
    jwt.RegisteredClaims
}

// generateJWT creates a short-lived JWT for a given username
func generateJWT(username string) (string, error) {
    claims := &Claims{
        Username: username,
        RegisteredClaims: jwt.RegisteredClaims{
            Issuer:    "wgrplane",
            ExpiresAt: jwt.NewNumericDate(time.Now().Add(15 * time.Minute)),
        },
    }
    token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
    return token.SignedString([]byte(jwtSecret))
}

func parseJWT(tokenStr string) (*Claims, error) {
    token, err := jwt.ParseWithClaims(tokenStr, &Claims{}, func(token *jwt.Token) (interface{}, error) {
        return []byte(jwtSecret), nil
    })
    if err != nil {
        return nil, err
    }
    if claims, ok := token.Claims.(*Claims); ok && token.Valid {
        return claims, nil
    }
    return nil, fmt.Errorf("invalid token")
}

// totpSetupHandler returns a new TOTP secret and provisioning URI for a user
func totpSetupHandler(w http.ResponseWriter, r *http.Request) {
    user := r.URL.Query().Get("user")
    if user == "" {
        http.Error(w, "missing user", http.StatusBadRequest)
        return
    }
    key, err := totp.Generate(totp.GenerateOpts{
        Issuer:      "WGRplane",
        AccountName: user,
    })
    if err != nil {
        http.Error(w, "failed to generate secret", http.StatusInternalServerError)
        return
    }
    secret := key.Secret()
    totpSecrets[user] = secret
    resp := map[string]string{
        "secret":          secret,
        "provisioning_uri": key.URL(),
    }
    w.Header().Set("Content-Type", "application/json")
    _ = json.NewEncoder(w).Encode(resp)
}

// AuthMiddleware protects selected routes with API Key or JWT + optional TOTP
func AuthMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 1) API Key path
        if key := r.Header.Get("wg-rplane-datadunia"); key != "" {
            if subtleConstantTimeEquals(key, apiKey) {
                // If user has a TOTP, require current OTP in header
                // The username isn't known from API Key alone; skip TOTP check here
                next.ServeHTTP(w, r)
                return
            }
        }

        // 2) JWT path
        authHeader := r.Header.Get("Authorization")
        if strings.HasPrefix(authHeader, "Bearer ") {
            token := strings.TrimPrefix(authHeader, "Bearer ")
            claims, err := parseJWT(token)
            if err != nil {
                http.Error(w, "Unauthorized", http.StatusUnauthorized)
                return
            }
            // If user has a TOTp secret, verify it
            if secret, ok := totpSecrets[claims.Username]; ok {
                otp := r.Header.Get("X-TOTP")
                if otp == "" || !totp.Validate(otp, secret) {
                    http.Error(w, "Unauthorized (TOTp)", http.StatusUnauthorized)
                    return
                }
            }
            // Attach username to context for downstream handlers if needed
            next.ServeHTTP(w, r)
            return
        }

        // No valid auth provided
        http.Error(w, "Unauthorized", http.StatusUnauthorized)
    })
}

// helper for constant-time string comparison
func subtleConstantTimeEquals(a, b string) bool {
    if len(a) != len(b) {
        return false
    }
    var diff byte
    for i := 0; i < len(a); i++ {
        diff |= a[i] ^ b[i]
    }
    return diff == 0
}