#!/usr/bin/env python3
"""Blind reviewer and adjudication tool for the Timmy training corpus (issue #26).

A reviewer reads one active sample, commits their own label without seeing what the model
said, and only then can the model answer be revealed. Adjudication is a second reader's
ruling, appended after the review it resolves; nothing in the review ledger is ever rewritten.
A withdrawn contribution cannot be queued, reviewed, revealed, adjudicated, or trained on.
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

# The tool sits beside corpus_manifest.py and is also imported by path from the test suite.
sys.path.insert(0, str(Path(__file__).resolve().parent))
from corpus_manifest import (  # noqa: E402
    COLORS,
    QUALITIES,
    CorpusError,
    active_records,
    load_manifest,
)

SCHEMA_VERSION = 1
REVIEW = "review"
ADJUDICATION = "adjudication"
DECISIONS = ("agreed", "corrected", "excluded")
ADJUDICATIONS = ("upheld", "corrected", "excluded")
LABEL_FIELDS = ("bristolType", "color", "quality")
_TIMESTAMP = re.compile(r"^\d{4}-\d{2}-\d{2}T")


class ReviewError(Exception):
    """A review action that must never happen is being attempted."""


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _manifest_path(dataset_root: Path) -> Path:
    return Path(dataset_root) / "manifest.jsonl"


def review_ledger_path(dataset_root: Path) -> Path:
    return Path(dataset_root) / "reviews.jsonl"


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


def _history(reviews: list[dict], contribution_id: str) -> list[dict]:
    return [record for record in reviews if record.get("contributionId") == contribution_id]


def _append_review(dataset_root: Path, record: dict) -> dict:
    path = review_ledger_path(dataset_root)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as ledger:
        ledger.write(json.dumps(record, separators=(",", ":"), sort_keys=False) + "\n")
    return record


def _active_sample(dataset_root: Path, contribution_id: str) -> dict:
    """The one active sample for this contribution, or a refusal that names why there is none."""
    records = load_manifest(_manifest_path(dataset_root))
    for sample in active_records(records):
        if sample.get("contributionId") == contribution_id:
            return sample
    if any(record.get("contributionId") == contribution_id for record in records):
        raise ReviewError(
            f"contribution {contribution_id} was withdrawn and cannot be reviewed, revealed, or trained on"
        )
    raise ReviewError(f"no active sample carries contribution {contribution_id}")


def _validate_labels(labels: object) -> dict:
    if not isinstance(labels, dict) or set(labels) != set(LABEL_FIELDS):
        raise ReviewError("a decided label must carry bristolType, color, and quality")
    if labels["bristolType"] not in range(1, 8):
        raise ReviewError("bristolType must be 1-7")
    if labels["color"] not in COLORS or labels["quality"] not in QUALITIES:
        raise ReviewError("color and quality must be from the documented vocabularies")
    return {field: labels[field] for field in LABEL_FIELDS}


def _decide(
    *,
    reviewer: str,
    decision: str,
    labels,
    exclusion_reason,
    proposed: dict,
    allowed: tuple[str, ...],
    change_subject: str = "proposed label",
) -> dict:
    """Validate one human decision and return the fields it contributes to the ledger."""
    reviewer_id = reviewer.strip() if isinstance(reviewer, str) else ""
    if not reviewer_id:
        raise ReviewError("reviewer is required")
    if decision not in allowed:
        raise ReviewError(f"decision must be one of {', '.join(allowed)}")
    if decision in ("agreed", "upheld"):
        return {"reviewer": reviewer_id, "decision": decision, "labels": dict(proposed)}
    if decision == "excluded":
        if not isinstance(exclusion_reason, str) or not exclusion_reason.strip():
            raise ReviewError("exclusion needs a recorded reason")
        return {
            "reviewer": reviewer_id,
            "decision": decision,
            "exclusionReason": exclusion_reason.strip(),
        }
    decided = _validate_labels(labels)
    if decided == proposed:
        raise ReviewError(f"a corrected decision must change the {change_subject}")
    return {"reviewer": reviewer_id, "decision": decision, "labels": decided}


def commit_review(
    dataset_root: Path,
    contribution_id: str,
    reviewer: str,
    decision: str,
    labels: dict | None = None,
    exclusion_reason: str | None = None,
    at: str | None = None,
) -> dict:
    """Append one human label decision for an active sample, before any model answer is shown."""
    dataset_root = Path(dataset_root)
    reviews = load_reviews(dataset_root)
    # A withdrawal outranks every other refusal: nothing may be reviewed out of a corpus that no
    # longer holds the sample, and the refusal has to say so.
    sample = _active_sample(dataset_root, contribution_id)
    if _history(reviews, contribution_id):
        raise ReviewError(
            f"contribution {contribution_id} is already reviewed; "
            "adjudication is the only way to change a decision"
        )
    proposed = dict(sample["labels"])
    committed = _decide(
        reviewer=reviewer,
        decision=decision,
        labels=labels,
        exclusion_reason=exclusion_reason,
        proposed=proposed,
        allowed=DECISIONS,
    )
    stamped = at or _now()
    if not _TIMESTAMP.match(stamped):
        raise ReviewError("committed timestamp must be ISO-8601")
    record = {
        "schemaVersion": SCHEMA_VERSION,
        "recordType": REVIEW,
        "contributionId": contribution_id,
        "imageId": sample["imageId"],
        "proposed": proposed,
        **committed,
        "committedAt": stamped,
    }
    return _append_review(dataset_root, record)


def load_predictions(path: Path) -> dict:
    """Read a model answer file: one JSON document, or one JSON object per line.

    A pretty-printed object is a document, not a set of lines, so it is parsed whole; only a file
    that is not valid JSON as a document is read line by line.
    """
    text = Path(path).read_text(encoding="utf-8")
    if not text.strip():
        return {"answers": {}}
    try:
        document = json.loads(text)
    except json.JSONDecodeError:
        document = None
    if isinstance(document, dict):
        if "answers" in document:
            return document
        if document and all(isinstance(value, dict) for value in document.values()):
            return {"answers": document}
        raise ReviewError("every prediction needs a contributionId or imageId")
    answers: dict[str, dict] = {}
    for line in text.splitlines():
        if not line.strip():
            continue
        entry = json.loads(line)
        key = entry.get("contributionId") or entry.get("imageId")
        if not key:
            raise ReviewError("every prediction needs a contributionId or imageId")
        answers[key] = entry
    return {"answers": answers}


def _model_answer(predictions: dict | None, contribution_id: str, image_id: str) -> dict | None:
    if not isinstance(predictions, dict):
        return None
    answers = predictions.get("answers", predictions)
    for key in (contribution_id, image_id):
        entry = answers.get(key)
        if isinstance(entry, dict):
            labels = {field: entry.get(field) for field in LABEL_FIELDS}
            if all(value is not None for value in labels.values()):
                answer = {"labels": labels}
                if "confidence" in entry:
                    answer["confidence"] = entry["confidence"]
                if "model" in predictions:
                    answer["model"] = predictions["model"]
                return answer
    return None


def reveal(dataset_root: Path, contribution_id: str, predictions: dict | None = None) -> dict:
    """Show the model's answer for a sample, but only after a human has committed their own.

    The order is the whole point: the reviewer's label is already fixed in the ledger when this
    runs, so a model answer can be read for disagreement but cannot steer the human into agreeing.
    """
    dataset_root = Path(dataset_root)
    sample = _active_sample(dataset_root, contribution_id)
    history = _history(load_reviews(dataset_root), contribution_id)
    if not history:
        raise ReviewError(
            f"the model answer for contribution {contribution_id} stays hidden until a review is committed"
        )
    reviewed = history[0]
    human_labels = dict(reviewed["labels"]) if "labels" in reviewed else None
    answer = _model_answer(predictions, contribution_id, sample["imageId"])
    disagree = []
    if answer and human_labels:
        disagree = [field for field in LABEL_FIELDS if answer["labels"][field] != human_labels[field]]
    return {
        "contributionId": contribution_id,
        "human": {"reviewer": reviewed["reviewer"], "decision": reviewed["decision"], "labels": human_labels},
        "model": answer,
        "disagreeFields": disagree,
        "disagreement": bool(disagree),
    }


def adjudicate(
    dataset_root: Path,
    contribution_id: str,
    reviewer: str,
    decision: str,
    labels: dict | None = None,
    exclusion_reason: str | None = None,
    at: str | None = None,
) -> dict:
    """Append a second reader's ruling on a review. The review it resolves stays in the ledger."""
    dataset_root = Path(dataset_root)
    reviews = load_reviews(dataset_root)
    sample = _active_sample(dataset_root, contribution_id)
    history = _history(reviews, contribution_id)
    if not history:
        raise ReviewError(f"no review to adjudicate for contribution {contribution_id}")
    if any(record.get("recordType") == ADJUDICATION for record in history):
        raise ReviewError(f"contribution {contribution_id} is already adjudicated")
    reviewed = history[0]
    if isinstance(reviewer, str) and reviewer.strip() == reviewed["reviewer"]:
        raise ReviewError("adjudication needs a second reviewer, not the reviewer who decided")
    reviewed_labels = dict(reviewed.get("labels") or {})
    ruling = _decide(
        reviewer=reviewer,
        decision=decision,
        labels=labels,
        exclusion_reason=exclusion_reason,
        proposed=reviewed_labels,
        allowed=ADJUDICATIONS,
        change_subject="reviewed label",
    )
    stamped = at or _now()
    if not _TIMESTAMP.match(stamped):
        raise ReviewError("committed timestamp must be ISO-8601")
    record = {
        "schemaVersion": SCHEMA_VERSION,
        "recordType": ADJUDICATION,
        "contributionId": contribution_id,
        "imageId": sample["imageId"],
        "resolvesRecord": reviews.index(reviewed) + 1,
        "resolves": {"reviewer": reviewed["reviewer"], "decision": reviewed["decision"]},
        **ruling,
        "adjudicatedAt": stamped,
    }
    return _append_review(dataset_root, record)


