package auth

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
)

// OIDCConfig holds the configuration for the Authelia OIDC client.
type OIDCConfig struct {
	IssuerURL    string
	ClientID     string
	ClientSecret string
	RedirectURL  string
	Domain       string
}

// OIDCHandler handles OIDC authentication flows with Authelia.
type OIDCHandler struct {
	config OIDCConfig
	store  *SessionStore
}

// NewOIDCHandler creates a new OIDC handler.
func NewOIDCHandler(config OIDCConfig, store *SessionStore) *OIDCHandler {
	return &OIDCHandler{
		config: config,
		store:  store,
	}
}

// LoginRedirect redirects the user to Authelia's OIDC authorization endpoint.
func (h *OIDCHandler) LoginRedirect(w http.ResponseWriter, r *http.Request) {
	state := generateToken(16)
	nonce := generateToken(16)

	// PKCE: generate code verifier and challenge
	verifier := generateToken(32)
	challenge := pkceChallenge(verifier)

	// Store state + verifier in cookies (shared across subdomains)
	http.SetCookie(w, &http.Cookie{
		Name:     "oidc_state",
		Value:    state,
		Path:     "/",
		Domain:   h.config.Domain,
		MaxAge:   300,
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
	})
	http.SetCookie(w, &http.Cookie{
		Name:     "oidc_verifier",
		Value:    verifier,
		Path:     "/",
		Domain:   h.config.Domain,
		MaxAge:   300,
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
	})

	authURL := fmt.Sprintf(
		"%s/api/oidc/authorize?response_type=code&client_id=%s&redirect_uri=%s&scope=openid+profile+email&state=%s&nonce=%s&code_challenge=%s&code_challenge_method=S256",
		h.config.IssuerURL,
		url.QueryEscape(h.config.ClientID),
		url.QueryEscape(h.config.RedirectURL),
		state,
		nonce,
		challenge,
	)

	http.Redirect(w, r, authURL, http.StatusFound)
}

// Callback handles the OIDC authorization code callback from Authelia.
func (h *OIDCHandler) Callback(w http.ResponseWriter, r *http.Request) {
	stateCookie, err := r.Cookie("oidc_state")
	if err != nil {
		http.Error(w, "missing state cookie", http.StatusBadRequest)
		return
	}

	// Get state from URL query (GET) or form body (POST)
	stateParam := r.URL.Query().Get("state")
	if stateParam == "" {
		r.ParseForm()
		stateParam = r.Form.Get("state")
	}
	if stateParam == "" || stateParam != stateCookie.Value {
		http.Error(w, "state mismatch", http.StatusForbidden)
		return
	}

	// Get PKCE verifier from cookie
	verifierCookie, _ := r.Cookie("oidc_verifier")
	verifier := ""
	if verifierCookie != nil {
		verifier = verifierCookie.Value
	}

	// Clear state cookies
	http.SetCookie(w, &http.Cookie{Name: "oidc_state", Value: "", Path: "/", MaxAge: -1, HttpOnly: true})
	http.SetCookie(w, &http.Cookie{Name: "oidc_verifier", Value: "", Path: "/", MaxAge: -1, HttpOnly: true})

	// Get code from URL query (GET) or form body (POST)
	code := r.URL.Query().Get("code")
	if code == "" {
		code = r.Form.Get("code")
	}
	if code == "" {
		http.Error(w, "missing authorization code", http.StatusBadRequest)
		return
	}

	// Exchange code for tokens (with PKCE verifier)
	username, err := h.exchangeCode(code, verifier)
	if err != nil {
		http.Error(w, "token exchange failed: "+err.Error(), http.StatusInternalServerError)
		return
	}

	token, err := h.store.CreateSession(username, 60)
	if err != nil {
		http.Error(w, "session creation failed", http.StatusInternalServerError)
		return
	}

	// Set session cookie
	http.SetCookie(w, &http.Cookie{
		Name:     "nextwks_session",
		Value:    token,
		Path:     "/",
		Domain:   h.config.Domain,
		MaxAge:   3600,
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
	})

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

// exchangeCode exchanges an OIDC authorization code for an ID token.
func (h *OIDCHandler) exchangeCode(code, verifier string) (string, error) {
	tokenURL := h.config.IssuerURL + "/api/oidc/token"

	data := url.Values{
		"grant_type":    {"authorization_code"},
		"code":          {code},
		"redirect_uri":  {h.config.RedirectURL},
		"client_id":     {h.config.ClientID},
		"code_verifier": {verifier},
	}

	resp, err := http.PostForm(tokenURL, data)
	if err != nil {
		return "", fmt.Errorf("token request failed: %w", err)
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("token endpoint returned %d: %s", resp.StatusCode, string(body))
	}

	var tokenResp struct {
		IDToken string `json:"id_token"`
	}
	if err := json.Unmarshal(body, &tokenResp); err != nil {
		return "", fmt.Errorf("parse token response: %w", err)
	}

	if tokenResp.IDToken == "" {
		return "", fmt.Errorf("no id_token in response")
	}

	username, err := decodeJWTSub(tokenResp.IDToken)
	if err != nil {
		return "", fmt.Errorf("decode id_token: %w", err)
	}

	return username, nil
}

// decodeJWTSub extracts the "sub" (subject/username) from a JWT without verifying the signature.
func decodeJWTSub(token string) (string, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return "", fmt.Errorf("invalid JWT format")
	}

	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return "", fmt.Errorf("decode JWT payload: %w", err)
	}

	var claims struct {
		Sub               string `json:"sub"`
		PreferredUsername string `json:"preferred_username"`
	}
	if err := json.Unmarshal(payload, &claims); err != nil {
		return "", fmt.Errorf("parse JWT claims: %w", err)
	}

	// Use preferred_username (actual username), fall back to sub (UUID)
	username := claims.PreferredUsername
	if username == "" {
		username = claims.Sub
	}
	if username == "" {
		return "", fmt.Errorf("missing username in id_token")
	}

	return username, nil
}

// pkceChallenge creates a PKCE S256 challenge from a verifier.
func pkceChallenge(verifier string) string {
	h := sha256.Sum256([]byte(verifier))
	return base64.RawURLEncoding.EncodeToString(h[:])
}

// AuthGateMiddleware protects routes behind OIDC authentication.
func (h *OIDCHandler) AuthGateMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, ok := GetUserID(r)
		if !ok {
			h.LoginRedirect(w, r)
			return
		}
		next.ServeHTTP(w, r)
	})
}