package server

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"math/big"
	"net/http"
	"net/url"
	"os"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	oidcCodeTTL  = 60 * time.Second
	oidcTokenTTL = 5 * time.Minute
	oidcParamMax = 512
)

// OIDCClient is a relying party allowed to log users in through misaki.
// A user may only sign in if they are approved and granted Service.
type OIDCClient struct {
	ID           string
	SecretSHA256 string // hex sha256 of the client secret
	RedirectURIs []string
	Service      string
	RequireEmail bool
}

type OIDCConfig struct {
	Issuer  string
	KeyFile string
	Clients []OIDCClient
}

type oidcGrant struct {
	clientID, redirectURI, nonce, challenge, scope, amr string
	userID                                              int64
	exp                                                 time.Time
}

type oidcProvider struct {
	issuer  string
	key     *ecdsa.PrivateKey
	kid     string
	clients map[string]OIDCClient

	mu     sync.Mutex
	codes  map[string]oidcGrant
	tokens map[string]oidcGrant
}

func newOIDCProvider(cfg *OIDCConfig) (*oidcProvider, error) {
	if cfg == nil || cfg.Issuer == "" {
		return nil, nil
	}
	if !strings.HasPrefix(cfg.Issuer, "https://") && !strings.HasPrefix(cfg.Issuer, "http://localhost") {
		return nil, fmt.Errorf("oidc issuer must be https: %q", cfg.Issuer)
	}
	key, err := loadOrCreateKey(cfg.KeyFile)
	if err != nil {
		return nil, err
	}
	p := &oidcProvider{
		issuer:  strings.TrimSuffix(cfg.Issuer, "/"),
		key:     key,
		kid:     jwkThumbprint(&key.PublicKey),
		clients: map[string]OIDCClient{},
		codes:   map[string]oidcGrant{},
		tokens:  map[string]oidcGrant{},
	}
	for _, c := range cfg.Clients {
		if c.ID == "" || c.Service == "" || len(c.RedirectURIs) == 0 {
			return nil, fmt.Errorf("oidc client %q: id, service and redirect_uris are required", c.ID)
		}
		if b, err := hex.DecodeString(c.SecretSHA256); err != nil || len(b) != sha256.Size {
			return nil, fmt.Errorf("oidc client %q: secret_sha256 must be 64 hex chars", c.ID)
		}
		for _, u := range c.RedirectURIs {
			if pu, err := url.Parse(u); err != nil || pu.Scheme != "https" || pu.Fragment != "" {
				return nil, fmt.Errorf("oidc client %q: redirect uri must be https without fragment: %q", c.ID, u)
			}
		}
		p.clients[c.ID] = c
	}
	return p, nil
}

func loadOrCreateKey(path string) (*ecdsa.PrivateKey, error) {
	if path == "" {
		return nil, errors.New("oidc key_file is required")
	}
	if b, err := os.ReadFile(path); err == nil {
		blk, _ := pem.Decode(b)
		if blk == nil {
			return nil, fmt.Errorf("oidc key %s: not PEM", path)
		}
		k, err := x509.ParsePKCS8PrivateKey(blk.Bytes)
		if err != nil {
			return nil, fmt.Errorf("oidc key %s: %w", path, err)
		}
		ek, ok := k.(*ecdsa.PrivateKey)
		if !ok || ek.Curve != elliptic.P256() {
			return nil, fmt.Errorf("oidc key %s: want P-256 EC key", path)
		}
		return ek, nil
	} else if !errors.Is(err, os.ErrNotExist) {
		return nil, err
	}
	k, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	der, err := x509.MarshalPKCS8PrivateKey(k)
	if err != nil {
		return nil, err
	}
	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	if err := pem.Encode(f, &pem.Block{Type: "PRIVATE KEY", Bytes: der}); err != nil {
		return nil, err
	}
	return k, f.Close()
}

var b64 = base64.RawURLEncoding

func pad32(n *big.Int) []byte { return n.FillBytes(make([]byte, 32)) }

