// Package webauthn implements the server side of the WebAuthn (FIDO2) protocol
// for passkey-based user verification. It handles both registration (creating
// a new credential) and assertion (verifying the user), using standard library
// crypto and a minimal CBOR decoder.
//
// Only ECDSA P-256 (ES256, COSE algorithm -7) credentials are supported.
// Attestation is read but not verified — we trust the authenticator attestation
// is honest for this deployment (single-user, self-hosted).
package webauthn

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math/big"
	"time"
)

// Config — the WebAuthn Relying Party parameters. Must match what the browser
// sees (origin = the page's origin, rpID = the effective domain).
type Config struct {
	Origin string // e.g. "https://maven.kvmx.ru"
	RPID   string // e.g. "maven.kvmx.ru"
	RPName string // e.g. "maven"
}

// CredentialLookup is the function signature the RP needs to load a stored
// credential for assertion verification. Returns the COSE public key bytes
// and the current sign count.
type CredentialLookup func(id string) (publicKey []byte, signCount int64, err error)

// CredentialSaver stores a newly registered credential.
type CredentialSaver func(id string, publicKey []byte, userID []byte, userDisplayName string) error

// SignCountUpdater persists an updated sign counter after a successful assertion.
type SignCountUpdater func(id string, count int64) error

// credentialRegistration is the in-memory state for an in-flight registration.
type credentialRegistration struct {
	Challenge string
	UserID    []byte
	CreatedAt time.Time
}

// credentialAssertion is the in-memory state for an in-flight assertion.
type credentialAssertion struct {
	Challenge string
	CreatedAt time.Time
}

// RP — the relying party instance. Holds config and transient challenge state.
// A single-user daemon has one RP.
type RP struct {
	cfg          Config
	regs         map[string]*credentialRegistration
	asserts      map[string]*credentialAssertion
	challengeTTL time.Duration
}

// NewRP creates a relying party with the given WebAuthn configuration.
func NewRP(cfg Config) *RP {
	return &RP{
		cfg:          cfg,
		regs:         make(map[string]*credentialRegistration),
		asserts:      make(map[string]*credentialAssertion),
		challengeTTL: 5 * time.Minute,
	}
}

// CleanExpired removes challenges older than the TTL.
func (rp *RP) CleanExpired() {
	now := time.Now()
	for k, r := range rp.regs {
		if now.Sub(r.CreatedAt) > rp.challengeTTL {
			delete(rp.regs, k)
		}
	}
	for k, a := range rp.asserts {
		if now.Sub(a.CreatedAt) > rp.challengeTTL {
			delete(rp.asserts, k)
		}
	}
}

// CreationOptions returns the PublicKeyCredentialCreationOptions as a
// JSON-serializable map for the browser to create a credential.
func (rp *RP) CreationOptions(userID []byte, userName string) (map[string]any, string, error) {
	challenge := make([]byte, 32)
	if _, err := rand.Read(challenge); err != nil {
		return nil, "", fmt.Errorf("webauthn: challenge: %w", err)
	}
	challengeB64 := base64.RawURLEncoding.EncodeToString(challenge)

	rp.CleanExpired()
	rp.regs[challengeB64] = &credentialRegistration{
		Challenge: challengeB64,
		UserID:    userID,
		CreatedAt: time.Now(),
	}

	return map[string]any{
		"rp": map[string]string{
			"name": rp.cfg.RPName,
			"id":   rp.cfg.RPID,
		},
		"user": map[string]any{
			"id":          base64.RawURLEncoding.EncodeToString(userID),
			"name":        userName,
			"displayName": userName,
		},
		"challenge": challengeB64,
		// ES256 only — parseCOSEKey verifies P-256/ES256 exclusively. Offering
		// RS256 here would let an authenticator register a key we can never
		// verify at assertion time (register-ok, assert-fail forever).
		"pubKeyCredParams": []map[string]any{
			{"type": "public-key", "alg": -7}, // ES256
		},
		"timeout":            60000,
		"attestation":        "none",
		"excludeCredentials": []any{},
	}, challengeB64, nil
}

