#!/usr/bin/env python3
"""Immutable, auditable corpus manifest for Timmy training contributions.

Every ingested sample is described by one append-only manifest record: who consented, under
which consent version, where the bytes came from, the hashes of the raw input and the derived
image, how the derived image was produced, who reviewed the labels, which split it belongs to,
and where it sits in its lifecycle. Nothing in this module rewrites history: records are only
ever appended, and a withdrawal is a tombstone, not a deletion of the trail.
"""
from __future__ import annotations

import hashlib
import json
import re
from pathlib import Path

SCHEMA_VERSION = 2
PIPELINE = "timmy-ingest/2"

SAMPLE = "sample"
TOMBSTONE = "tombstone"
RECORD_TYPES = frozenset({SAMPLE, TOMBSTONE})

ACTIVE = "active"
REVOKED = "revoked"
STATES = frozenset({ACTIVE, REVOKED})

SPLITS = ("train", "validation", "test")
REVIEWERS = ("user-confirmed", "clinician-reviewed")
COLORS = ("brown", "green", "yellow", "pale", "red", "black")
QUALITIES = ("good", "fair", "poor")

_SHA256 = re.compile(r"^[0-9a-f]{64}$")
_CONTRIBUTION_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{7,127}$")


class CorpusError(Exception):
    """A record cannot be built or written as asked."""


class IllegalTransition(CorpusError):
    """A lifecycle change that must never happen is being attempted."""


