package acme

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
)

// JWSFlat is the wire format of a flattened JWS (RFC 7515 §7.2.2).
type JWSFlat struct {
	Protected string `json:"protected"`
	Payload   string `json:"payload"`
	Signature string `json:"signature"`
}

// JWSHeader is the decoded JOSE protected header from an ACME request.
type JWSHeader struct {
	Alg   string          `json:"alg"`
	Nonce string          `json:"nonce"`
	URL   string          `json:"url"`
	JWK   json.RawMessage `json:"jwk,omitempty"` // present for new-account / new-order with new key
	KID   string          `json:"kid,omitempty"` // present for subsequent requests
}

// ParsedJWS is the result of parsing a JWS request body.
// Signature verification is NOT performed here; call VerifyJWS separately.
type ParsedJWS struct {
	Header        JWSHeader
	Payload       []byte // decoded payload bytes; empty-string payload decodes to nil
	SigningInput   []byte // protected + "." + payload (ASCII; used for signature verification)
	RawSignature  []byte // decoded signature bytes
	RawBody       JWSFlat
}

// ParseJWS decodes the flattened JWS from body without verifying the signature.
// Callers must call VerifyJWS after resolving the public key.
func ParseJWS(body []byte) (*ParsedJWS, error) {
	var flat JWSFlat
	if err := json.Unmarshal(body, &flat); err != nil {
		return nil, fmt.Errorf("invalid JWS: %w", err)
	}

	headerBytes, err := base64.RawURLEncoding.DecodeString(flat.Protected)
	if err != nil {
		return nil, fmt.Errorf("invalid JWS protected header encoding: %w", err)
	}
	var header JWSHeader
	if err := json.Unmarshal(headerBytes, &header); err != nil {
		return nil, fmt.Errorf("invalid JWS protected header: %w", err)
	}

	var payload []byte
	if flat.Payload != "" {
		payload, err = base64.RawURLEncoding.DecodeString(flat.Payload)
		if err != nil {
			return nil, fmt.Errorf("invalid JWS payload encoding: %w", err)
		}
	}

	sig, err := base64.RawURLEncoding.DecodeString(flat.Signature)
	if err != nil {
		return nil, fmt.Errorf("invalid JWS signature encoding: %w", err)
	}

	// Signing input is the ASCII bytes of "protected.payload" (not decoded).
	signingInput := []byte(flat.Protected + "." + flat.Payload)

	return &ParsedJWS{
		Header:       header,
		Payload:      payload,
		SigningInput:  signingInput,
		RawSignature: sig,
		RawBody:      flat,
	}, nil
}

// VerifyJWS verifies the signature of a parsed JWS against the given public key.
// Supported algorithms: ES256, ES384, ES512 (ECDSA), RS256 (RSA-PKCS1v15).
func VerifyJWS(parsed *ParsedJWS, pubKey crypto.PublicKey) error {
	switch parsed.Header.Alg {
	case "ES256":
		return verifyECDSA(parsed.SigningInput, parsed.RawSignature, pubKey, crypto.SHA256)
	case "ES384":
		return verifyECDSA(parsed.SigningInput, parsed.RawSignature, pubKey, crypto.SHA384)
	case "ES512":
		return verifyECDSA(parsed.SigningInput, parsed.RawSignature, pubKey, crypto.SHA512)
	case "RS256":
		return verifyRSA(parsed.SigningInput, parsed.RawSignature, pubKey)
	default:
		return fmt.Errorf("unsupported algorithm: %s", parsed.Header.Alg)
	}
}

func verifyECDSA(input, sig []byte, pubKey crypto.PublicKey, hash crypto.Hash) error {
	ecKey, ok := pubKey.(*ecdsa.PublicKey)
	if !ok {
		return errors.New("expected ECDSA public key")
	}

	var digest []byte
	switch hash {
	case crypto.SHA256:
		h := sha256.Sum256(input)
		digest = h[:]
	case crypto.SHA384:
		h := sha512.Sum384(input)
		digest = h[:]
	case crypto.SHA512:
		h := sha512.Sum512(input)
		digest = h[:]
	default:
		return errors.New("unsupported hash")
	}

	// ECDSA JWS signatures are the concatenation of R and S, each padded to
	// the curve's byte length (RFC 7518 §3.4).
	keyBytes := (ecKey.Curve.Params().BitSize + 7) / 8
	if len(sig) != 2*keyBytes {
		return fmt.Errorf("ECDSA signature has wrong length: got %d, want %d", len(sig), 2*keyBytes)
	}
	r := new(big.Int).SetBytes(sig[:keyBytes])
	s := new(big.Int).SetBytes(sig[keyBytes:])
	if !ecdsa.Verify(ecKey, digest, r, s) {
		return errors.New("ECDSA signature verification failed")
	}
	return nil
}

func verifyRSA(input, sig []byte, pubKey crypto.PublicKey) error {
	rsaKey, ok := pubKey.(*rsa.PublicKey)
	if !ok {
		return errors.New("expected RSA public key")
	}
	digest := sha256.Sum256(input)
	return rsa.VerifyPKCS1v15(rsaKey, crypto.SHA256, digest[:], sig)
}

