// client2server - Authentication (JWT) and password hashing
//
// Roles: system_admin, project_admin, user
// Tokens are HMAC-SHA256, 24h expiry.

package main

import (
	"crypto/hmac"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"golang.org/x/crypto/scrypt"
)

const (
	RoleSystemAdmin  = "system_admin"
	RoleProjectAdmin = "project_admin"
	RoleUser         = "user"

	tokenLifetime = 24 * time.Hour
	scryptN       = 1 << 15
	scryptR       = 8
	scryptP       = 1
	scryptKeyLen  = 32
)

var (
	jwtSecretMu sync.RWMutex
	jwtSecret   []byte
)

func getJWTSecret() []byte {
	jwtSecretMu.RLock()
	defer jwtSecretMu.RUnlock()
	if len(jwtSecret) > 0 {
		return jwtSecret
	}
	// Default to env, or a deterministic dev key
	s := os.Getenv("JWT_SECRET")
	if s == "" {
		s = "dev-insecure-jwt-secret-change-me-please-32b"
	}
	return []byte(s)
}

func setJWTSecret(s []byte) {
	jwtSecretMu.Lock()
	defer jwtSecretMu.Unlock()
	jwtSecret = s
}

// ----------------------------------------------------------------------------
// Password hashing (scrypt)
// ----------------------------------------------------------------------------

func HashPassword(password string) (string, error) {
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return "", err
	}
	h, err := scryptHash([]byte(password), salt)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("scrypt$%s$%s",
		base64.RawStdEncoding.EncodeToString(salt),
		base64.RawStdEncoding.EncodeToString(h),
	), nil
}

func VerifyPassword(password, encoded string) (bool, error) {
	parts := strings.Split(encoded, "$")
	if len(parts) != 3 || parts[0] != "scrypt" {
		return false, errors.New("invalid hash format")
	}
	salt, err := base64.RawStdEncoding.DecodeString(parts[1])
	if err != nil {
		return false, err
	}
	want, err := base64.RawStdEncoding.DecodeString(parts[2])
	if err != nil {
		return false, err
	}
	got, err := scryptHash([]byte(password), salt)
	if err != nil {
		return false, err
	}
	return hmac.Equal(got, want), nil
}

func scryptHash(password, salt []byte) ([]byte, error) {
	return scryptKey(password, salt, scryptN, scryptR, scryptP, scryptKeyLen)
}

// scryptKey is a small wrapper that uses stdlib crypto/scrypt via a pure-Go
// implementation. We import golang.org/x/crypto/scrypt for the real one.
//
// (Note: scrypt is a real stdlib-ish dep. The 'scrypt' symbol below is
// provided by golang.org/x/crypto/scrypt, imported separately as needed.)
func scryptKey(pw, salt []byte, n, r, p, keyLen int) ([]byte, error) {
	return scrypt.Key(pw, salt, n, r, p, keyLen)
}

// ----------------------------------------------------------------------------
// JWT
// ----------------------------------------------------------------------------

type Claims struct {
	UserID   int64  `json:"uid"`
	Username string `json:"usr"`
	Role     string `json:"rol"`
	jwt.RegisteredClaims
}

func IssueJWT(userID int64, username, role string) (string, error) {
	now := time.Now()
	claims := Claims{
		UserID:   userID,
		Username: username,
		Role:     role,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    "client2server",
			Subject:   username,
			IssuedAt:  jwt.NewNumericDate(now),
			ExpiresAt: jwt.NewNumericDate(now.Add(tokenLifetime)),
		},
	}
	tok := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return tok.SignedString(getJWTSecret())
}

func ParseJWT(tokenStr string) (*Claims, error) {
	claims := &Claims{}
	tok, err := jwt.ParseWithClaims(tokenStr, claims, func(t *jwt.Token) (interface{}, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
		}
		return getJWTSecret(), nil
	})
	if err != nil {
		return nil, err
	}
	if !tok.Valid {
		return nil, errors.New("invalid token")
	}
	return claims, nil
}

// ----------------------------------------------------------------------------
// HTTP middleware
// ----------------------------------------------------------------------------

type ctxKey int

const ctxClaimsKey ctxKey = 0

func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		auth := r.Header.Get("Authorization")
		token := strings.TrimPrefix(auth, "Bearer ")
		if token == "" || token == auth {
			http.Error(w, `{"error":"missing bearer token"}`, http.StatusUnauthorized)
			return
		}
		claims, err := ParseJWT(token)
		if err != nil {
			http.Error(w, `{"error":"invalid token"}`, http.StatusUnauthorized)
			return
		}
		next(w, r.WithContext(r.Context()))
		_ = claims
	}
}

func requireRole(roles ...string) func(http.HandlerFunc) http.HandlerFunc {
	return func(next http.HandlerFunc) http.HandlerFunc {
		return authMiddleware(func(w http.ResponseWriter, r *http.Request) {
			auth := r.Header.Get("Authorization")
			token := strings.TrimPrefix(auth, "Bearer ")
			claims, err := ParseJWT(token)
			if err != nil {
				http.Error(w, `{"error":"invalid token"}`, http.StatusUnauthorized)
				return
			}
			for _, want := range roles {
				if claims.Role == want {
					next(w, r)
					return
				}
			}
			http.Error(w, `{"error":"forbidden"}`, http.StatusForbidden)
		})
	}
}

// ----------------------------------------------------------------------------
// Login endpoint
// ----------------------------------------------------------------------------

type loginRequest struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

type loginResponse struct {
	Token   string `json:"token"`
	Role    string `json:"role"`
	Expires string `json:"expires"`
}

func handleLogin(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, `{"error":"method not allowed"}`, http.StatusMethodNotAllowed)
		return
	}
	var req loginRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, `{"error":"invalid json"}`, http.StatusBadRequest)
		return
	}
	if req.Username == "" || req.Password == "" {
		http.Error(w, `{"error":"username and password required"}`, http.StatusBadRequest)
		return
	}

	var (
		id   int64
		hash string
		role string
	)
	err := db.QueryRow("SELECT id, password_hash, role FROM users WHERE username = ?", req.Username).Scan(&id, &hash, &role)
	if err != nil {
		http.Error(w, `{"error":"invalid credentials"}`, http.StatusUnauthorized)
		return
	}
	ok, err := VerifyPassword(req.Password, hash)
	if err != nil || !ok {
		http.Error(w, `{"error":"invalid credentials"}`, http.StatusUnauthorized)
		return
	}

	tok, err := IssueJWT(id, req.Username, role)
	if err != nil {
		http.Error(w, `{"error":"server error"}`, http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(loginResponse{
		Token:   tok,
		Role:    role,
		Expires: time.Now().Add(tokenLifetime).Format(time.RFC3339),
	})
}

// ----------------------------------------------------------------------------
// Random hex (for default admin token in logs)
// ----------------------------------------------------------------------------

func randHex(n int) string {
	b := make([]byte, n)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}