func (p *oidcProvider) jwk() map[string]string {
	pub := p.key.PublicKey
	return map[string]string{
		"kty": "EC", "crv": "P-256", "use": "sig", "alg": "ES256", "kid": p.kid,
		"x": b64.EncodeToString(pad32(pub.X)), "y": b64.EncodeToString(pad32(pub.Y)),
	}
}

// RFC 7638 thumbprint.
func jwkThumbprint(pub *ecdsa.PublicKey) string {
	j := fmt.Sprintf(`{"crv":"P-256","kty":"EC","x":"%s","y":"%s"}`,
		b64.EncodeToString(pad32(pub.X)), b64.EncodeToString(pad32(pub.Y)))
	sum := sha256.Sum256([]byte(j))
	return b64.EncodeToString(sum[:])
}

func (p *oidcProvider) sign(claims map[string]any) (string, error) {
	hdr, _ := json.Marshal(map[string]string{"alg": "ES256", "typ": "JWT", "kid": p.kid})
	body, err := json.Marshal(claims)
	if err != nil {
		return "", err
	}
	input := b64.EncodeToString(hdr) + "." + b64.EncodeToString(body)
	sum := sha256.Sum256([]byte(input))
	r, s, err := ecdsa.Sign(rand.Reader, p.key, sum[:])
	if err != nil {
		return "", err
	}
	return input + "." + b64.EncodeToString(append(pad32(r), pad32(s)...)), nil
}

func (p *oidcProvider) put(m map[string]oidcGrant, g oidcGrant) (string, error) {
	tok, err := randomToken()
	if err != nil {
		return "", err
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	now := time.Now()
	for k, v := range m {
		if now.After(v.exp) {
			delete(m, k)
		}
	}
	m[hashTok(tok)] = g
	return tok, nil
}

// take returns and removes a grant (codes are single-use).
func (p *oidcProvider) take(m map[string]oidcGrant, tok string, remove bool) (oidcGrant, bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	k := hashTok(tok)
	g, ok := m[k]
	if remove {
		delete(m, k)
	}
	if !ok || time.Now().After(g.exp) {
		return oidcGrant{}, false
	}
	return g, true
}

func hashTok(t string) string {
	sum := sha256.Sum256([]byte(t))
	return hex.EncodeToString(sum[:])
}

func hasScope(scope, want string) bool { return slices.Contains(strings.Fields(scope), want) }

func (s *Server) oidcRoutes(mux *http.ServeMux) {
	if s.oidc == nil {
		return
	}
	mux.HandleFunc("GET /.well-known/openid-configuration", s.handleOIDCDiscovery)
	mux.HandleFunc("GET /oauth/jwks", s.handleOIDCJWKS)
	mux.HandleFunc("GET /oauth/authorize", s.handleOIDCAuthorize)
	mux.HandleFunc("POST /oauth/token", s.handleOIDCToken)
	mux.HandleFunc("GET /oauth/userinfo", s.handleOIDCUserinfo)
	mux.HandleFunc("POST /oauth/userinfo", s.handleOIDCUserinfo)
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "no-store")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(v)
}

func (s *Server) handleOIDCDiscovery(w http.ResponseWriter, r *http.Request) {
	iss := s.oidc.issuer
	writeJSON(w, http.StatusOK, map[string]any{
		"issuer":                                iss,
		"authorization_endpoint":                iss + "/oauth/authorize",
		"token_endpoint":                        iss + "/oauth/token",
		"userinfo_endpoint":                     iss + "/oauth/userinfo",
		"jwks_uri":                              iss + "/oauth/jwks",
		"response_types_supported":              []string{"code"},
		"response_modes_supported":              []string{"query"},
		"grant_types_supported":                 []string{"authorization_code"},
		"subject_types_supported":               []string{"public"},
		"id_token_signing_alg_values_supported": []string{"ES256"},
		"scopes_supported":                      []string{"openid", "profile", "email"},
		"token_endpoint_auth_methods_supported": []string{"client_secret_basic", "client_secret_post"},
		"code_challenge_methods_supported":      []string{"S256"},
		"claims_supported":                      []string{"sub", "iss", "aud", "exp", "iat", "nonce", "amr", "name", "preferred_username", "nickname", "email"},
	})
}