def trainable_records(dataset_root: Path) -> list[dict]:
    """The training set this corpus is allowed to produce right now.

    A sample is trainable only if it is still active, a human reviewed it, and the final decision
    was not an exclusion. Unreviewed proposals are never training data, and an adjudication
    replaces the label it resolved without deleting it.
    """
    dataset_root = Path(dataset_root)
    reviews = load_reviews(dataset_root)
    trainable = []
    for sample in active_records(load_manifest(_manifest_path(dataset_root))):
        history = _history(reviews, sample["contributionId"])
        if not history:
            continue
        final = history[-1]
        if final["decision"] == "excluded":
            continue
        entry = {
            "contributionId": sample["contributionId"],
            "imageId": sample["imageId"],
            "relativePath": sample["relativePath"],
            "split": sample["split"],
            "labels": dict(final["labels"]),
            "reviewedBy": history[0]["reviewer"],
        }
        if final.get("recordType") == ADJUDICATION:
            entry["adjudicatedBy"] = final["reviewer"]
        trainable.append(entry)
    return trainable


def _decision_from(args: argparse.Namespace, *, uphold: bool) -> tuple[str, dict | None, str | None]:
    """Translate the flag shape into one decision, refusing a half-specified review."""
    label_flags = [args.bristol, args.color, args.quality]
    if args.exclude:
        if any(flag is not None for flag in label_flags):
            raise ReviewError("an exclusion cannot carry labels")
        return "excluded", None, args.exclude
    if uphold and args.upheld:
        if any(flag is not None for flag in label_flags):
            raise ReviewError("an upheld review keeps the reviewed label; pass no labels with it")
        return "upheld", None, None
    if not uphold and args.agree:
        if any(flag is not None for flag in label_flags):
            raise ReviewError("agreeing keeps the proposed label; pass no labels with it")
        return "agreed", None, None
    if all(flag is not None for flag in label_flags):
        return "corrected", {"bristolType": args.bristol, "color": args.color, "quality": args.quality}, None
    if any(flag is not None for flag in label_flags):
        raise ReviewError("a corrected decision needs --bristol, --color, and --quality together")
    hint = "--upheld, a corrected label, or --exclude" if uphold else "--agree, a corrected label, or --exclude"
    raise ReviewError(f"a decision is required: pass {hint}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Blind review and adjudication for the Timmy training corpus.",
    )
    commands = parser.add_subparsers(dest="command", required=True)

    queue = commands.add_parser("queue", help="list the blind review queue")
    queue.add_argument("--dataset", required=True)
    queue.add_argument("--limit", type=int)

    review = commands.add_parser("review", help="commit a human label decision")
    review.add_argument("--dataset", required=True)
    review.add_argument("--contribution", required=True)
    review.add_argument("--reviewer", required=True)
    review.add_argument("--agree", action="store_true", help="the proposed label is correct")
    review.add_argument("--bristol", type=int, choices=range(1, 8))
    review.add_argument("--color", choices=COLORS)
    review.add_argument("--quality", choices=QUALITIES)
    review.add_argument("--exclude", metavar="REASON")

    reveal = commands.add_parser("reveal", help="show the model answer after the human commits")
    reveal.add_argument("--dataset", required=True)
    reveal.add_argument("--contribution", required=True)
    reveal.add_argument("--predictions", help="JSON or JSONL of model answers; omit if none exist")

    adjudicate_parser = commands.add_parser("adjudicate", help="append a second reader's ruling")
    adjudicate_parser.add_argument("--dataset", required=True)
    adjudicate_parser.add_argument("--contribution", required=True)
    adjudicate_parser.add_argument("--reviewer", required=True)
    adjudicate_parser.add_argument("--bristol", type=int, choices=range(1, 8))
    adjudicate_parser.add_argument("--color", choices=COLORS)
    adjudicate_parser.add_argument("--quality", choices=QUALITIES)
    adjudicate_parser.add_argument("--exclude", metavar="REASON")
    adjudicate_parser.add_argument("--upheld", action="store_true", help="keep the reviewed label")

    trainable = commands.add_parser("trainable", help="list what a trainer is allowed to read")
    trainable.add_argument("--dataset", required=True)

    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        if args.command == "queue":
            items = build_queue(Path(args.dataset), limit=args.limit)
            print(json.dumps({"count": len(items), "items": items}, indent=2))
            return 0
        if args.command == "review":
            decision, labels, exclusion_reason = _decision_from(args, uphold=False)
            record = commit_review(
                Path(args.dataset), args.contribution, args.reviewer,
                decision, labels, exclusion_reason,
            )
            print(json.dumps(record, indent=2))
            return 0
        if args.command == "reveal":
            predictions = load_predictions(args.predictions) if args.predictions else None
            print(json.dumps(reveal(Path(args.dataset), args.contribution, predictions), indent=2))
            return 0
        if args.command == "adjudicate":
            decision, labels, exclusion_reason = _decision_from(args, uphold=True)
            record = adjudicate(
                Path(args.dataset), args.contribution, args.reviewer,
                decision, labels, exclusion_reason,
            )
            print(json.dumps(record, indent=2))
            return 0
        if args.command == "trainable":
            items = trainable_records(Path(args.dataset))
            print(json.dumps({"count": len(items), "items": items}, indent=2))
            return 0
    except (ReviewError, CorpusError) as error:
        print(str(error), file=sys.stderr)
        return 2
    return 0


def build_queue(dataset_root: Path, limit: int | None = None) -> list[dict]:
    """Blind review queue: what a reviewer must judge, and nothing about what a model thinks.

    This function deliberately takes no predictions. Blindness is structural here: the queue
    cannot leak a model answer because it never receives one.
    """
    dataset_root = Path(dataset_root)
    reviews = load_reviews(dataset_root)
    reviewed = {record.get("contributionId") for record in reviews}
    queue: list[dict] = []
    for sample in active_records(load_manifest(_manifest_path(dataset_root))):
        contribution_id = sample["contributionId"]
        if contribution_id in reviewed:
            continue
        image = dataset_root / sample["relativePath"]
        if not image.is_file():
            raise ReviewError(f"the dataset lists an active sample whose image is missing: {sample['relativePath']}")
        queue.append({
            "schemaVersion": SCHEMA_VERSION,
            "contributionId": contribution_id,
            "imageId": sample["imageId"],
            "relativePath": sample["relativePath"],
            "split": sample["split"],
            "proposed": dict(sample["labels"]),
            "blindToModel": True,
        })
        if limit is not None and len(queue) >= limit:
            break
    return queue


if __name__ == "__main__":
    sys.exit(main())