package ucp

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"
)

// ---------------------------------------------------------------------------
// SigningConfig — ECDSA P-256 key for RFC 9421 HTTP Message Signatures
// ---------------------------------------------------------------------------

// SigningConfig holds the ECDSA P-256 private key and key ID for signing UCP
// HTTP responses with RFC 9421 HTTP Message Signatures. The corresponding
// public key is published in the UCP profile (signing_keys) and served at
// /.well-known/ucp so any UCP platform can verify response authenticity.
type SigningConfig struct {
	PrivateKey *ecdsa.PrivateKey
	KeyID      string // matches the kid in the UCP profile signing_keys
}

// NewSigningConfigFromPEM parses a PEM-encoded ECDSA P-256 private key and
// returns a SigningConfig ready for use.
//
//	data — PEM block containing an EC PRIVATE KEY (PKCS#8 or SEC1)
//	keyID — the kid matched by the UCP profile signing_keys entry
func NewSigningConfigFromPEM(pemData []byte, keyID string) (*SigningConfig, error) {
	block, _ := pem.Decode(pemData)
	if block == nil {
		return nil, fmt.Errorf("ucp signing: no PEM block found")
	}

	var priv any
	var err error
	switch block.Type {
	case "EC PRIVATE KEY":
		// SEC1 format (openssl ecparam -genkey -name prime256v1)
		priv, err = x509.ParseECPrivateKey(block.Bytes)
	case "PRIVATE KEY":
		// PKCS#8 format
		priv, err = x509.ParsePKCS8PrivateKey(block.Bytes)
	default:
		return nil, fmt.Errorf("ucp signing: unsupported PEM type %q", block.Type)
	}
	if err != nil {
		return nil, fmt.Errorf("ucp signing: parse key: %w", err)
	}

	ecKey, ok := priv.(*ecdsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("ucp signing: key is not ECDSA")
	}
	if ecKey.Curve != elliptic.P256() {
		return nil, fmt.Errorf("ucp signing: key is not P-256 (got %s)", ecKey.Curve.Params().Name)
	}

	return &SigningConfig{PrivateKey: ecKey, KeyID: keyID}, nil
}

// ---------------------------------------------------------------------------
// WithSigning — HTTP middleware wrapper for RFC 9421 response signing
// ---------------------------------------------------------------------------

// WithSigning wraps h with an HTTP middleware that adds RFC 9421 HTTP Message
// Signatures to every response. The Signature-Input and Signature headers are
// computed over @status, content-type, x-ucp-timestamp, and content-digest,
// signed with the configured ECDSA P-256 key.
//
// Mount it on any existing UCP handler:
//
//	mux.Handle("/ucp/v1/carts", ucp.WithSigning(ucp.CartHandler(pool), cfg))
func WithSigning(h http.Handler, cfg *SigningConfig) http.Handler {
	if cfg == nil || cfg.PrivateKey == nil {
		return h // pass through without signing
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sw := newSignedWriter(w, cfg)
		h.ServeHTTP(sw, r)
		sw.flush()
	})
}

// ---------------------------------------------------------------------------
// signedWriter — response interceptor that adds signature headers
// ---------------------------------------------------------------------------

type signedWriter struct {
	responseWriter http.ResponseWriter
	header         http.Header
	body           bytes.Buffer
	cfg            *SigningConfig
	statusCode     int
	wroteHeader    bool
}

func newSignedWriter(w http.ResponseWriter, cfg *SigningConfig) *signedWriter {
	return &signedWriter{
		responseWriter: w,
		header:         make(http.Header),
		cfg:            cfg,
	}
}

func (w *signedWriter) Header() http.Header {
	return w.header
}

func (w *signedWriter) WriteHeader(statusCode int) {
	if w.wroteHeader {
		return
	}
	w.wroteHeader = true
	w.statusCode = statusCode
}

func (w *signedWriter) Write(p []byte) (int, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	return w.body.Write(p)
}

func (w *signedWriter) flush() {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	if w.header.Get("Content-Type") == "" && w.body.Len() > 0 {
		w.header.Set("Content-Type", http.DetectContentType(w.body.Bytes()))
	}

	created := time.Now().Unix()
	w.header.Set("x-ucp-timestamp", strconv.FormatInt(created, 10))
	w.header.Set("Content-Digest", buildContentDigest(w.body.Bytes()))

	// Add signature headers before flushing.
	w.addSignatureHeaders(created)

	dst := w.responseWriter.Header()
	for k, vals := range w.header {
		dst[k] = append([]string(nil), vals...)
	}
	w.responseWriter.WriteHeader(w.statusCode)
	_, _ = w.responseWriter.Write(w.body.Bytes())
}

// addSignatureHeaders computes and writes the Signature-Input and Signature
// headers per RFC 9421 §3.2 / §3.3.
func (w *signedWriter) addSignatureHeaders(created int64) {
	contentType := w.header.Get("Content-Type")
	sigTimestamp := strconv.FormatInt(created, 10)
	contentDigest := w.header.Get("Content-Digest")

	// Covered components (in order). RFC 9421 §3.2: derived component names
	// (@status) are bare identifiers; HTTP header field names are sf-strings
	// and MUST be quoted.
	components := []string{
		"@status",
		`"content-type"`,
		`"x-ucp-timestamp"`,
		`"content-digest"`,
	}

	// Build the signature base string (RFC 9421 §2.2).
	var base strings.Builder
	base.WriteString(fmt.Sprintf(`"@status": %d`, w.statusCode))
	base.WriteByte('\n')
	base.WriteString(fmt.Sprintf(`"content-type": %s`, contentType))
	base.WriteByte('\n')
	base.WriteString(fmt.Sprintf(`"x-ucp-timestamp": %s`, sigTimestamp))
	base.WriteByte('\n')
	base.WriteString(fmt.Sprintf(`"content-digest": %s`, contentDigest))
	base.WriteByte('\n')

	// Append the signature-params pseudo-line (RFC 9421 §2.2).
	compList := strings.Join(components, " ")
	base.WriteString(fmt.Sprintf(`"@signature-params": (%s);created=%d;keyid=%q;alg="ecdsa-p256"`,
		compList, created, w.cfg.KeyID))

	// Sign the SHA-256 digest.
	digest := sha256.Sum256([]byte(base.String()))
	r, s, err := ecdsa.Sign(rand.Reader, w.cfg.PrivateKey, digest[:])
	if err != nil {
		// Fail open — response is sent without signature headers.
		return
	}

	// ECDSA P-256 signature raw bytes: r||s concatenation, each 32 bytes.
	rBytes := r.Bytes()
	sBytes := s.Bytes()
	rawSig := make([]byte, 64)
	copy(rawSig[32-len(rBytes):32], rBytes)
	copy(rawSig[64-len(sBytes):64], sBytes)

	sigB64 := base64.RawURLEncoding.EncodeToString(rawSig)

	// Signature-Input: sig1=(components);created=N;keyid="...";alg="ecdsa-p256"
	sigInput := fmt.Sprintf(`sig1=(%s);created=%d;keyid=%q;alg="ecdsa-p256"`,
		compList, created, w.cfg.KeyID)

	// Signature: sig1=:base64url:
	sigValue := fmt.Sprintf("sig1=:%s:", sigB64)

	w.header.Set("Signature-Input", sigInput)
	w.header.Set("Signature", sigValue)
}

func buildContentDigest(body []byte) string {
	sum := sha256.Sum256(body)
	return "sha-256=:" + base64.StdEncoding.EncodeToString(sum[:]) + ":"
}