"""The witness — an append-only cosigner for a Merkle transparency log.

A witness holds its OWN signing key (separate from the log operator's) and
remembers the last tree head it cosigned. Before it will cosign a new head, it
requires an RFC 6962 **consistency proof** showing the new tree is an append-only
extension of the old one. An operator who rewrites, truncates, or forks the log
cannot obtain the witness's signature on the forged head — the consistency check
fails first.

This is the seam the whole neutrality story plugs into. Run on a *second party's*
infrastructure it provides genuine non-equivocation (the 1-of-N-honest-witness
model behind Certificate Transparency and Sigstore). Run self-hosted it is
**defense-in-depth + a migration scaffold**, not split-view protection — an
operator who holds both keys signs both. We say that plainly rather than
overclaim; the honest upgrade path is an independent witness network or an
external anchor (see ``anchor.py`` for RFC 3161 timestamping, a genuinely
third-party existence proof that ships today).

Pure stdlib + the existing signer/merkle primitives, so the whole thing is
unit-testable with no network and no external service.
"""
from __future__ import annotations

from .audit import canonical_json
from .merkle import verify_consistency_hex
from .signer import Signer, sign_payload, verify_payload


class WitnessRejected(Exception):
    """Raised when the witness refuses to cosign a head (not append-only)."""


class Witness:
    """An append-only cosigner. Construct with the witness's own ``Signer``."""

    def __init__(self, signer: Signer):
        self._signer = signer
        self._last_size: int = 0
        self._last_root: str | None = None

    @property
    def key_id(self) -> str:
        return self._signer.key_id()

    @property
    def public_key_b64(self) -> str:
        return self._signer.public_key_b64()

    # -- state persistence (so a CLI witness enforces append-only across runs) --
    def state(self) -> dict:
        """The witness's checkpoint — the last (size, root) it cosigned. Persist
        this so a later process refuses a log that was rewritten while it was gone."""
        return {"last_size": self._last_size, "last_root": self._last_root}

    def restore(self, state: dict) -> None:
        """Restore a checkpoint saved by :meth:`state`."""
        self._last_size = state.get("last_size", 0)
        self._last_root = state.get("last_root")

    # -- the core --------------------------------------------------------------
    def cosign(self, head: dict, *, consistency: dict | None = None) -> dict:
        """Cosign ``head`` (a ``TreeHead.as_dict()``) iff it is an append-only
        extension of the last head this witness cosigned.

        Rules:
          * shrinking (new size < last cosigned size) is always rejected;
          * growing (new size > last size) requires a valid consistency proof
            from the last size to the new size;
          * re-presenting the exact same (size, root) is idempotent;
          * a fresh witness (nothing cosigned yet) accepts any first head.
        """
        size = head["tree_size"]
        root = head["root_hash"]

        if self._last_root is not None:
            if size < self._last_size:
                raise WitnessRejected(
                    f"refusing to cosign a smaller tree ({size} < {self._last_size}): "
                    "a log may only grow (append-only)"
                )
            if size == self._last_size:
                if root != self._last_root:
                    raise WitnessRejected(
                        "refusing to cosign a different root at the same size: "
                        "the log was rewritten (not an append-only extension)"
                    )
                # identical head — idempotent re-cosign
            else:  # size > self._last_size: require proof of append-only growth
                if consistency is None:
                    raise WitnessRejected(
                        f"consistency proof required to cosign growth "
                        f"{self._last_size} -> {size}"
                    )
                if not verify_consistency_hex(self._last_root, consistency):
                    raise WitnessRejected(
                        f"invalid consistency proof {self._last_size} -> {size}: "
                        "the new tree is not an append-only extension of the last"
                    )

        payload = self._payload(size, root)
        envelope = sign_payload(self._signer, payload)
        self._last_size, self._last_root = size, root
        return {
            "tree_size": size,
            "root_hash": root,
            "witness_key_id": envelope["key_id"],
            "signature": envelope["value"],
            "public_key": envelope["pub"],
            "alg": envelope["alg"],
        }

    # -- verification (what a third party runs) --------------------------------
    @staticmethod
    def _payload(size: int, root: str) -> bytes:
        # domain-separated so a witness cosignature can't be confused with any
        # other signature over the same bytes
        return canonical_json(
            {"abom_witness": "v1", "tree_size": size, "root_hash": root}
        ).encode("utf-8")

    @classmethod
    def verify_cosignature(cls, cosig: dict, *, witness_public_key_b64: str) -> bool:
        """Verify a witness cosignature against the witness's public key.

        A relying party runs this with only the cosignature and the witness's
        published key — no trust in the operator required.
        """
        envelope = {
            "alg": cosig.get("alg", "ed25519"),
            "pub": cosig.get("public_key") or witness_public_key_b64,
            "value": cosig.get("signature"),
        }
        if envelope["pub"] != witness_public_key_b64:
            return False
        payload = cls._payload(cosig["tree_size"], cosig["root_hash"])
        return verify_payload(payload, envelope)