func (s *Server) handleOIDCJWKS(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, map[string]any{"keys": []any{s.oidc.jwk()}})
}

// oidcAllowed reports why user u may not use client c ("" = allowed).
func (s *Server) oidcDenied(uid int64, c OIDCClient) string {
	u, err := s.db.GetUserByID(uid)
	if err != nil || u.Status != "approved" {
		return "account not approved"
	}
	if on, _, err := s.db.GetUserService(u.ID, c.Service); err != nil || !on {
		return "no access to " + c.Service
	}
	if c.RequireEmail {
		if e, _ := s.db.GetUserEmail(u.ID); e == "" {
			return "no email for " + c.Service
		}
	}
	return ""
}

func (s *Server) handleOIDCAuthorize(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	c, ok := s.oidc.clients[q.Get("client_id")]
	redirect := q.Get("redirect_uri")
	if !ok || !slices.Contains(c.RedirectURIs, redirect) {
		http.Error(w, "unknown client or redirect_uri", http.StatusBadRequest)
		return
	}
	state, nonce := q.Get("state"), q.Get("nonce")
	fail := func(code, desc string) {
		v := url.Values{"error": {code}, "error_description": {desc}, "iss": {s.oidc.issuer}}
		if state != "" {
			v.Set("state", state)
		}
		http.Redirect(w, r, redirect+sep(redirect)+v.Encode(), http.StatusFound)
	}
	challenge, method := q.Get("code_challenge"), q.Get("code_challenge_method")
	switch {
	case q.Get("response_type") != "code":
		fail("unsupported_response_type", "only code is supported")
		return
	case !hasScope(q.Get("scope"), "openid"):
		fail("invalid_scope", "openid scope required")
		return
	case len(state) > oidcParamMax || len(nonce) > oidcParamMax || len(challenge) > 128:
		fail("invalid_request", "parameter too long")
		return
	case challenge != "" && method != "S256":
		fail("invalid_request", "only S256 PKCE is supported")
		return
	}

	u, sess := s.session(r)
	if u == nil {
		if q.Get("prompt") == "none" {
			fail("login_required", "not logged in")
			return
		}
		next := s.oidc.issuer + r.URL.RequestURI()
		http.Redirect(w, r, "/login?next="+url.QueryEscape(next), http.StatusSeeOther)
		return
	}
	if why := s.oidcDenied(u.ID, c); why != "" {
		msg := "request access to " + c.Service + " on your dashboard first"
		if strings.HasPrefix(why, "no email") {
			msg = "add a contact email for " + c.Service + " on your dashboard first"
		}
		http.Redirect(w, r, "/dashboard?err="+url.QueryEscape(msg), http.StatusSeeOther)
		return
	}

	code, err := s.oidc.put(s.oidc.codes, oidcGrant{
		clientID: c.ID, redirectURI: redirect, nonce: nonce, challenge: challenge,
		scope: q.Get("scope"), amr: sess.AMR, userID: u.ID, exp: time.Now().Add(oidcCodeTTL),
	})
	if err != nil {
		s.internalError(w)
		return
	}
	v := url.Values{"code": {code}, "iss": {s.oidc.issuer}}
	if state != "" {
		v.Set("state", state)
	}
	http.Redirect(w, r, redirect+sep(redirect)+v.Encode(), http.StatusFound)
}

func sep(u string) string {
	if strings.Contains(u, "?") {
		return "&"
	}
	return "?"
}

func (s *Server) oidcClientAuth(r *http.Request) (OIDCClient, bool) {
	id, secret, basic := r.BasicAuth()
	if basic {
		id, _ = url.QueryUnescape(id)
		secret, _ = url.QueryUnescape(secret)
	} else {
		id, secret = r.PostFormValue("client_id"), r.PostFormValue("client_secret")
	}
	c, ok := s.oidc.clients[id]
	want, _ := hex.DecodeString(c.SecretSHA256)
	got := sha256.Sum256([]byte(secret))
	if subtle.ConstantTimeCompare(got[:], want) != 1 || !ok || secret == "" {
		return OIDCClient{}, false
	}
	return c, true
}

