#!/usr/bin/env python3
"""Source-bound AI review gate for biomedical and life-science subtitles.

The active session model supplies a compact domain profile and reviews one
bounded batch at a time.  This script validates every decision, preserves the
initial translation, applies only evidence-backed corrections, and emits a
checksum-bound reviewed translation set plus an auditable report.
"""

from __future__ import annotations

import argparse
from collections import Counter
import hashlib
import importlib.util
import json
from pathlib import Path
import re
import sys
from typing import Any, Sequence


SCRIPT_DIR = Path(__file__).resolve().parent
PIPELINE_PATH = SCRIPT_DIR / "subtitle_pipeline.py"
SPEC = importlib.util.spec_from_file_location("materialsub_subtitle_pipeline", PIPELINE_PATH)
if SPEC is None or SPEC.loader is None:  # pragma: no cover - installation failure
    raise RuntimeError("Could not load subtitle_pipeline.py")
pipeline = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(pipeline)


SCHEMA_VERSION = 1
REVIEW_CONTRACT_VERSION = 1
REVIEW_BATCH_SIZE = 40
CONTEXT_SEGMENTS = 2
PROFILE_NAME = "domain-profile.json"
PROFILE_INPUT_NAME = "domain-profile-input.json"
REPORT_NAME = "report.json"
REPORT_MARKDOWN_NAME = "report.md"
REVIEWED_DIR_NAME = "reviewed-translations"
INPUT_DIR_NAME = "input"
OUTPUT_DIR_NAME = "output"
RELEVANCE = {"primary", "secondary"}
STATUSES = {"approved", "corrected", "flagged"}
SEVERITIES = {"none", "low", "medium", "high"}
CATEGORIES = {
    "none",
    "terminology",
    "anatomy",
    "procedure",
    "experimental_animal",
    "drug_dose_route",
    "number_unit",
    "gene_protein_vector",
    "cell_molecular",
    "imaging_instrument",
    "statistics_results",
    "logic_negation_sequence",
    "source_text_suspected",
    "language_clarity",
    "other_scientific",
}
DISCLOSURE = (
    "本字幕经过 AI 辅助医学与生命科学术语、语义及实验参数一致性审校，"
    "未经相关专业人员人工审核。"
)


class ReviewError(RuntimeError):
    """A malformed profile, review decision, or review job."""


def _canonical(value: Any) -> bytes:
    return json.dumps(
        value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
    ).encode("utf-8")


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


def _sha256_json(value: Any) -> str:
    return _sha256_bytes(_canonical(value))


def _read_json(path: Path) -> Any:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except FileNotFoundError as exc:
        raise ReviewError(f"file not found: {path}") from exc
    except json.JSONDecodeError as exc:
        raise ReviewError(f"invalid JSON in {path}: {exc}") from exc


def _write_json(path: Path, value: Any) -> None:
    pipeline._atomic_write(path, pipeline._json_bytes(value))


def _manifest_title(manifest_path: Path) -> str:
    download_manifest = manifest_path.parent.parent / "download-manifest.json"
    if not download_manifest.is_file():
        return "Untitled scientific video"
    value = _read_json(download_manifest)
    source = value.get("source") if isinstance(value, dict) else None
    title = source.get("title") if isinstance(source, dict) else None
    return title.strip() if isinstance(title, str) and title.strip() else "Untitled scientific video"


def _segment_items(manifest: dict[str, Any], translations: dict[str, str]) -> list[dict[str, str]]:
    cue_by_id = {cue["id"]: cue for cue in manifest["cues"]}
    return [
        {
            "id": segment["id"],
            "source": pipeline._segment_source_text(segment, cue_by_id),
            "translation": translations[segment["id"]],
        }
        for segment in manifest["segments"]
    ]


def _profile_samples(items: list[dict[str, str]]) -> list[dict[str, str]]:
    if len(items) <= 18:
        return items
    indices = list(range(6))
    middle = len(items) // 2
    indices.extend(range(max(6, middle - 3), min(len(items) - 6, middle + 3)))
    indices.extend(range(len(items) - 6, len(items)))
    return [items[index] for index in sorted(set(indices))]


def _profile_input(
    manifest_path: Path, manifest: dict[str, Any], items: list[dict[str, str]]
) -> dict[str, Any]:
    return {
        "scientific_review_contract_version": REVIEW_CONTRACT_VERSION,
        "task": "classify_biomedical_life_science_domains",
        "title": _manifest_title(manifest_path),
        "source_language": manifest["source_language"],
        "target_language": manifest["target_language"],
        "representative_samples": _profile_samples(items),
        "output_fields": ["domains", "review_focus", "terminology"],
    }


