package acme

import (
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/asn1"
	"encoding/base64"
	"encoding/json"
	"io"
	"log/slog"
	"math/big"
	"strings"
	"sync"
	"testing"
	"time"

	"git.wntrmute.dev/mc/metacrypt/internal/barrier"
)

// memBarrier is an in-memory barrier for testing.
type memBarrier struct {
	data map[string][]byte
	mu   sync.RWMutex
}

func newMemBarrier() *memBarrier {
	return &memBarrier{data: make(map[string][]byte)}
}

func (m *memBarrier) Unseal(_ []byte) error { return nil }
func (m *memBarrier) Seal() error           { return nil }
func (m *memBarrier) IsSealed() bool        { return false }

func (m *memBarrier) Get(_ context.Context, path string) ([]byte, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	v, ok := m.data[path]
	if !ok {
		return nil, barrier.ErrNotFound
	}
	cp := make([]byte, len(v))
	copy(cp, v)
	return cp, nil
}

func (m *memBarrier) Put(_ context.Context, path string, value []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	cp := make([]byte, len(value))
	copy(cp, value)
	m.data[path] = cp
	return nil
}

func (m *memBarrier) Delete(_ context.Context, path string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.data, path)
	return nil
}

func (m *memBarrier) List(_ context.Context, prefix string) ([]string, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	var paths []string
	for k := range m.data {
		if strings.HasPrefix(k, prefix) {
			paths = append(paths, strings.TrimPrefix(k, prefix))
		}
	}
	return paths, nil
}

// generateES256Key generates an ECDSA P-256 key pair and returns the private
// key along with a JWK JSON representation of the public key.
func generateES256Key(t *testing.T) (*ecdsa.PrivateKey, json.RawMessage) {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("generate ES256 key: %v", err)
	}
	byteLen := (key.Curve.Params().BitSize + 7) / 8
	xBytes := key.PublicKey.X.Bytes()
	yBytes := key.PublicKey.Y.Bytes()
	// Pad to curve byte length.
	for len(xBytes) < byteLen {
		xBytes = append([]byte{0}, xBytes...)
	}
	for len(yBytes) < byteLen {
		yBytes = append([]byte{0}, yBytes...)
	}
	jwk, err := json.Marshal(map[string]string{
		"kty": "EC",
		"crv": "P-256",
		"x":   base64.RawURLEncoding.EncodeToString(xBytes),
		"y":   base64.RawURLEncoding.EncodeToString(yBytes),
	})
	if err != nil {
		t.Fatalf("marshal ES256 JWK: %v", err)
	}
	return key, json.RawMessage(jwk)
}

// generateRSA2048Key generates an RSA 2048-bit key pair and returns the
// private key along with a JWK JSON representation of the public key.
func generateRSA2048Key(t *testing.T) (*rsa.PrivateKey, json.RawMessage) {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate RSA 2048 key: %v", err)
	}
	nBytes := key.PublicKey.N.Bytes()
	eBytes := big.NewInt(int64(key.PublicKey.E)).Bytes()
	jwk, err := json.Marshal(map[string]string{
		"kty": "RSA",
		"n":   base64.RawURLEncoding.EncodeToString(nBytes),
		"e":   base64.RawURLEncoding.EncodeToString(eBytes),
	})
	if err != nil {
		t.Fatalf("marshal RSA JWK: %v", err)
	}
	return key, json.RawMessage(jwk)
}

// ecdsaSigASN1 is used to decode an ASN.1 ECDSA signature into R and S.
type ecdsaSigASN1 struct {
	R *big.Int
	S *big.Int
}

// signJWS creates a valid JWS in flattened serialization.
func signJWS(t *testing.T, key crypto.Signer, alg string, header JWSHeader, payload []byte) []byte {
	t.Helper()

	headerJSON, err := json.Marshal(header)
	if err != nil {
		t.Fatalf("marshal JWS header: %v", err)
	}
	protected := base64.RawURLEncoding.EncodeToString(headerJSON)

	var encodedPayload string
	if payload != nil {
		encodedPayload = base64.RawURLEncoding.EncodeToString(payload)
	}

	signingInput := []byte(protected + "." + encodedPayload)

	var sig []byte
	switch alg {
	case "ES256", "ES384":
		var hashFunc crypto.Hash
		if alg == "ES256" {
			hashFunc = crypto.SHA256
		} else {
			hashFunc = crypto.SHA384
		}
		h := hashFunc.New()
		h.Write(signingInput)
		digest := h.Sum(nil)

		derSig, err := ecdsa.SignASN1(rand.Reader, key.(*ecdsa.PrivateKey), digest)
		if err != nil {
			t.Fatalf("sign ECDSA: %v", err)
		}

		var parsed ecdsaSigASN1
		if _, err := asn1.Unmarshal(derSig, &parsed); err != nil {
			t.Fatalf("unmarshal ECDSA ASN.1 signature: %v", err)
		}

		ecKey := key.(*ecdsa.PrivateKey)
		byteLen := (ecKey.Curve.Params().BitSize + 7) / 8
		rBytes := parsed.R.Bytes()
		sBytes := parsed.S.Bytes()
		for len(rBytes) < byteLen {
			rBytes = append([]byte{0}, rBytes...)
		}
		for len(sBytes) < byteLen {
			sBytes = append([]byte{0}, sBytes...)
		}
		sig = append(rBytes, sBytes...)

	case "RS256":
		digest := sha256.Sum256(signingInput)
		rsaSig, err := rsa.SignPKCS1v15(rand.Reader, key.(*rsa.PrivateKey), crypto.SHA256, digest[:])
		if err != nil {
			t.Fatalf("sign RSA: %v", err)
		}
		sig = rsaSig

	default:
		t.Fatalf("unsupported algorithm: %s", alg)
	}

	flat := JWSFlat{
		Protected: protected,
		Payload:   encodedPayload,
		Signature: base64.RawURLEncoding.EncodeToString(sig),
	}
	out, err := json.Marshal(flat)
	if err != nil {
		t.Fatalf("marshal JWSFlat: %v", err)
	}
	return out
}

