package dpop

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

	"github.com/golang-jwt/jwt/v5"
)

// KeyPair holds an ES256 key pair for DPoP use.
type KeyPair struct {
	Private *ecdsa.PrivateKey
}

// Generate creates a new ES256 key pair.
func Generate() (*KeyPair, error) {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generate DPoP key: %w", err)
	}
	return &KeyPair{Private: priv}, nil
}

// PublicJWK returns the public key as a JWK map (for embedding in DPoP JWT header).
func (kp *KeyPair) PublicJWK() map[string]interface{} {
	pub := kp.Private.PublicKey
	return map[string]interface{}{
		"kty": "EC",
		"crv": "P-256",
		"x":   base64.RawURLEncoding.EncodeToString(pub.X.Bytes()),
		"y":   base64.RawURLEncoding.EncodeToString(pub.Y.Bytes()),
	}
}

// MarshalPrivate serializes the private key to JWK JSON for session storage.
func (kp *KeyPair) MarshalPrivate() ([]byte, error) {
	pub := kp.Private.PublicKey
	jwk := map[string]interface{}{
		"kty": "EC",
		"crv": "P-256",
		"x":   base64.RawURLEncoding.EncodeToString(pub.X.Bytes()),
		"y":   base64.RawURLEncoding.EncodeToString(pub.Y.Bytes()),
		"d":   base64.RawURLEncoding.EncodeToString(kp.Private.D.Bytes()),
	}
	return json.Marshal(jwk)
}

// UnmarshalPrivate deserializes a private key from JWK JSON.
func UnmarshalPrivate(data []byte) (*KeyPair, error) {
	var jwk map[string]json.RawMessage
	if err := json.Unmarshal(data, &jwk); err != nil {
		return nil, err
	}

	decode := func(key string) ([]byte, error) {
		var s string
		if err := json.Unmarshal(jwk[key], &s); err != nil {
			return nil, err
		}
		return base64.RawURLEncoding.DecodeString(s)
	}

	xb, err := decode("x")
	if err != nil {
		return nil, err
	}
	yb, err := decode("y")
	if err != nil {
		return nil, err
	}
	db, err := decode("d")
	if err != nil {
		return nil, err
	}

	priv := &ecdsa.PrivateKey{
		PublicKey: ecdsa.PublicKey{
			Curve: elliptic.P256(),
			X:     new(big.Int).SetBytes(xb),
			Y:     new(big.Int).SetBytes(yb),
		},
		D: new(big.Int).SetBytes(db),
	}
	return &KeyPair{Private: priv}, nil
}

// randJTI generates a random 16-byte hex string for use as the JWT ID.
func randJTI() string {
	b := make([]byte, 16)
	rand.Read(b)
	return fmt.Sprintf("%x", b)
}

// Proof builds a DPoP proof JWT for the given HTTP method and URL.
// nonce is optional. accessToken, when non-empty, causes an "ath" claim to be
// included (base64url SHA-256 of the token) as required by RFC 9449 §4.2 when
// presenting a DPoP-bound access token.
func (kp *KeyPair) Proof(method, htu, nonce, accessToken string) (string, error) {
	claims := jwt.MapClaims{
		"jti": randJTI(),
		"htm": method,
		"htu": htu,
		"iat": time.Now().Unix(),
	}
	if nonce != "" {
		claims["nonce"] = nonce
	}
	if accessToken != "" {
		h := sha256.Sum256([]byte(accessToken))
		claims["ath"] = base64.RawURLEncoding.EncodeToString(h[:])
	}

	token := jwt.NewWithClaims(jwt.SigningMethodES256, claims)
	token.Header["typ"] = "dpop+jwt"
	token.Header["jwk"] = kp.PublicJWK()

	return token.SignedString(kp.Private)
}