def _validate_profile(value: Any, path: Path) -> dict[str, Any]:
    if not isinstance(value, dict) or set(value) != {"domains", "review_focus", "terminology"}:
        raise ReviewError(f"domain profile has forbidden or missing fields: {path}")
    domains = value["domains"]
    if not isinstance(domains, list) or not 1 <= len(domains) <= 8:
        raise ReviewError("domain profile must contain 1-8 domains")
    seen: set[str] = set()
    for domain in domains:
        if not isinstance(domain, dict) or set(domain) != {"id", "label", "relevance"}:
            raise ReviewError("each domain must contain only id, label, and relevance")
        if not all(isinstance(item, str) and item.strip() for item in domain.values()):
            raise ReviewError("domain fields must be non-empty strings")
        if not re.fullmatch(r"[a-z][a-z0-9-]{1,47}", domain["id"]):
            raise ReviewError(f"invalid domain id: {domain['id']!r}")
        if domain["id"] in seen:
            raise ReviewError(f"duplicate domain id: {domain['id']}")
        seen.add(domain["id"])
        if domain["relevance"] not in RELEVANCE:
            raise ReviewError(f"invalid domain relevance: {domain['relevance']}")
    focus = value["review_focus"]
    if not isinstance(focus, list) or not 1 <= len(focus) <= 16:
        raise ReviewError("review_focus must contain 1-16 strings")
    if any(not isinstance(item, str) or not item.strip() for item in focus):
        raise ReviewError("review_focus entries must be non-empty strings")
    terminology = value["terminology"]
    if not isinstance(terminology, list) or len(terminology) > 80:
        raise ReviewError("terminology must be an array of at most 80 entries")
    for term in terminology:
        if not isinstance(term, dict) or set(term) != {"source", "preferred", "category", "note"}:
            raise ReviewError("terminology entries have forbidden or missing fields")
        if not all(isinstance(item, str) for item in term.values()):
            raise ReviewError("terminology fields must be strings")
        if not term["source"].strip() or not term["preferred"].strip():
            raise ReviewError("terminology source and preferred fields cannot be empty")
    return value


_NUMBER_TOKEN = re.compile(r"\d+(?:[.,:/-]\d+)*(?:\s?%|\s?°[CF])?")
_SCIENTIFIC_TOKEN = re.compile(
    r"\b(?=[A-Za-z0-9+_.-]{2,}\b)(?=[A-Za-z0-9+_.-]*(?:[A-Z]{2}|\d))"
    r"[A-Za-z][A-Za-z0-9+_.-]*\b"
)


def _protected_tokens(text: str) -> Counter[str]:
    return Counter([*_NUMBER_TOKEN.findall(text), *_SCIENTIFIC_TOKEN.findall(text)])


def _review_records(value: Any, path: Path) -> list[dict[str, str]]:
    if not isinstance(value, dict) or set(value) != {"reviews"} or not isinstance(value["reviews"], list):
        raise ReviewError(f"{path} must contain only a reviews array")
    records: list[dict[str, str]] = []
    required = {"id", "status", "translation", "severity", "category", "reason"}
    for index, record in enumerate(value["reviews"], start=1):
        if not isinstance(record, dict) or set(record) != required:
            raise ReviewError(f"review {index} in {path} has forbidden or missing fields")
        if not all(isinstance(item, str) for item in record.values()):
            raise ReviewError(f"review {index} in {path} fields must be strings")
        if any("\n" in item or "\r" in item for item in record.values()):
            raise ReviewError(f"review {index} in {path} fields must be single-line strings")
        records.append(record)
    return records


def _validate_review_output(
    path: Path, batch: dict[str, Any]
) -> list[dict[str, str]]:
    records = _review_records(_read_json(path), path)
    expected = batch["items"]
    if [record["id"] for record in records] != [item["id"] for item in expected]:
        raise ReviewError(f"review output IDs mismatch: {path}")
    validated: list[dict[str, str]] = []
    for record, item in zip(records, expected):
        status = record["status"]
        severity = record["severity"]
        category = record["category"]
        reason = record["reason"].strip()
        translation = record["translation"].strip()
        initial = item["translation"].strip()
        if status not in STATUSES or severity not in SEVERITIES or category not in CATEGORIES:
            raise ReviewError(f"invalid review status, severity, or category for {record['id']}")
        if not translation:
            raise ReviewError(f"reviewed translation is empty for {record['id']}")
        if status == "approved":
            if translation != initial or severity != "none" or category != "none" or reason:
                raise ReviewError(f"approved review must preserve the initial translation: {record['id']}")
        elif status == "corrected":
            if translation == initial or severity == "none" or category == "none" or not reason:
                raise ReviewError(f"corrected review lacks a justified change: {record['id']}")
            if _protected_tokens(translation) != _protected_tokens(initial):
                raise ReviewError(
                    f"correction changes protected numbers or scientific names; flag it instead: {record['id']}"
                )
        else:
            if translation != initial or severity not in {"medium", "high"} or category == "none" or not reason:
                raise ReviewError(f"flagged review must conservatively preserve its initial translation: {record['id']}")
        validated.append({**record, "translation": translation, "reason": reason})
    return validated