// VerifyEAB verifies the External Account Binding inner JWS (RFC 8555 §7.3.4).
// The inner JWS is a MAC over the account JWK bytes using the EAB HMAC key.
// kid must match the KID in the inner JWS header.
// accountJWK is the canonical JSON of the new account's public key (the outer JWK).
func VerifyEAB(eabJWS json.RawMessage, kid string, hmacKey, accountJWK []byte) error {
	parsed, err := ParseJWS(eabJWS)
	if err != nil {
		return fmt.Errorf("invalid EAB JWS: %w", err)
	}
	if parsed.Header.Alg != "HS256" {
		return fmt.Errorf("EAB must use HS256, got %s", parsed.Header.Alg)
	}
	if parsed.Header.KID != kid {
		return fmt.Errorf("EAB kid mismatch: got %q, want %q", parsed.Header.KID, kid)
	}
	// The EAB payload must be the account JWK (base64url-encoded).
	expectedPayload := base64.RawURLEncoding.EncodeToString(accountJWK)
	if parsed.RawBody.Payload != expectedPayload {
		// Canonical form may differ; compare decoded bytes.
		if string(parsed.Payload) != string(accountJWK) {
			return errors.New("EAB payload does not match account JWK")
		}
	}
	// Verify HMAC-SHA256.
	mac := hmac.New(sha256.New, hmacKey)
	mac.Write(parsed.SigningInput)
	expected := mac.Sum(nil)
	sig, err := base64.RawURLEncoding.DecodeString(parsed.RawBody.Signature)
	if err != nil {
		return fmt.Errorf("invalid EAB signature encoding: %w", err)
	}
	if !hmac.Equal(sig, expected) {
		return errors.New("EAB HMAC verification failed")
	}
	return nil
}

// JWK types for JSON unmarshaling.
type jwkEC struct {
	Kty string `json:"kty"`
	Crv string `json:"crv"`
	X   string `json:"x"`
	Y   string `json:"y"`
}

type jwkRSA struct {
	Kty string `json:"kty"`
	N   string `json:"n"`
	E   string `json:"e"`
}

// ParseJWK parses a JWK (JSON Web Key) into a Go public key.
// Supports EC keys (P-256, P-384, P-521) and RSA keys.
func ParseJWK(jwk json.RawMessage) (crypto.PublicKey, error) {
	var kty struct {
		Kty string `json:"kty"`
	}
	if err := json.Unmarshal(jwk, &kty); err != nil {
		return nil, fmt.Errorf("invalid JWK: %w", err)
	}
	switch kty.Kty {
	case "EC":
		var key jwkEC
		if err := json.Unmarshal(jwk, &key); err != nil {
			return nil, fmt.Errorf("invalid EC JWK: %w", err)
		}
		var curve elliptic.Curve
		switch key.Crv {
		case "P-256":
			curve = elliptic.P256()
		case "P-384":
			curve = elliptic.P384()
		case "P-521":
			curve = elliptic.P521()
		default:
			return nil, fmt.Errorf("unsupported EC curve: %s", key.Crv)
		}
		xBytes, err := base64.RawURLEncoding.DecodeString(key.X)
		if err != nil {
			return nil, fmt.Errorf("invalid EC JWK x: %w", err)
		}
		yBytes, err := base64.RawURLEncoding.DecodeString(key.Y)
		if err != nil {
			return nil, fmt.Errorf("invalid EC JWK y: %w", err)
		}
		return &ecdsa.PublicKey{
			Curve: curve,
			X:     new(big.Int).SetBytes(xBytes),
			Y:     new(big.Int).SetBytes(yBytes),
		}, nil

	case "RSA":
		var key jwkRSA
		if err := json.Unmarshal(jwk, &key); err != nil {
			return nil, fmt.Errorf("invalid RSA JWK: %w", err)
		}
		nBytes, err := base64.RawURLEncoding.DecodeString(key.N)
		if err != nil {
			return nil, fmt.Errorf("invalid RSA JWK n: %w", err)
		}
		eBytes, err := base64.RawURLEncoding.DecodeString(key.E)
		if err != nil {
			return nil, fmt.Errorf("invalid RSA JWK e: %w", err)
		}
		// e is a big-endian integer, typically 65537.
		eInt := 0
		for _, b := range eBytes {
			eInt = eInt<<8 | int(b)
		}
		return &rsa.PublicKey{
			N: new(big.Int).SetBytes(nBytes),
			E: eInt,
		}, nil

	default:
		return nil, fmt.Errorf("unsupported key type: %s", kty.Kty)
	}
}

// ThumbprintJWK computes the RFC 7638 JWK thumbprint (SHA-256) of a public key JWK.
// Used to compute the key authorization for challenges.
func ThumbprintJWK(jwk json.RawMessage) (string, error) {
	var kty struct {
		Kty string `json:"kty"`
	}
	if err := json.Unmarshal(jwk, &kty); err != nil {
		return "", fmt.Errorf("invalid JWK: %w", err)
	}

	// RFC 7638 §3: thumbprint input is a JSON object with only the required
	// members in lexicographic order, with no whitespace.
	var canonical []byte
	var err error
	switch kty.Kty {
	case "EC":
		var key jwkEC
		if err = json.Unmarshal(jwk, &key); err != nil {
			return "", err
		}
		canonical, err = json.Marshal(map[string]string{
			"crv": key.Crv,
			"kty": key.Kty,
			"x":   key.X,
			"y":   key.Y,
		})
	case "RSA":
		var key jwkRSA
		if err = json.Unmarshal(jwk, &key); err != nil {
			return "", err
		}
		canonical, err = json.Marshal(map[string]string{
			"e":   key.E,
			"kty": key.Kty,
			"n":   key.N,
		})
	default:
		return "", fmt.Errorf("unsupported key type: %s", kty.Kty)
	}
	if err != nil {
		return "", err
	}
	digest := sha256.Sum256(canonical)
	return base64.RawURLEncoding.EncodeToString(digest[:]), nil
}

// KeyAuthorization computes the key authorization string for a challenge token
// and account JWK: token + "." + base64url(SHA-256(JWK thumbprint)).
func KeyAuthorization(token string, accountJWK json.RawMessage) (string, error) {
	thumbprint, err := ThumbprintJWK(accountJWK)
	if err != nil {
		return "", err
	}
	return token + "." + thumbprint, nil
}