"""Output provenance: bind an output to the certificates of its composition.

A produced output is bound, by digest and by Ed25519 signature, to the
certificates of the sources that composed it, on two axes:

* weight-source: the capability (for example a fine-tuned adapter) whose
  weights produced the output;
* context-source: the licensed context or retrieved material that informed it.

The binding is signed by the producing capability's private key, so the
record cannot be forged or edited, and each referenced certificate is pinned
by content digest, so a certificate cannot be silently swapped for another
with the same serial.

Verification checks the output digest, the signature, and the full chain of
every referenced certificate against the trust store AND the revocation
ledger. The consequence, demonstrated in the examples and tests: an output
verifies VALID while its composition is intact, and flips to INVALID the
moment any referenced certificate is revoked or crypto-erased. Provenance is
not a static stamp; it is a live claim against the current state of the
authority.

Scope note (honest bound): this is OUTPUT-LEVEL binding. Per-unit-of-text
attribution of which source produced which span is a separate measured
research result, referenced from the README by result id, and is not
reimplemented here.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Dict, Iterable, Mapping, Optional, Sequence, Tuple

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

from .certificates import Certificate, verify_chain
from .revocation import RevocationList

AXIS_WEIGHT_SOURCE = "weight-source"
AXIS_CONTEXT_SOURCE = "context-source"


@dataclass(frozen=True)
class SourceReference:
    """A composition source, pinned to an exact certificate by digest."""

    axis: str
    serial: str
    cert_sha256: str


@dataclass(frozen=True)
class ProvenanceRecord:
    output_sha256: str
    sources: Tuple[SourceReference, ...]
    signer_serial: str
    issued_at: str
    signature_hex: str = ""

    def tbs_dict(self) -> Dict:
        return {
            "output_sha256": self.output_sha256,
            "sources": [
                {"axis": s.axis, "serial": s.serial, "cert_sha256": s.cert_sha256}
                for s in self.sources
            ],
            "signer_serial": self.signer_serial,
            "issued_at": self.issued_at,
        }

    def tbs_bytes(self) -> bytes:
        return json.dumps(self.tbs_dict(), sort_keys=True, separators=(",", ":")).encode(
            "utf-8"
        )


@dataclass(frozen=True)
class ProvenanceResult:
    valid: bool
    reason: str
    checks: Tuple[str, ...] = ()

    def __bool__(self) -> bool:  # pragma: no cover - convenience only
        return self.valid


def bind_output(
    output: bytes,
    sources: Sequence[Tuple[str, Certificate]],
    signer_certificate: Certificate,
    signer_private_key: Ed25519PrivateKey,
    at: Optional[datetime] = None,
) -> ProvenanceRecord:
    """Create a signed provenance record binding ``output`` to its sources.

    ``sources`` is a sequence of ``(axis, certificate)`` pairs, one per
    composition source. The signer is the producing capability (normally the
    weight-source), holding its own leaf private key: the authority does not
    sign provenance, the capability does, and any verifier checks it against
    the capability's certificate.
    """
    moment = at or datetime.now(timezone.utc)
    unsigned = ProvenanceRecord(
        output_sha256=hashlib.sha256(output).hexdigest(),
        sources=tuple(
            SourceReference(axis=axis, serial=cert.serial, cert_sha256=cert.sha256())
            for axis, cert in sources
        ),
        signer_serial=signer_certificate.serial,
        issued_at=moment.astimezone(timezone.utc).isoformat(timespec="seconds"),
    )
    signature = signer_private_key.sign(unsigned.tbs_bytes()).hex()
    return ProvenanceRecord(
        output_sha256=unsigned.output_sha256,
        sources=unsigned.sources,
        signer_serial=unsigned.signer_serial,
        issued_at=unsigned.issued_at,
        signature_hex=signature,
    )


def verify_output_provenance(
    output: bytes,
    record: ProvenanceRecord,
    registry: Mapping[str, Certificate],
    intermediates: Iterable[Certificate],
    trusted_roots: Iterable[Certificate],
    revocation_list: Optional[RevocationList] = None,
    at: Optional[datetime] = None,
) -> ProvenanceResult:
    """Verify an output against its provenance record and the live authority state.

    Returns VALID only if every check passes:

    1. the output's digest matches the bound digest;
    2. the signer's certificate is known and its signature over the record
       verifies;
    3. the signer's chain verifies to a trusted root and is not revoked;
    4. every referenced source certificate is known, byte-identical to the
       one bound at composition time, chain-valid, and not revoked or erased.
    """
    checks = []
    intermediates = tuple(intermediates)
    trusted_roots = tuple(trusted_roots)

    def fail(reason: str) -> ProvenanceResult:
        return ProvenanceResult(False, reason, tuple(checks))

    digest = hashlib.sha256(output).hexdigest()
    if digest != record.output_sha256:
        return fail(
            "output does not match the bound digest (tampered or substituted output)"
        )
    checks.append(f"output digest matches ({digest[:16]}...)")

    signer = registry.get(record.signer_serial)
    if signer is None:
        return fail(f"signer certificate not found: {record.signer_serial}")
    try:
        signer.public_key().verify(
            bytes.fromhex(record.signature_hex), record.tbs_bytes()
        )
    except (InvalidSignature, ValueError):
        return fail(
            f"provenance signature does not verify against signer {signer.serial} "
            "(forged or edited record)"
        )
    checks.append(f"record signature by {signer.serial} verifies")

    signer_chain = verify_chain(
        signer, intermediates, trusted_roots, at=at, revocation_list=revocation_list
    )
    if not signer_chain.valid:
        return fail(f"signer chain invalid: {signer_chain.reason}")
    checks.append("signer chain verifies to trusted root")

    for source in record.sources:
        certificate = registry.get(source.serial)
        if certificate is None:
            return fail(f"{source.axis} certificate not found: {source.serial}")
        if certificate.sha256() != source.cert_sha256:
            return fail(
                f"{source.axis} certificate {source.serial} does not match the "
                "content bound at composition time (substituted certificate)"
            )
        chain = verify_chain(
            certificate,
            intermediates,
            trusted_roots,
            at=at,
            revocation_list=revocation_list,
        )
        if not chain.valid:
            return fail(
                f"{source.axis} certificate {source.serial} failed verification: "
                f"{chain.reason}"
            )
        checks.append(f"{source.axis} {source.serial}: intact, chain valid, not revoked")

    return ProvenanceResult(True, "output provenance verifies", tuple(checks))