// FinishRegistration parses the browser's response and stores the credential.
func (rp *RP) FinishRegistration(save CredentialSaver, challengeB64 string, resp map[string]any) (string, error) {
	rp.CleanExpired()
	reg, ok := rp.regs[challengeB64]
	if !ok {
		return "", fmt.Errorf("webauthn: unknown or expired challenge")
	}
	delete(rp.regs, challengeB64)

	credID := rawString(resp, "id")
	if credID == "" {
		return "", fmt.Errorf("webauthn: missing credential id")
	}

	rawResponse, ok := resp["response"].(map[string]any)
	if !ok {
		return "", fmt.Errorf("webauthn: missing response")
	}

	cdjB64 := rawString(rawResponse, "clientDataJSON")
	cdjRaw, err := base64.RawURLEncoding.DecodeString(cdjB64)
	if err != nil {
		return "", fmt.Errorf("webauthn: clientDataJSON: %w", err)
	}
	if err := verifyClientDataBytes(cdjRaw, "webauthn.create", challengeB64, rp.cfg.Origin); err != nil {
		return "", err
	}

	attObjB64 := rawString(rawResponse, "attestationObject")
	attObj, err := base64.RawURLEncoding.DecodeString(attObjB64)
	if err != nil {
		return "", fmt.Errorf("webauthn: attestationObject: %w", err)
	}

	publicKey, err := extractPublicKey(attObj)
	if err != nil {
		return "", fmt.Errorf("webauthn: extract key: %w", err)
	}

	if err := save(credID, publicKey, reg.UserID, "maven user"); err != nil {
		return "", fmt.Errorf("webauthn: store credential: %w", err)
	}

	return credID, nil
}

// AssertionOptions returns a JSON-serializable map for browser authentication.
func (rp *RP) AssertionOptions() (map[string]any, string, error) {
	challenge := make([]byte, 32)
	if _, err := rand.Read(challenge); err != nil {
		return nil, "", fmt.Errorf("webauthn: challenge: %w", err)
	}
	challengeB64 := base64.RawURLEncoding.EncodeToString(challenge)

	rp.CleanExpired()
	rp.asserts[challengeB64] = &credentialAssertion{
		Challenge: challengeB64,
		CreatedAt: time.Now(),
	}

	return map[string]any{
		"challenge":        challengeB64,
		"timeout":          60000,
		"rpId":             rp.cfg.RPID,
		"allowCredentials": []any{},
		"userVerification": "required",
	}, challengeB64, nil
}

// FinishAssertion verifies the browser's assertion response and returns the
// verified credential ID.
func (rp *RP) FinishAssertion(lookup CredentialLookup, updateSignCount SignCountUpdater, challengeB64 string, resp map[string]any) (string, error) {
	rp.CleanExpired()
	if _, ok := rp.asserts[challengeB64]; !ok {
		return "", fmt.Errorf("webauthn: unknown or expired challenge")
	}
	delete(rp.asserts, challengeB64)

	credID := rawString(resp, "id")
	if credID == "" {
		return "", fmt.Errorf("webauthn: missing credential id")
	}

	rawResponse, ok := resp["response"].(map[string]any)
	if !ok {
		return "", fmt.Errorf("webauthn: missing response")
	}

	cdjRaw, err := base64.RawURLEncoding.DecodeString(rawString(rawResponse, "clientDataJSON"))
	if err != nil {
		return "", fmt.Errorf("webauthn: clientDataJSON: %w", err)
	}
	if err := verifyClientDataBytes(cdjRaw, "webauthn.get", challengeB64, rp.cfg.Origin); err != nil {
		return "", err
	}

	authenticatorData, err := base64.RawURLEncoding.DecodeString(rawString(rawResponse, "authenticatorData"))
	if err != nil {
		return "", fmt.Errorf("webauthn: authenticatorData: %w", err)
	}
	// Bind the assertion to this RP and require a verified user gesture. Origin
	// is checked via clientDataJSON above; rpIdHash + flags bind the
	// authenticator half. We requested userVerification:required, so UV must be
	// set — that IS the step-up gesture (biometric/PIN).
	if len(authenticatorData) < 37 {
		return "", fmt.Errorf("webauthn: authenticatorData too short (%d)", len(authenticatorData))
	}
	rpIDHash := sha256.Sum256([]byte(rp.cfg.RPID))
	if !bytes.Equal(authenticatorData[:32], rpIDHash[:]) {
		return "", fmt.Errorf("webauthn: rpIdHash mismatch")
	}
	const flagUP, flagUV = 1 << 0, 1 << 2
	if authenticatorData[32]&flagUP == 0 {
		return "", fmt.Errorf("webauthn: user-present flag not set")
	}
	if authenticatorData[32]&flagUV == 0 {
		return "", fmt.Errorf("webauthn: user-verification flag not set")
	}

	sig, err := base64.RawURLEncoding.DecodeString(rawString(rawResponse, "signature"))
	if err != nil {
		return "", fmt.Errorf("webauthn: signature: %w", err)
	}

	pubKeyBytes, signCount, err := lookup(credID)
	if err != nil {
		return "", fmt.Errorf("webauthn: credential not found: %w", err)
	}

	clientDataHash := sha256.Sum256(cdjRaw)
	sigData := append(authenticatorData, clientDataHash[:]...)

	pubKey, err := parseCOSEKey(pubKeyBytes)
	if err != nil {
		return "", fmt.Errorf("webauthn: parse key: %w", err)
	}
	if !ecdsa.VerifyASN1(pubKey, sigData, sig) {
		return "", fmt.Errorf("webauthn: signature verification failed")
	}

	if len(authenticatorData) >= 37 {
		counter := int64(binary.BigEndian.Uint32(authenticatorData[33:37]))
		if counter > 0 && counter <= signCount {
			return "", fmt.Errorf("webauthn: sign count not greater (old=%d, new=%d)", signCount, counter)
		}
		if counter > 0 {
			if err := updateSignCount(credID, counter); err != nil {
				return "", fmt.Errorf("webauthn: update sign count: %w", err)
			}
		}
	}

	return credID, nil
}

