package api

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/coreos/go-oidc/v3/oidc"
	"golang.org/x/oauth2"
)

const browserSessionCookie = "__Host-fancydns-session"
const loginCookie = "__Host-fancydns-login"

type oidcLogin struct {
	verifier string
	next     string
	expires  time.Time
}

// The registered callback is configuration, never an untrusted Host header.
func (s *Server) browserRedirect() (*url.URL, error) {
	raw := s.auth.oidcRedirectURI
	if raw == "" {
		u, err := url.Parse(s.auth.oidcAudience)
		if err != nil || u.Scheme != "https" || u.Host == "" {
			return nil, errors.New("configure api.auth.oidc_redirect_uri with the HTTPS console callback URL")
		}
		raw = u.Scheme + "://" + u.Host + "/auth/oidc/callback"
	}
	u, err := url.Parse(raw)
	if err != nil || u.Scheme != "https" || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" {
		return nil, errors.New("OIDC callback must be an absolute HTTPS URL without query or fragment")
	}
	return u, nil
}

func (s *Server) cookieRequestAllowed(r *http.Request) bool {
	if r.Method == http.MethodGet || r.Method == http.MethodHead {
		return true
	}
	u, err := s.browserRedirect()
	return err == nil && r.Header.Get("Origin") == u.Scheme+"://"+u.Host
}

func (s *Server) setBrowserCookie(w http.ResponseWriter, name, value string, maxAge int) {
	http.SetCookie(w, &http.Cookie{Name: name, Value: value, Path: "/", Secure: true, HttpOnly: true, SameSite: http.SameSiteLaxMode, MaxAge: maxAge})
}

func (s *Server) browserOAuth(ctx context.Context) (*oauth2.Config, error) {
	if s.auth.mode != "oidc" {
		return nil, errors.New("OIDC login is not enabled")
	}
	redirect, err := s.browserRedirect()
	if err != nil {
		return nil, err
	}
	provider, err := oidc.NewProvider(ctx, s.auth.oidcIssuer)
	if err != nil {
		return nil, err
	}
	endpoint := provider.Endpoint()
	// Discovery is trusted only over TLS; never send codes to a downgraded endpoint.
	for _, raw := range []string{endpoint.AuthURL, endpoint.TokenURL} {
		u, err := url.Parse(raw)
		if err != nil || u.Scheme != "https" || u.Host == "" || u.User != nil {
			return nil, errors.New("OIDC endpoints must use HTTPS")
		}
	}
	endpoint.AuthStyle = oauth2.AuthStyleInParams
	return &oauth2.Config{ClientID: s.auth.oidcClientID, RedirectURL: redirect.String(), Endpoint: endpoint, Scopes: []string{"openid", "profile", "email"}}, nil
}

func safeLoginNext(next string) string {
	if !strings.HasPrefix(next, "/") || strings.HasPrefix(next, "//") || strings.ContainsAny(next, "\\\r\n") || strings.HasPrefix(next, "/auth/") {
		return "/"
	}
	return next
}

func (s *Server) handleOIDCStart(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", "no-store")
	ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
	defer cancel()
	cfg, err := s.browserOAuth(ctx)
	if err != nil {
		writeError(w, http.StatusServiceUnavailable, "could not start identity provider login")
		return
	}
	stateBytes := make([]byte, 32)
	if _, err := rand.Read(stateBytes); err != nil {
		writeError(w, 500, "could not create login state")
		return
	}
	state := base64.RawURLEncoding.EncodeToString(stateBytes)
	verifier := oauth2.GenerateVerifier()
	now := time.Now()
	s.auth.mu.Lock()
	if s.auth.oidcPending == nil {
		s.auth.oidcPending = make(map[string]oidcLogin)
	}
	for key, login := range s.auth.oidcPending {
		if now.After(login.expires) {
			delete(s.auth.oidcPending, key)
		}
	}
	if len(s.auth.oidcPending) >= 1024 {
		s.auth.mu.Unlock()
		writeError(w, 429, "too many pending logins")
		return
	}
	s.auth.oidcPending[state] = oidcLogin{verifier: verifier, next: safeLoginNext(r.URL.Query().Get("next")), expires: now.Add(5 * time.Minute)}
	s.auth.mu.Unlock()
	s.setBrowserCookie(w, loginCookie, state, 300)
	http.Redirect(w, r, cfg.AuthCodeURL(state, oauth2.S256ChallengeOption(verifier)), http.StatusFound)
}

func (s *Server) handleOIDCCallback(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", "no-store")
	if !s.cookieRequestAllowed(r) {
		writeError(w, 403, "invalid login origin")
		return
	}
	var in struct {
		Code  string `json:"code"`
		State string `json:"state"`
	}
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 8192)).Decode(&in); err != nil || in.Code == "" || in.State == "" {
		writeError(w, 400, "code and state are required")
		return
	}
	cookie, err := r.Cookie(loginCookie)
	if err != nil || subtle.ConstantTimeCompare([]byte(cookie.Value), []byte(in.State)) != 1 {
		writeError(w, 400, "login state check failed; start sign in again")
		return
	}
	s.auth.mu.Lock()
	login, ok := s.auth.oidcPending[in.State]
	delete(s.auth.oidcPending, in.State)
	s.auth.mu.Unlock()
	s.setBrowserCookie(w, loginCookie, "", -1)
	if !ok || time.Now().After(login.expires) {
		writeError(w, 400, "login expired; start sign in again")
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 15*time.Second)
	defer cancel()
	cfg, err := s.browserOAuth(ctx)
	if err != nil {
		writeError(w, 502, "could not discover identity provider")
		return
	}
	token, err := cfg.Exchange(ctx, in.Code, oauth2.VerifierOption(login.verifier))
	if err != nil {
		writeError(w, 401, "identity provider rejected login; start sign in again")
		return
	}
	p, err := s.auth.oidcPrincipal(token.AccessToken)
	if err != nil {
		writeError(w, 401, "identity provider returned an invalid access token")
		return
	}
	// Bound the session by both the configured local TTL and the signed JWT expiry.
	parts := strings.Split(token.AccessToken, ".")
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	var claims struct {
		Exp int64 `json:"exp"`
	}
	if err != nil || json.Unmarshal(payload, &claims) != nil || claims.Exp <= time.Now().Unix() {
		writeError(w, 401, "access token has expired")
		return
	}
	local, expires := s.auth.issue(p)
	if local == "" {
		writeError(w, 500, "could not create session")
		return
	}
	if jwtExpiry := time.Unix(claims.Exp, 0); jwtExpiry.Before(expires) {
		expires = jwtExpiry
	}
	s.auth.mu.Lock()
	s.auth.sessions[local] = session{principal: p, expires: expires, accessToken: token.AccessToken}
	s.auth.mu.Unlock()
	s.setBrowserCookie(w, browserSessionCookie, local, int(time.Until(expires).Seconds()))
	writeJSON(w, 200, map[string]string{"next": login.next})
}