#!/usr/bin/env python3
"""Build deterministic, fully filled prompt examples for all 41 classes.

The examples are selected from the transformed training split. Canonical labels
are shown only as documentation metadata; the fenced system/user/assistant
blocks are copied byte-for-byte from each selected row's actual conversation.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Any, Dict, Iterable, List, Mapping


ROOT = Path(__file__).resolve().parents[1]
EXPECTED_LABELS = 41
EXPECTED_TRAIN_ROWS = 15686


def resolve(value: str) -> Path:
    path = Path(value)
    return path if path.is_absolute() else ROOT / path


def load_jsonl(path: Path) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, 1):
            if line.strip():
                try:
                    rows.append(json.loads(line))
                except json.JSONDecodeError as exc:
                    raise RuntimeError(f"Invalid JSON at {path}:{line_number}") from exc
    return rows


def dump_jsonl(rows: Iterable[Mapping[str, Any]], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(dict(row), ensure_ascii=False) + "\n")


def table_text(value: Any) -> str:
    return str(value).replace("|", "\\|").replace("\n", "<br>")


def validate_row(row: Mapping[str, Any]) -> None:
    messages = row.get("messages")
    prompt_messages = row.get("prompt_messages")
    if not isinstance(messages, list) or len(messages) != 3:
        raise RuntimeError(f"Row {row.get('id')} does not contain three labeled messages")
    if [message.get("role") for message in messages] != ["system", "user", "assistant"]:
        raise RuntimeError(f"Row {row.get('id')} has unexpected message roles")
    if messages[:2] != prompt_messages:
        raise RuntimeError(f"Row {row.get('id')} prompt/messages mismatch")
    codes = row.get("option_codes")
    options = row.get("allowed_options_ar")
    labels = row.get("allowed_relation_full_labels")
    if not isinstance(codes, list) or not isinstance(options, list) or not isinstance(labels, list):
        raise RuntimeError(f"Row {row.get('id')} has malformed option arrays")
    if not (len(codes) == len(options) == len(labels)):
        raise RuntimeError(f"Row {row.get('id')} option arrays are not aligned")
    if not options or options[-1] != "لا توجد علاقة" or labels[-1] != "no_relation":
        raise RuntimeError(f"Row {row.get('id')} does not keep no_relation last")
    gold_index = row.get("gold_option_index")
    if not isinstance(gold_index, int) or not 0 <= gold_index < len(codes):
        raise RuntimeError(f"Row {row.get('id')} has an invalid gold index")
    if row.get("gold_answer_code") != codes[gold_index]:
        raise RuntimeError(f"Row {row.get('id')} gold code/index mismatch")
    if row.get("gold_answer_ar") != options[gold_index]:
        raise RuntimeError(f"Row {row.get('id')} gold Arabic/index mismatch")
    if row.get("gold_relation_full") != labels[gold_index]:
        raise RuntimeError(f"Row {row.get('id')} gold relation/index mismatch")
    if messages[-1] != {"role": "assistant", "content": row.get("gold_answer_code")}:
        raise RuntimeError(f"Row {row.get('id')} assistant target mismatch")


def selected_examples(rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    if len(rows) != EXPECTED_TRAIN_ROWS:
        raise RuntimeError(f"Expected {EXPECTED_TRAIN_ROWS} training rows, found {len(rows)}")
    by_label: Dict[str, List[Dict[str, Any]]] = {}
    for row in rows:
        validate_row(row)
        by_label.setdefault(str(row["gold_relation_full"]), []).append(row)
    if len(by_label) != EXPECTED_LABELS:
        raise RuntimeError(
            f"Expected {EXPECTED_LABELS} labeled classes in training, found {len(by_label)}"
        )

    labels = sorted(by_label, key=lambda label: (label == "no_relation", label))
    chosen: List[Dict[str, Any]] = []
    for label in labels:
        candidates = by_label[label]
        # Short real prompts make the appendix readable. The row ID is the
        # deterministic tie-breaker, so identical input bytes regenerate the
        # same appendix.
        candidates.sort(
            key=lambda row: (
                len(str(row["prompt_messages"][1]["content"])),
                len(row["allowed_options_ar"]),
                str(row["id"]),
            )
        )
        chosen.append(candidates[0])
    return chosen


def compact_record(row: Mapping[str, Any]) -> Dict[str, Any]:
    keep = (
        "id",
        "sentence_id",
        "triple_id",
        "sentence",
        "subject",
        "object",
        "subject_start",
        "subject_end",
        "object_start",
        "object_end",
        "subject_type",
        "object_type",
        "first_type_ar",
        "second_type_ar",
        "marked_sentence",
        "allowed_relation_full_labels",
        "allowed_relation_ontology_ids",
        "allowed_options_ar",
        "option_codes",
        "gold_relation_full",
        "gold_relation_ontology_id",
        "gold_answer_ar",
        "gold_option_index",
        "gold_answer_code",
        "prompt_messages",
        "messages",
        "prompt_version",
    )
    return {key: row.get(key) for key in keep}


def markdown(examples: List[Dict[str, Any]], dataset_revision: str) -> str:
    lines = [
        "# Fully filled examples for all 41 output classes",
        "",
        "This appendix contains one real transformed **training** row for every",
        "one of the 40 positive relations and for `no_relation`. It is generated",
        "deterministically by `tools/build_examples_appendix.py`; no prompt or",
        "answer below was invented for documentation.",
        "",
        f"Dataset revision: `{dataset_revision}`.",
        "",
        "> The metadata and mapping tables are explanations for humans. The model",
        "> receives only the exact fenced **system** and **user** messages and is",
        "> trained to emit only the fenced **assistant** letter. Canonical English",
        "> labels never occur inside model conversation content.",
        "",
        "## Coverage index",
        "",
        "| # | Metadata class | Row ID | Entity types | Gold Arabic answer | Output |",
        "|---:|---|---|---|---|:---:|",
    ]
    for number, row in enumerate(examples, 1):
        type_pair = f"{row['subject_type']} → {row['object_type']}"
        lines.append(
            "| "
            + " | ".join(
                [
                    str(number),
                    f"`{table_text(row['gold_relation_full'])}`",
                    f"`{table_text(row['id'])}`",
                    f"`{table_text(type_pair)}`",
                    table_text(row["gold_answer_ar"]),
                    table_text(row["gold_answer_code"]),
                ]
            )
            + " |"
        )

    for number, row in enumerate(examples, 1):
        gold_index = int(row["gold_option_index"])
        system = str(row["messages"][0]["content"])
        user = str(row["messages"][1]["content"])
        assistant = str(row["messages"][2]["content"])
        lines.extend(
            [
                "",
                f"## {number}. `{row['gold_relation_full']}`",
                "",
                "| Property | Value |",
                "|---|---|",
                f"| Real transformed row | `{table_text(row['id'])}` |",
                f"| Directional canonical types | `{table_text(row['subject_type'])} → {table_text(row['object_type'])}` |",
                f"| Arabic types shown | {table_text(row['first_type_ar'])} → {table_text(row['second_type_ar'])} |",
                f"| Gold Arabic relation | {table_text(row['gold_answer_ar'])} |",
                f"| Zero-based gold index | `{gold_index}` |",
                f"| Exact assistant token | `{table_text(row['gold_answer_code'])}` |",
                "",
                "### Exact system message sent to Yehia",
                "",
                "```text",
                system,
                "```",
                "",
                "### Exact user message sent to Yehia",
                "",
                "```text",
                user,
                "```",
                "",
                "### Exact expected assistant output",
                "",
                "```text",
                assistant,
                "```",
                "",
                "### Row-local decoding table",
                "",
                "| Code | Arabic option displayed in the prompt | Metadata class | Gold? |",
                "|:---:|---|---|:---:|",
            ]
        )
        for index, (code, option, label) in enumerate(
            zip(
                row["option_codes"],
                row["allowed_options_ar"],
                row["allowed_relation_full_labels"],
            )
        ):
            lines.append(
                f"| {table_text(code)} | {table_text(option)} | `{table_text(label)}` | "
                f"{'yes' if index == gold_index else ''} |"
            )
        lines.extend(
            [
                "",
                f"For this row, output `{assistant}` maps through index `{gold_index}` to",
                f"`{row['gold_relation_full']}`. That letter has no permanent class meaning",
                "outside this row's displayed list.",
            ]
        )

    lines.extend(
        [
            "",
            "## Regeneration check",
            "",
            "From the repository root, after downloading the pinned transformed dataset:",
            "",
            "```bash",
            "python tools/build_examples_appendix.py",
            "```",
            "",
            "The machine-readable selected rows are in",
            "`examples/all_41_class_examples.jsonl`.",
            "",
        ]
    )
    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataset-dir",
        default=os.environ.get("CHOICE_DATASET_DIR", "data/Yehia-RE-SFT"),
    )
    parser.add_argument("--markdown", default="examples/ALL_41_CASES.md")
    parser.add_argument("--jsonl", default="examples/all_41_class_examples.jsonl")
    parser.add_argument(
        "--dataset-revision",
        default=os.environ.get(
            "SFT_DATASET_REVISION", "a060e47f56025778b97344d4d3de60a8fd53be7c"
        ),
    )
    args = parser.parse_args()

    train_path = resolve(args.dataset_dir) / "train.jsonl"
    examples = selected_examples(load_jsonl(train_path))
    markdown_path = resolve(args.markdown)
    jsonl_path = resolve(args.jsonl)
    markdown_path.parent.mkdir(parents=True, exist_ok=True)
    markdown_path.write_text(markdown(examples, args.dataset_revision), encoding="utf-8")
    dump_jsonl((compact_record(row) for row in examples), jsonl_path)
    print(
        json.dumps(
            {
                "status": "passed",
                "source": str(train_path),
                "classes": len(examples),
                "markdown": str(markdown_path),
                "jsonl": str(jsonl_path),
                "no_relation_is_final_documented_case": (
                    examples[-1]["gold_relation_full"] == "no_relation"
                ),
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()