package webserver

import (
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"log"
	"net/http"
	"slices"
	"strings"
)

// sessionKey is the context key for the session token.
type sessionKey struct{}

// tokenFromContext retrieves the bearer token from context.
func tokenFromContext(ctx context.Context) string {
	s, _ := ctx.Value(sessionKey{}).(string)
	return s
}

// contextWithToken stores a bearer token in the context.
func contextWithToken(ctx context.Context, token string) context.Context {
	return context.WithValue(ctx, sessionKey{}, token)
}

// sessionMiddleware checks for a valid mcr_session cookie and adds the
// token to the request context. If no session is present, it redirects
// to the login page.
func (s *Server) sessionMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cookie, err := r.Cookie("mcr_session")
		if err != nil || cookie.Value == "" {
			http.Redirect(w, r, "/login", http.StatusSeeOther)
			return
		}

		ctx := contextWithToken(r.Context(), cookie.Value)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// handleLoginPage renders the login form.
func (s *Server) handleLoginPage(w http.ResponseWriter, r *http.Request) {
	csrf := s.generateCSRFToken(w)
	s.templates.render(w, "login", map[string]any{
		"CSRFToken": csrf,
		"Session":   false,
	})
}

// handleLoginSubmit processes the login form.
func (s *Server) handleLoginSubmit(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, 1<<20) // 1 MiB limit
	if err := r.ParseForm(); err != nil {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}

	if !s.validateCSRFToken(r) {
		csrf := s.generateCSRFToken(w)
		s.templates.render(w, "login", map[string]any{
			"Error":     "Invalid or expired form submission. Please try again.",
			"CSRFToken": csrf,
			"Session":   false,
		})
		return
	}

	username := r.FormValue("username")
	password := r.FormValue("password")

	if username == "" || password == "" {
		csrf := s.generateCSRFToken(w)
		s.templates.render(w, "login", map[string]any{
			"Error":     "Username and password are required.",
			"CSRFToken": csrf,
			"Session":   false,
		})
		return
	}

	token, _, err := s.loginFn(username, password)
	if err != nil {
		log.Printf("login failed for user %q: %v", username, err)
		csrf := s.generateCSRFToken(w)
		s.templates.render(w, "login", map[string]any{
			"Error":     "Invalid username or password.",
			"CSRFToken": csrf,
			"Session":   false,
		})
		return
	}

	// Validate the token to check roles. Guest accounts are not
	// permitted to use the web interface.
	roles, err := s.validateFn(token)
	if err != nil {
		log.Printf("login token validation failed for user %q: %v", username, err)
		csrf := s.generateCSRFToken(w)
		s.templates.render(w, "login", map[string]any{
			"Error":     "Login failed. Please try again.",
			"CSRFToken": csrf,
			"Session":   false,
		})
		return
	}

	if slices.Contains(roles, "guest") {
		log.Printf("login denied for guest user %q", username)
		csrf := s.generateCSRFToken(w)
		s.templates.render(w, "login", map[string]any{
			"Error":     "Guest accounts are not permitted to access the web interface.",
			"CSRFToken": csrf,
			"Session":   false,
		})
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     "mcr_session",
		Value:    token,
		Path:     "/",
		HttpOnly: true,
		Secure:   true,
		SameSite: http.SameSiteStrictMode,
	})

	http.Redirect(w, r, "/", http.StatusSeeOther)
}

// handleLogout clears the session and redirects to login.
func (s *Server) handleLogout(w http.ResponseWriter, r *http.Request) {
	http.SetCookie(w, &http.Cookie{
		Name:     "mcr_session",
		Value:    "",
		Path:     "/",
		MaxAge:   -1,
		HttpOnly: true,
		Secure:   true,
		SameSite: http.SameSiteStrictMode,
	})
	http.Redirect(w, r, "/login", http.StatusSeeOther)
}

// generateCSRFToken creates a random token, signs it with HMAC, stores
// the signed value in a cookie, and returns the token for form embedding.
func (s *Server) generateCSRFToken(w http.ResponseWriter) string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		// Crypto RNG failure is fatal; this should never happen.
		log.Printf("csrf: failed to generate random bytes: %v", err)
		return ""
	}

	token := hex.EncodeToString(b)
	sig := s.signCSRF(token)
	cookieVal := token + "." + sig

	http.SetCookie(w, &http.Cookie{
		Name:     "csrf_token",
		Value:    cookieVal,
		Path:     "/",
		HttpOnly: true,
		Secure:   true,
		SameSite: http.SameSiteStrictMode,
	})

	return token
}

// validateCSRFToken verifies the form _csrf field matches the cookie and
// the HMAC signature is valid.
func (s *Server) validateCSRFToken(r *http.Request) bool {
	formToken := r.FormValue("_csrf")
	if formToken == "" {
		return false
	}

	cookie, err := r.Cookie("csrf_token")
	if err != nil || cookie.Value == "" {
		return false
	}

	parts := strings.SplitN(cookie.Value, ".", 2)
	if len(parts) != 2 {
		return false
	}

	cookieToken := parts[0]
	cookieSig := parts[1]

	// Verify the form token matches the cookie token.
	if !hmac.Equal([]byte(formToken), []byte(cookieToken)) {
		return false
	}

	// Verify the HMAC signature.
	expectedSig := s.signCSRF(cookieToken)
	return hmac.Equal([]byte(cookieSig), []byte(expectedSig))
}

// signCSRF computes an HMAC-SHA256 signature for a CSRF token.
func (s *Server) signCSRF(token string) string {
	mac := hmac.New(sha256.New, s.csrfKey)
	mac.Write([]byte(token))
	return hex.EncodeToString(mac.Sum(nil))
}