// Package signing provides ECDSA P-256 artifact signing and verification.
//
// Every artifact uploaded through the API is automatically signed by the
// server's signing key. The resulting Bundle is self-contained: it carries
// the payload JSON, the base64-encoded ASN.1 signature, and the signer's
// public key PEM, so any verifier can reconstruct the check without needing
// access to the server's private key.
package signing

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"log"
	"time"
)

// KeyStore holds the server signing key pair.
type KeyStore struct {
	privateKey   *ecdsa.PrivateKey
	publicKeyPEM string
	keyID        string // 16-char hex fingerprint of the DER public key
}

// New creates a KeyStore from a PEM-encoded ECDSA private key.
func New(privateKeyPEM string) (*KeyStore, error) {
	block, _ := pem.Decode([]byte(privateKeyPEM))
	if block == nil {
		return nil, errors.New("signing: invalid PEM block")
	}
	key, err := x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("signing: parse private key: %w", err)
	}
	return newFromKey(key)
}

// Generate creates a fresh ephemeral ECDSA P-256 key pair.
// Logs a warning — not suitable for production; use ARTIFACT_SIGNING_KEY env var.
func Generate() (*KeyStore, error) {
	log.Println("signing: ARTIFACT_SIGNING_KEY not set — generating ephemeral key (signatures will not survive restart)")
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("signing: generate key: %w", err)
	}
	return newFromKey(key)
}

func newFromKey(key *ecdsa.PrivateKey) (*KeyStore, error) {
	pubDER, err := x509.MarshalPKIXPublicKey(&key.PublicKey)
	if err != nil {
		return nil, fmt.Errorf("signing: marshal public key: %w", err)
	}
	pubPEM := string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubDER}))
	sum := sha256.Sum256(pubDER)
	keyID := fmt.Sprintf("%x", sum[:8])
	return &KeyStore{privateKey: key, publicKeyPEM: pubPEM, keyID: keyID}, nil
}

// PrivateKeyPEM serialises the private key so callers can persist it.
func (ks *KeyStore) PrivateKeyPEM() (string, error) {
	der, err := x509.MarshalECPrivateKey(ks.privateKey)
	if err != nil {
		return "", err
	}
	return string(pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: der})), nil
}

// PublicKeyPEM returns the signer's PEM public key (embedded in every bundle).
func (ks *KeyStore) PublicKeyPEM() string { return ks.publicKeyPEM }

// KeyID returns the short fingerprint of the public key.
func (ks *KeyStore) KeyID() string { return ks.keyID }

// ─── Bundle types ─────────────────────────────────────────────────────────────

const bundleMediaType = "application/vnd.forgebucket.signature.bundle+json"

// Bundle is the self-contained signature artifact stored alongside each upload.
type Bundle struct {
	MediaType string        `json:"mediaType"`
	Payload   BundlePayload `json:"payload"`
	Signature string        `json:"signature"` // base64(ASN.1 DER ECDSA signature)
	PublicKey string        `json:"publicKey"` // PEM-encoded ECDSA public key
	KeyID     string        `json:"keyId"`
}

// BundlePayload is the data that was signed (JSON-serialised before hashing).
type BundlePayload struct {
	ArtifactID int64  `json:"artifactId"`
	Name       string `json:"name"`
	Digest     string `json:"digest"`   // "sha256:<hex>"
	SignedAt   string `json:"signedAt"` // RFC 3339
}

// ─── Sign ─────────────────────────────────────────────────────────────────────

// Sign computes SHA-256(rawContent), builds a BundlePayload, signs
// SHA-256(JSON(payload)) with the private key, and returns the Bundle.
func (ks *KeyStore) Sign(artifactID int64, name string, rawContent []byte) (*Bundle, error) {
	contentDigest := sha256.Sum256(rawContent)

	payload := BundlePayload{
		ArtifactID: artifactID,
		Name:       name,
		Digest:     fmt.Sprintf("sha256:%x", contentDigest),
		SignedAt:   time.Now().UTC().Format(time.RFC3339),
	}
	payloadJSON, err := json.Marshal(payload)
	if err != nil {
		return nil, fmt.Errorf("signing: marshal payload: %w", err)
	}

	payloadHash := sha256.Sum256(payloadJSON)
	sigDER, err := ecdsa.SignASN1(rand.Reader, ks.privateKey, payloadHash[:])
	if err != nil {
		return nil, fmt.Errorf("signing: ecdsa sign: %w", err)
	}

	return &Bundle{
		MediaType: bundleMediaType,
		Payload:   payload,
		Signature: base64.StdEncoding.EncodeToString(sigDER),
		PublicKey: ks.publicKeyPEM,
		KeyID:     ks.keyID,
	}, nil
}

// ─── Verify ───────────────────────────────────────────────────────────────────

// VerifyResult is returned by both verification functions.
type VerifyResult struct {
	Valid       bool   `json:"valid"`
	Digest      string `json:"digest"`
	SignedAt    string `json:"signedAt"`
	KeyID       string `json:"keyId"`
	KeyMatches  bool   `json:"keyMatchesServer"` // true if bundle public key == server public key
}

// Verify parses bundleJSON, verifies the embedded signature against the
// embedded public key, and returns a VerifyResult.
// The caller should also check KeyMatches to confirm it was signed by this server.
func (ks *KeyStore) Verify(bundleJSON []byte) (*VerifyResult, error) {
	var b Bundle
	if err := json.Unmarshal(bundleJSON, &b); err != nil {
		return nil, fmt.Errorf("signing: parse bundle: %w", err)
	}

	block, _ := pem.Decode([]byte(b.PublicKey))
	if block == nil {
		return nil, errors.New("signing: invalid public key PEM in bundle")
	}
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("signing: parse public key: %w", err)
	}
	pub, ok := pubInterface.(*ecdsa.PublicKey)
	if !ok {
		return nil, errors.New("signing: public key is not ECDSA")
	}

	payloadJSON, err := json.Marshal(b.Payload)
	if err != nil {
		return nil, fmt.Errorf("signing: marshal payload: %w", err)
	}
	payloadHash := sha256.Sum256(payloadJSON)

	sigDER, err := base64.StdEncoding.DecodeString(b.Signature)
	if err != nil {
		return nil, fmt.Errorf("signing: decode signature base64: %w", err)
	}

	valid := ecdsa.VerifyASN1(pub, payloadHash[:], sigDER)

	return &VerifyResult{
		Valid:      valid,
		Digest:     b.Payload.Digest,
		SignedAt:   b.Payload.SignedAt,
		KeyID:      b.KeyID,
		KeyMatches: b.PublicKey == ks.publicKeyPEM,
	}, nil
}