def _review_batch_payload(
    items: list[dict[str, str]], start: int, profile: dict[str, Any]
) -> dict[str, Any]:
    selected = items[start : start + REVIEW_BATCH_SIZE]
    end = start + len(selected)
    return {
        "scientific_review_contract_version": REVIEW_CONTRACT_VERSION,
        "task": "source_bound_biomedical_life_science_review",
        "domain_profile": profile,
        "context": {
            "before": items[max(0, start - CONTEXT_SEGMENTS) : start],
            "after": items[end : end + CONTEXT_SEGMENTS],
        },
        "items": selected,
        "output_fields": ["id", "status", "translation", "severity", "category", "reason"],
        "allowed_statuses": sorted(STATUSES),
        "allowed_severities": sorted(SEVERITIES),
        "allowed_categories": sorted(CATEGORIES),
    }


def _prepare_batches(
    review_dir: Path, items: list[dict[str, str]], profile: dict[str, Any]
) -> list[tuple[Path, Path, dict[str, Any]]]:
    input_dir = review_dir / INPUT_DIR_NAME
    output_dir = review_dir / OUTPUT_DIR_NAME
    input_dir.mkdir(parents=True, exist_ok=True)
    output_dir.mkdir(parents=True, exist_ok=True)
    batches: list[tuple[Path, Path, dict[str, Any]]] = []
    for number, start in enumerate(range(0, len(items), REVIEW_BATCH_SIZE), start=1):
        payload = _review_batch_payload(items, start, profile)
        input_path = input_dir / f"batch-{number:04d}.json"
        output_path = output_dir / input_path.name
        encoded = _canonical(payload) + b"\n"
        if input_path.exists() and input_path.read_bytes() != encoded:
            raise ReviewError(f"existing scientific review input changed: {input_path}")
        if not input_path.exists():
            pipeline._atomic_write(input_path, encoded)
        batches.append((input_path, output_path, payload))
    return batches


def next_batch(
    manifest_path: Path, translations_dir: Path, review_dir: Path
) -> dict[str, Any]:
    manifest_path = manifest_path.expanduser().resolve()
    review_dir = review_dir.expanduser().resolve()
    manifest = pipeline.validate_manifest(manifest_path)
    translations = pipeline.load_translations(manifest, translations_dir)
    items = _segment_items(manifest, translations)
    review_dir.mkdir(parents=True, exist_ok=True)
    profile_input_path = review_dir / PROFILE_INPUT_NAME
    profile_payload = _profile_input(manifest_path, manifest, items)
    encoded_profile_input = _canonical(profile_payload) + b"\n"
    if profile_input_path.exists() and profile_input_path.read_bytes() != encoded_profile_input:
        raise ReviewError("existing scientific domain profile input changed")
    if not profile_input_path.exists():
        pipeline._atomic_write(profile_input_path, encoded_profile_input)
    profile_path = review_dir / PROFILE_NAME
    if not profile_path.exists():
        return {
            "done": False,
            "stage": "domain_profile_required",
            "input_path": str(profile_input_path),
            "output_path": str(profile_path),
            "profile": profile_payload,
        }
    profile = _validate_profile(_read_json(profile_path), profile_path)
    batches = _prepare_batches(review_dir, items, profile)
    pending: list[tuple[Path, Path, dict[str, Any]]] = []
    for input_path, output_path, payload in batches:
        if not output_path.exists():
            pending.append((input_path, output_path, payload))
        else:
            _validate_review_output(output_path, payload)
    if not pending:
        return {
            "done": True,
            "stage": "finalize_required",
            "remaining": 0,
            "review_dir": str(review_dir),
        }
    input_path, output_path, payload = pending[0]
    return {
        "done": False,
        "stage": "scientific_review_required",
        "remaining": len(pending),
        "input_path": str(input_path),
        "output_path": str(output_path),
        "batch": payload,
    }


def _markdown_report(report: dict[str, Any]) -> str:
    lines = [
        "# 科研专业审校报告",
        "",
        report["disclosure"],
        "",
        f"- 审校条目：{report['counts']['total']}",
        f"- 自动修订：{report['counts']['corrected']}",
        f"- 保守标记：{report['counts']['flagged']}",
        f"- 未解决高风险：{report['counts']['unresolved_high']}",
        "",
        "## 领域标签",
        "",
    ]
    lines.extend(
        f"- {domain['label']}（{domain['relevance']}）"
        for domain in report["profile"]["domains"]
    )
    lines.extend(["", "## 修改与疑点", ""])
    if not report["issues"]:
        lines.append("未发现需要修改或保守标记的专业问题。")
    for issue in report["issues"]:
        lines.extend(
            [
                f"### {issue['id']} · {issue['status']} · {issue['severity']}",
                "",
                f"- 类别：{issue['category']}",
                f"- 原译：{issue['initial_translation']}",
                f"- 审校后：{issue['translation']}",
                f"- 理由：{issue['reason']}",
                "",
            ]
        )
    return "\n".join(lines).rstrip() + "\n"