// signEAB creates a valid EAB inner JWS (RFC 8555 section 7.3.4).
func signEAB(t *testing.T, kid string, hmacKey []byte, accountJWK json.RawMessage) json.RawMessage {
	t.Helper()

	header := map[string]string{
		"alg": "HS256",
		"kid": kid,
		"url": "https://acme.test/acme/test-pki/new-account",
	}
	headerJSON, err := json.Marshal(header)
	if err != nil {
		t.Fatalf("marshal EAB header: %v", err)
	}
	protected := base64.RawURLEncoding.EncodeToString(headerJSON)
	encodedPayload := base64.RawURLEncoding.EncodeToString(accountJWK)

	signingInput := []byte(protected + "." + encodedPayload)

	mac := hmac.New(sha256.New, hmacKey)
	mac.Write(signingInput)
	sig := mac.Sum(nil)

	flat := JWSFlat{
		Protected: protected,
		Payload:   encodedPayload,
		Signature: base64.RawURLEncoding.EncodeToString(sig),
	}
	out, err := json.Marshal(flat)
	if err != nil {
		t.Fatalf("marshal EAB JWSFlat: %v", err)
	}
	return json.RawMessage(out)
}

// testHandler creates a Handler with in-memory barrier for testing.
func testHandler(t *testing.T) *Handler {
	t.Helper()

	b := newMemBarrier()
	h := &Handler{
		mount:   "test-pki",
		barrier: b,
		engines: nil,
		nonces:  NewNonceStore(),
		baseURL: "https://acme.test",
		logger:  slog.New(slog.NewTextHandler(io.Discard, nil)),
	}

	// Store a default ACME config with a test issuer.
	cfg := &ACMEConfig{DefaultIssuer: "test-issuer"}
	cfgData, err := json.Marshal(cfg)
	if err != nil {
		t.Fatalf("marshal ACME config: %v", err)
	}
	if err := b.Put(context.Background(), h.barrierPrefix()+"config.json", cfgData); err != nil {
		t.Fatalf("store ACME config: %v", err)
	}
	return h
}

// createTestAccount creates an ACME account in the handler's barrier,
// bypassing the HTTP handler. Returns the account, private key, and JWK.
func createTestAccount(t *testing.T, h *Handler) (*Account, *ecdsa.PrivateKey, json.RawMessage) {
	t.Helper()
	ctx := context.Background()

	key, jwk := generateES256Key(t)

	// Create an EAB credential.
	eab, err := h.CreateEAB(ctx, "test-user")
	if err != nil {
		t.Fatalf("create EAB: %v", err)
	}
	// Mark the EAB as used since we're storing the account directly.
	if err := h.MarkEABUsed(ctx, eab.KID); err != nil {
		t.Fatalf("mark EAB used: %v", err)
	}

	// Compute the account ID from the JWK thumbprint.
	accountID := thumbprintKey(jwk)

	acc := &Account{
		ID:            accountID,
		Status:        StatusValid,
		Contact:       []string{"mailto:test@example.com"},
		JWK:           jwk,
		CreatedAt:     time.Now(),
		MCIASUsername: "test-user",
	}
	data, err := json.Marshal(acc)
	if err != nil {
		t.Fatalf("marshal account: %v", err)
	}
	path := h.barrierPrefix() + "accounts/" + accountID + ".json"
	if err := h.barrier.Put(ctx, path, data); err != nil {
		t.Fatalf("store account: %v", err)
	}
	return acc, key, jwk
}

// getNonce issues a nonce from the handler's nonce store and returns it.
func getNonce(t *testing.T, h *Handler) string {
	t.Helper()
	nonce, err := h.nonces.Issue()
	if err != nil {
		t.Fatalf("issue nonce: %v", err)
	}
	return nonce
}