def sha256_hex(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _require_sha256(value: object, field: str) -> str:
    if not isinstance(value, str) or not _SHA256.match(value):
        raise CorpusError(f"{field} must be a lowercase hex sha256 digest")
    return value


def _require_iso(value: object, field: str) -> str:
    if not isinstance(value, str) or not re.match(r"^\d{4}-\d{2}-\d{2}T", value):
        raise CorpusError(f"{field} must be an ISO-8601 timestamp")
    return value


def build_record(
    *,
    contribution_id: str,
    consent_version: str,
    subject_key: str,
    split: str,
    labels: dict,
    review_status: str,
    derived_sha256: str,
    raw_sha256: str,
    ingested_at: str,
    source_kind: str = "cli-ingest",
) -> dict:
    """Build one active sample record carrying its full provenance."""
    if not isinstance(contribution_id, str) or not _CONTRIBUTION_ID.match(contribution_id):
        raise CorpusError("contribution id is required and must be a stable opaque identifier")
    if not isinstance(consent_version, str) or not consent_version.strip():
        raise CorpusError("consent version is required")
    if not isinstance(subject_key, str) or not re.match(r"^[0-9a-f]{16}$", subject_key):
        raise CorpusError("subject key must be the 16-character pseudonymous digest")
    if split not in SPLITS:
        raise CorpusError(f"split must be one of {', '.join(SPLITS)}")
    if review_status not in REVIEWERS:
        raise CorpusError(f"review status must be one of {', '.join(REVIEWERS)}")
    if not isinstance(source_kind, str) or not source_kind.strip():
        raise CorpusError("source kind is required")
    if not isinstance(labels, dict) or set(labels) != {"bristolType", "color", "quality"}:
        raise CorpusError("labels must carry bristolType, color, and quality")
    if labels["bristolType"] not in range(1, 8):
        raise CorpusError("bristolType must be 1-7")
    if labels["color"] not in COLORS or labels["quality"] not in QUALITIES:
        raise CorpusError("color and quality must be from the documented vocabularies")
    derived = _require_sha256(derived_sha256, "derived sha256")
    raw = _require_sha256(raw_sha256, "raw input sha256")
    stamped = _require_iso(ingested_at, "ingestedAt")

    return {
        "schemaVersion": SCHEMA_VERSION,
        "recordType": SAMPLE,
        "contributionId": contribution_id,
        "imageId": derived,
        "relativePath": f"images/{derived}.jpg",
        "subjectKey": subject_key,
        "split": split,
        "consentVersion": consent_version.strip(),
        "consent": {"version": consent_version.strip(), "recordedAt": stamped},
        "source": {"kind": source_kind, "inputSha256": raw},
        "hashes": {"rawInput": raw, "derived": derived},
        "derivation": {
            "format": "jpeg",
            "maxDimension": 1200,
            "metadataStripped": True,
            "pipeline": PIPELINE,
        },
        "labels": {
            "bristolType": labels["bristolType"],
            "color": labels["color"],
            "quality": labels["quality"],
        },
        "review": {"status": review_status, "reviewedAt": stamped},
        "lifecycle": {"state": ACTIVE, "updatedAt": stamped},
        "ingestedAt": stamped,
    }


def assert_admissible(manifest_path: Path, record: dict) -> None:
    """Refuse a sample that duplicates a contribution or re-labels the same derived image."""
    if record.get("recordType", SAMPLE) != SAMPLE:
        return
    for prior in load_manifest(manifest_path):
        if prior.get("recordType", SAMPLE) != SAMPLE:
            continue
        if prior.get("contributionId") == record.get("contributionId"):
            raise IllegalTransition(f"contribution {record['contributionId']} was already ingested")
        if prior.get("imageId") == record.get("imageId") and prior.get("labels") != record.get("labels"):
            raise CorpusError("the same derived image cannot carry two different label sets")


def append_record(manifest_path: Path, record: dict) -> None:
    """Append one record to the manifest without touching what is already there."""
    assert_admissible(manifest_path, record)
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    with manifest_path.open("a", encoding="utf-8") as manifest:
        manifest.write(json.dumps(record, separators=(",", ":"), sort_keys=False) + "\n")


def load_manifest(manifest_path: Path) -> list[dict]:
    """Read the append-only manifest in order. A missing manifest is an empty corpus."""
    if not manifest_path.exists():
        return []
    with manifest_path.open("r", encoding="utf-8") as manifest:
        return [json.loads(line) for line in manifest.read().splitlines() if line.strip()]


def transition(record: dict, state: str, at: str, reason: str | None = None) -> dict:
    """Return a copy of the record in the next lifecycle state, or refuse the move."""
    current = record.get("lifecycle", {}).get("state")
    if state not in STATES:
        raise IllegalTransition(f"unknown lifecycle state: {state!r}")
    if current not in STATES:
        raise IllegalTransition(f"record carries no legal lifecycle state: {current!r}")
    if current == state:
        raise IllegalTransition(f"record is already {state}")
    if current == REVOKED:
        raise IllegalTransition("a revoked contribution cannot return to the active dataset")
    _require_iso(at, "lifecycle timestamp")
    if state == REVOKED and (not isinstance(reason, str) or not reason.strip()):
        raise IllegalTransition("revocation needs a recorded reason")
    updated = {**record, "lifecycle": {"state": state, "updatedAt": at}}
    if state == REVOKED:
        updated["lifecycle"]["reason"] = reason.strip()
    return updated


def active_records(records: list[dict]) -> list[dict]:
    """Samples still usable for training: active, and never revoked afterwards."""
    revoked = {
        record.get("contributionId")
        for record in records
        if record.get("recordType") == TOMBSTONE or record.get("lifecycle", {}).get("state") == REVOKED
    }
    return [
        record
        for record in records
        if record.get("recordType", SAMPLE) == SAMPLE
        and record.get("lifecycle", {}).get("state") == ACTIVE
        and record.get("contributionId") not in revoked
    ]


def _safe_dataset_path(dataset_root: Path, relative: str) -> Path:
    path = (dataset_root / relative).resolve()
    if not path.is_relative_to(dataset_root.resolve()):
        raise CorpusError(f"refusing to touch a path outside the dataset: {relative}")
    return path


def revoke(dataset_root: Path, contribution_id: str, at: str, reason: str) -> dict:
    """Revoke one contribution: append a tombstone and delete its bytes from the dataset."""
    manifest_path = dataset_root / "manifest.jsonl"
    records = load_manifest(manifest_path)
    samples = [
        record
        for record in records
        if record.get("recordType", SAMPLE) == SAMPLE and record.get("contributionId") == contribution_id
    ]
    if not isinstance(reason, str) or not reason.strip():
        raise IllegalTransition("revocation needs a recorded reason")
    if any(
        record.get("recordType") == TOMBSTONE and record.get("contributionId") == contribution_id
        for record in records
    ):
        raise IllegalTransition(f"contribution {contribution_id} is already revoked")
    if not samples:
        raise IllegalTransition(f"no ingested sample carries contribution {contribution_id}")
    sample = samples[-1]
    if sample.get("lifecycle", {}).get("state") != ACTIVE:
        raise IllegalTransition(f"contribution {contribution_id} is not active")

    deleted = []
    candidates = [(sample["relativePath"], sample["hashes"]["derived"])]
    # A raw pre-derivation copy is identified by the raw input digest, never by a user path.
    for path in sorted((dataset_root / "raw").glob(f"{sample['hashes']['rawInput']}*")):
        candidates.append((str(path.relative_to(dataset_root)), sample["hashes"]["rawInput"]))
    for relative, digest in candidates:
        target = _safe_dataset_path(dataset_root, relative)
        existed = target.is_file()
        if existed:
            target.unlink()
        deleted.append({"relativePath": relative, "sha256": digest, "deleted": existed})

    tombstone = {
        "schemaVersion": SCHEMA_VERSION,
        "recordType": TOMBSTONE,
        "contributionId": contribution_id,
        "imageId": sample["imageId"],
        "deletedArtifacts": deleted,
        "lifecycle": {"state": REVOKED, "updatedAt": _require_iso(at, "lifecycle timestamp"), "reason": reason.strip()},
    }
    append_record(manifest_path, tombstone)
    return tombstone