def finalize(
    manifest_path: Path, translations_dir: Path, review_dir: Path
) -> dict[str, Any]:
    manifest_path = manifest_path.expanduser().resolve()
    translations_dir = translations_dir.expanduser().resolve()
    review_dir = review_dir.expanduser().resolve()
    manifest = pipeline.validate_manifest(manifest_path)
    initial = pipeline.load_translations(manifest, translations_dir)
    items = _segment_items(manifest, initial)
    profile_path = review_dir / PROFILE_NAME
    profile = _validate_profile(_read_json(profile_path), profile_path)
    batches = _prepare_batches(review_dir, items, profile)
    all_reviews: list[dict[str, str]] = []
    for _, output_path, payload in batches:
        if not output_path.is_file():
            raise ReviewError(f"scientific review output is missing: {output_path}")
        all_reviews.extend(_validate_review_output(output_path, payload))
    if [item["id"] for item in all_reviews] != [item["id"] for item in items]:
        raise ReviewError("scientific reviews do not cover every subtitle segment exactly once")
    reviewed = {record["id"]: record["translation"] for record in all_reviews}
    suppressed_segment_ids = [
        record["id"] for record in all_reviews if record["status"] == "flagged"
    ]
    reviewed_dir = review_dir / REVIEWED_DIR_NAME
    reviewed_dir.mkdir(parents=True, exist_ok=True)
    reviewed_path = reviewed_dir / "translations.json"
    reviewed_payload = {
        "translations": [
            {"id": item["id"], "translation": reviewed[item["id"]]}
            for item in items
        ]
    }
    _write_json(reviewed_path, reviewed_payload)
    initial_by_id = {item["id"]: item["translation"] for item in items}
    issues = [
        {
            **record,
            "initial_translation": initial_by_id[record["id"]],
        }
        for record in all_reviews
        if record["status"] != "approved"
    ]
    counts = Counter(record["status"] for record in all_reviews)
    report = {
        "schema_version": SCHEMA_VERSION,
        "scientific_review_contract_version": REVIEW_CONTRACT_VERSION,
        "status": "complete",
        "review_method": "active_session_model_source_bound",
        "human_expert_reviewed": False,
        "disclosure": DISCLOSURE,
        "subtitle_manifest_sha256": _sha256_bytes(manifest_path.read_bytes()),
        "initial_translation_sha256": _sha256_json(initial),
        "reviewed_translation_sha256": _sha256_json(reviewed),
        "reviewed_translations_dir": str(reviewed_dir.resolve()),
        "suppressed_segment_ids": suppressed_segment_ids,
        "profile": profile,
        "counts": {
            "total": len(all_reviews),
            "approved": counts["approved"],
            "corrected": counts["corrected"],
            "flagged": counts["flagged"],
            "unresolved_high": sum(
                record["status"] == "flagged" and record["severity"] == "high"
                for record in all_reviews
            ),
        },
        "issues": issues,
    }
    report_path = review_dir / REPORT_NAME
    _write_json(report_path, report)
    pipeline._atomic_write(
        review_dir / REPORT_MARKDOWN_NAME, _markdown_report(report).encode("utf-8")
    )
    return {
        "report": str(report_path),
        "report_markdown": str(review_dir / REPORT_MARKDOWN_NAME),
        "reviewed_translations_dir": str(reviewed_dir),
        "counts": report["counts"],
    }


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Run source-bound AI scientific review over completed subtitle translations."
    )
    commands = parser.add_subparsers(dest="command", required=True)
    for name, help_text in (
        ("next-batch", "return the next domain-profile or scientific-review batch"),
        ("finalize", "validate all reviews and build reviewed translations and reports"),
    ):
        command = commands.add_parser(name, help=help_text)
        command.add_argument("--manifest", type=Path, required=True)
        command.add_argument("--translations-dir", type=Path, required=True)
        command.add_argument("--review-dir", type=Path, required=True)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = _parser().parse_args(argv)
    try:
        if args.command == "next-batch":
            payload = next_batch(args.manifest, args.translations_dir, args.review_dir)
        else:
            payload = {
                "done": True,
                "stage": "scientific_review_complete",
                **finalize(args.manifest, args.translations_dir, args.review_dir),
            }
        print(json.dumps({"ok": True, **payload}, ensure_ascii=False, sort_keys=True))
        return 0
    except (ReviewError, pipeline.PipelineError, OSError, UnicodeError) as exc:
        print(f"scientific review error: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())