func (s *Server) handleOIDCToken(w http.ResponseWriter, r *http.Request) {
	ip := clientIP(r)
	if s.tokenLimiter.blocked(ip) {
		writeJSON(w, http.StatusTooManyRequests, map[string]string{"error": "slow_down"})
		return
	}
	r.Body = http.MaxBytesReader(w, r.Body, 8<<10)
	bad := func(code int, err string) {
		s.tokenLimiter.record(ip)
		writeJSON(w, code, map[string]string{"error": err})
	}
	c, ok := s.oidcClientAuth(r)
	if !ok {
		w.Header().Set("WWW-Authenticate", `Basic realm="misaki"`)
		bad(http.StatusUnauthorized, "invalid_client")
		return
	}
	if r.PostFormValue("grant_type") != "authorization_code" {
		bad(http.StatusBadRequest, "unsupported_grant_type")
		return
	}
	g, ok := s.oidc.take(s.oidc.codes, r.PostFormValue("code"), true)
	if !ok || g.clientID != c.ID || g.redirectURI != r.PostFormValue("redirect_uri") {
		bad(http.StatusBadRequest, "invalid_grant")
		return
	}
	if g.challenge != "" {
		sum := sha256.Sum256([]byte(r.PostFormValue("code_verifier")))
		if subtle.ConstantTimeCompare([]byte(b64.EncodeToString(sum[:])), []byte(g.challenge)) != 1 {
			bad(http.StatusBadRequest, "invalid_grant")
			return
		}
	}
	if s.oidcDenied(g.userID, c) != "" {
		bad(http.StatusBadRequest, "invalid_grant")
		return
	}

	claims, err := s.oidcClaims(g)
	if err != nil {
		s.internalError(w)
		return
	}
	now := time.Now()
	claims["iss"], claims["aud"] = s.oidc.issuer, c.ID
	claims["iat"], claims["exp"] = now.Unix(), now.Add(oidcTokenTTL).Unix()
	if g.nonce != "" {
		claims["nonce"] = g.nonce
	}
	if g.amr != "" {
		claims["amr"] = strings.Fields(g.amr)
	}
	idToken, err := s.oidc.sign(claims)
	if err != nil {
		s.internalError(w)
		return
	}
	g.exp = now.Add(oidcTokenTTL)
	access, err := s.oidc.put(s.oidc.tokens, g)
	if err != nil {
		s.internalError(w)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"access_token": access,
		"token_type":   "Bearer",
		"expires_in":   int(oidcTokenTTL.Seconds()),
		"id_token":     idToken,
		"scope":        g.scope,
	})
}

func (s *Server) oidcClaims(g oidcGrant) (map[string]any, error) {
	u, err := s.db.GetUserByID(g.userID)
	if err != nil {
		return nil, err
	}
	claims := map[string]any{"sub": strconv.FormatInt(u.ID, 10)}
	if hasScope(g.scope, "profile") {
		claims["preferred_username"] = u.Username
		claims["nickname"] = u.Username
		claims["name"] = u.Username
	}
	if hasScope(g.scope, "email") {
		if e, _ := s.db.GetUserEmail(u.ID); e != "" {
			claims["email"] = e
		}
	}
	return claims, nil
}

func (s *Server) handleOIDCUserinfo(w http.ResponseWriter, r *http.Request) {
	auth := r.Header.Get("Authorization")
	tok, found := strings.CutPrefix(auth, "Bearer ")
	g, ok := s.oidc.take(s.oidc.tokens, tok, false)
	if !found || !ok || s.oidcDenied(g.userID, s.oidc.clients[g.clientID]) != "" {
		w.Header().Set("WWW-Authenticate", `Bearer error="invalid_token"`)
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "invalid_token"})
		return
	}
	claims, err := s.oidcClaims(g)
	if err != nil {
		s.internalError(w)
		return
	}
	writeJSON(w, http.StatusOK, claims)
}