func verifyClientDataBytes(clientDataJSON []byte, expectedType, expectedChallenge, expectedOrigin string) error {
	var cdj struct {
		Type      string `json:"type"`
		Challenge string `json:"challenge"`
		Origin    string `json:"origin"`
	}
	if err := json.Unmarshal(clientDataJSON, &cdj); err != nil {
		return fmt.Errorf("webauthn: parse clientDataJSON: %w", err)
	}
	if cdj.Type != expectedType {
		return fmt.Errorf("webauthn: unexpected type %q", cdj.Type)
	}
	if cdj.Challenge != expectedChallenge {
		return fmt.Errorf("webauthn: challenge mismatch")
	}
	if cdj.Origin != expectedOrigin {
		return fmt.Errorf("webauthn: origin mismatch: %q != %q", cdj.Origin, expectedOrigin)
	}
	return nil
}

func extractPublicKey(attObj []byte) ([]byte, error) {
	v, err := decodeCBOR(attObj)
	if err != nil {
		return nil, fmt.Errorf("webauthn: decode attestation: %w", err)
	}
	m, err := v.MapText()
	if err != nil {
		return nil, fmt.Errorf("webauthn: attestation is not a map: %w", err)
	}

	authDataV, ok := m["authData"]
	if !ok {
		return nil, fmt.Errorf("webauthn: attestation missing authData")
	}
	authData, err := authDataV.Bytes()
	if err != nil {
		return nil, fmt.Errorf("webauthn: authData not bytes: %w", err)
	}
	return extractCOSEKeyFromAuthData(authData)
}

func extractCOSEKeyFromAuthData(authData []byte) ([]byte, error) {
	if len(authData) < 37 {
		return nil, fmt.Errorf("webauthn: authData too short (%d)", len(authData))
	}
	flags := authData[32]
	if flags&(1<<6) == 0 {
		return nil, fmt.Errorf("webauthn: AT flag not set in authData")
	}
	acd := authData[37:]
	if len(acd) < 18 {
		return nil, fmt.Errorf("webauthn: attested credential data too short (%d)", len(acd))
	}
	credIDLen := int(binary.BigEndian.Uint16(acd[16:18]))
	coseKeyOff := 18 + credIDLen
	if coseKeyOff > len(acd) {
		return nil, fmt.Errorf("webauthn: credential ID length %d exceeds data (%d)", credIDLen, len(acd))
	}
	return acd[coseKeyOff:], nil
}

func parseCOSEKey(raw []byte) (*ecdsa.PublicKey, error) {
	v, err := decodeCBOR(raw)
	if err != nil {
		return nil, fmt.Errorf("cose: decode: %w", err)
	}
	m, err := v.Map()
	if err != nil {
		return nil, fmt.Errorf("cose: not a map: %w", err)
	}

	kty, ok := m[1]
	if !ok {
		return nil, fmt.Errorf("cose: missing kty")
	}
	ktyV, err := kty.Int()
	if err != nil {
		return nil, fmt.Errorf("cose: kty: %w", err)
	}
	if ktyV != 2 {
		return nil, fmt.Errorf("cose: unsupported kty %d (expected 2=EC2)", ktyV)
	}

	crv, ok := m[-1]
	if !ok {
		return nil, fmt.Errorf("cose: missing crv")
	}
	crvV, err := crv.Int()
	if err != nil {
		return nil, fmt.Errorf("cose: crv: %w", err)
	}
	if crvV != 1 {
		return nil, fmt.Errorf("cose: unsupported crv %d (expected 1=P-256)", crvV)
	}

	xV, ok := m[-2]
	if !ok {
		return nil, fmt.Errorf("cose: missing x coordinate")
	}
	x, err := xV.Bytes()
	if err != nil {
		return nil, fmt.Errorf("cose: x: %w", err)
	}

	yV, ok := m[-3]
	if !ok {
		return nil, fmt.Errorf("cose: missing y coordinate")
	}
	y, err := yV.Bytes()
	if err != nil {
		return nil, fmt.Errorf("cose: y: %w", err)
	}

	return &ecdsa.PublicKey{
		Curve: elliptic.P256(),
		X:     new(big.Int).SetBytes(x),
		Y:     new(big.Int).SetBytes(y),
	}, nil
}

func rawString(m map[string]any, key string) string {
	v, ok := m[key]
	if !ok {
		return ""
	}
	s, _ := v.(string)
	return s
}