package middleware

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"time"

	"aamos-admin/config"
	"aamos-admin/models"
)

type contextKey string

const userKey contextKey = "user"

type jwtHeader struct {
	Alg string `json:"alg"`
	Typ string `json:"typ"`
}

type jwtClaims struct {
	Sub   string     `json:"sub"`
	Email string     `json:"email"`
	Name  string     `json:"name"`
	Role  models.Role `json:"role"`
	Exp   int64      `json:"exp"`
	Iat   int64      `json:"iat"`
}

var (
	errMalformed = errors.New("malformed token")
	errSignature = errors.New("invalid signature")
	errExpired   = errors.New("token expired")
	errAlgorithm = errors.New("unexpected algorithm")
)

// Auth returns a middleware that validates a Bearer JWT (HMAC-SHA256).
// On success the authenticated *models.User is placed in the request context.
// On failure a 401 JSON response is written and the chain is halted.
func Auth(cfg *config.Config) func(http.Handler) http.Handler {
	secret := []byte(cfg.JWTSecret)
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			raw, ok := bearerToken(r)
			if !ok {
				writeUnauthorized(w, "missing bearer token")
				return
			}

			claims, err := parseJWT(raw, secret)
			if err != nil {
				msg := "invalid token"
				if errors.Is(err, errExpired) {
					msg = "token expired"
				}
				writeUnauthorized(w, msg)
				return
			}

			user := &models.User{
				ID:    claims.Sub,
				Email: claims.Email,
				Name:  claims.Name,
				Role:  claims.Role,
			}
			next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), userKey, user)))
		})
	}
}

// UserFromContext returns the authenticated user stored in ctx, or nil.
func UserFromContext(ctx context.Context) *models.User {
	u, _ := ctx.Value(userKey).(*models.User)
	return u
}

func bearerToken(r *http.Request) (string, bool) {
	h := r.Header.Get("Authorization")
	t, found := strings.CutPrefix(h, "Bearer ")
	if !found || t == "" {
		return "", false
	}
	return t, true
}

func parseJWT(token string, secret []byte) (*jwtClaims, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return nil, errMalformed
	}

	headerJSON, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return nil, errMalformed
	}
	var hdr jwtHeader
	if err := json.Unmarshal(headerJSON, &hdr); err != nil {
		return nil, errMalformed
	}
	if hdr.Alg != "HS256" {
		return nil, errAlgorithm
	}

	mac := hmac.New(sha256.New, secret)
	mac.Write([]byte(parts[0] + "." + parts[1]))
	expected := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
	if !hmac.Equal([]byte(expected), []byte(parts[2])) {
		return nil, errSignature
	}

	payloadJSON, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return nil, errMalformed
	}
	var claims jwtClaims
	if err := json.Unmarshal(payloadJSON, &claims); err != nil {
		return nil, errMalformed
	}

	if claims.Exp > 0 && time.Now().Unix() > claims.Exp {
		return nil, errExpired
	}

	return &claims, nil
}

func writeUnauthorized(w http.ResponseWriter, msg string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusUnauthorized)
	json.NewEncoder(w).Encode(struct {
		Error string `json:"error"`
	}{Error: msg})
}