"""Hash-chain integrity for operational custody (AUDIT-WP-0007).

See ``docs/integrity.md`` for the proof bound. This module is the shared
arithmetic; backends persist and walk rows.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable, Mapping

SCHEMA = "audit-core.chain-head.v1"
GENESIS = "0" * 64
# Documented advisory-lock key so concurrent accepts cannot fork the head.
CHAIN_LOCK_KEY = 0xA0D17007


def chain_link(previous: str, payload_hash: str, event_id: str) -> str:
    """Next ``chain_hash`` = SHA-256(previous || payload_hash || event_id)."""
    material = f"{previous}|{payload_hash}|{event_id}".encode("utf-8")
    return hashlib.sha256(material).hexdigest()


def utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


@dataclass(frozen=True)
class ChainRow:
    event_id: str
    payload_hash: str
    chain_hash: str
    chain_prev: str
    accepted_at: str | None = None


@dataclass(frozen=True)
class ChainReport:
    intact: bool
    events: int
    head: str
    head_event_id: str | None
    head_accepted_at: str | None
    first_break: str | None
    attestation_match: bool | None = None

    def as_dict(self) -> dict[str, Any]:
        return {
            "intact": self.intact,
            "events": self.events,
            "head": self.head,
            "head_event_id": self.head_event_id,
            "head_accepted_at": self.head_accepted_at,
            "first_break": self.first_break,
            "attestation_match": self.attestation_match,
        }


def verify_rows(
    rows: Iterable[ChainRow],
    *,
    attestation: Mapping[str, Any] | None = None,
) -> ChainReport:
    """Walk accept order and recompute every link.

    ``first_break`` is an event id, or ``attestation_mismatch`` when a cited
    head is not present in the live chain.
    """
    ordered = list(rows)
    prev = GENESIS
    first_break: str | None = None
    for row in ordered:
        expected = chain_link(prev, row.payload_hash, row.event_id)
        if row.chain_prev != prev or row.chain_hash != expected:
            first_break = row.event_id
            break
        prev = row.chain_hash

    intact = first_break is None
    last = ordered[-1] if ordered else None
    head = last.chain_hash if last else GENESIS
    report = ChainReport(
        intact=intact,
        events=len(ordered),
        head=head,
        head_event_id=last.event_id if last else None,
        head_accepted_at=last.accepted_at if last else None,
        first_break=first_break,
    )
    if attestation is None:
        return report
    return _apply_attestation(report, ordered, attestation)


def _apply_attestation(
    report: ChainReport,
    rows: list[ChainRow],
    attestation: Mapping[str, Any],
) -> ChainReport:
    cited = str(attestation.get("chain_hash") or "")
    live_hashes = {row.chain_hash for row in rows}
    empty_ok = not rows and cited in {"", GENESIS}
    matched = cited in live_hashes or empty_ok
    first_break = report.first_break
    intact = report.intact
    if not matched:
        intact = False
        if first_break is None:
            first_break = "attestation_mismatch"
    return ChainReport(
        intact=intact,
        events=report.events,
        head=report.head,
        head_event_id=report.head_event_id,
        head_accepted_at=report.head_accepted_at,
        first_break=first_break,
        attestation_match=matched,
    )


def attestation_from_report(report: ChainReport, *, observed_at: str | None = None) -> dict[str, Any]:
    return {
        "schema": SCHEMA,
        "genesis": GENESIS,
        "chain_hash": report.head,
        "event_id": report.head_event_id,
        "accepted_at": report.head_accepted_at,
        "event_count": report.events,
        "observed_at": observed_at or utc_now(),
    }


def load_attestation(path: str | Path) -> dict[str, Any]:
    payload = json.loads(Path(path).read_text())
    if not isinstance(payload, dict) or not payload.get("chain_hash"):
        raise ValueError("attestation must be an object with chain_hash")
    return payload


def write_attestation(path: str | Path, report: ChainReport) -> dict[str, Any]:
    body = attestation_from_report(report)
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(json.dumps(body, indent=2, sort_keys=True) + "\n")
    return body