#!/usr/bin/env python
from __future__ import annotations

import argparse
import hashlib
import json
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.data.datasets import CILDataset  # noqa: E402
from dovla_cil.generation.tangent_targets import (  # noqa: E402
    DEFAULT_BASE_CANDIDATE_TYPES,
    action_matrix,
    action_shape,
    choose_base_record,
    label_from_delta_utility,
    record_utility,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Build CIL chart data accounting tables.")
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, default=Path("runs/data_accounting"))
    parser.add_argument("--epsilon", type=float, default=0.05)
    parser.add_argument("--split-fractions", default="0.70,0.15,0.15")
    parser.add_argument("--split-seed", type=int, default=0)
    parser.add_argument(
        "--base-candidate-types",
        default=",".join(DEFAULT_BASE_CANDIDATE_TYPES),
    )
    parser.add_argument(
        "--no-markdown-report",
        action="store_true",
        help="Do not write report.md; persistent prose is consolidated in README.md.",
    )
    args = parser.parse_args(argv)

    fractions = _parse_split_fractions(args.split_fractions)
    base_candidate_types = _parse_csv(args.base_candidate_types)
    dataset = CILDataset(args.dataset)
    rows: dict[str, dict[str, Any]] = {
        split: _empty_row(split) for split in ("train", "val", "test")
    }
    skip_counts: Counter[str] = Counter()

    for group in dataset.iter_groups():
        if not group:
            skip_counts["empty_group"] += 1
            continue
        split = _assign_split(group[0].group_id, fractions=fractions, seed=args.split_seed)
        row = rows[split]
        base = choose_base_record(group, base_candidate_types=base_candidate_types)
        if base is None:
            skip_counts["missing_base"] += 1
            continue
        base_action = action_matrix(base)
        if not base_action:
            skip_counts["empty_base_action"] += 1
            continue
        base_shape = action_shape(base_action)
        base_utility = record_utility(base)
        row["chart_ids"].add(group[0].group_id)
        row["state_hashes"].add(group[0].state_hash)
        row["task_ids"].add(group[0].task_id)
        seed = group[0].metadata.get("episode_seed", group[0].metadata.get("random_seed"))
        if seed is not None:
            row["seeds"].add(str(seed))
        for record in group:
            action = action_matrix(record)
            if not action or action_shape(action) != base_shape:
                skip_counts["unusable_branch"] += 1
                continue
            row["num_branches"] += 1
            candidate_type = str(record.candidate_type)
            row["candidate_type_counts"][candidate_type] += 1
            if record.record_id == base.record_id:
                row["num_base_branches"] += 1
            if candidate_type == "expert" or candidate_type.endswith("_expert"):
                row["num_expert"] += 1
            if record.record_id == base.record_id:
                label = "neutral"
            else:
                label = label_from_delta_utility(
                    record_utility(record) - base_utility,
                    epsilon=args.epsilon,
                )
            row[f"num_{label}"] += 1

    materialized = [_finalize_row(row) for row in rows.values()]
    payload = {
        "schema_version": 1,
        "dataset": str(args.dataset),
        "epsilon": args.epsilon,
        "split_fractions": fractions,
        "split_seed": args.split_seed,
        "split_hash": _split_hash(materialized),
        "skip_counts": dict(sorted(skip_counts.items())),
        "rows": materialized,
    }
    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "table.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    (out_dir / "table.tex").write_text(_latex_table(materialized) + "\n")
    _write_markdown_report(out_dir, payload, no_markdown_report=args.no_markdown_report)
    print(json.dumps({"out_dir": str(out_dir), "split_hash": payload["split_hash"]}, indent=2))
    return 0


def _write_markdown_report(
    out_dir: Path,
    payload: dict[str, Any],
    *,
    no_markdown_report: bool,
) -> None:
    report_path = out_dir / "report.md"
    if no_markdown_report:
        report_path.unlink(missing_ok=True)
        return
    report_path.write_text(_markdown_report(payload) + "\n")


def _empty_row(split: str) -> dict[str, Any]:
    return {
        "split": split,
        "task_ids": set(),
        "seeds": set(),
        "chart_ids": set(),
        "state_hashes": set(),
        "num_branches": 0,
        "num_base_branches": 0,
        "num_positive": 0,
        "num_neutral": 0,
        "num_negative": 0,
        "num_expert": 0,
        "candidate_type_counts": Counter(),
    }


def _finalize_row(row: dict[str, Any]) -> dict[str, Any]:
    split = row["split"]
    num_branches = int(row["num_branches"])
    return {
        "split": split,
        "num_tasks": len(row["task_ids"]),
        "num_seeds": len(row["seeds"]),
        "num_charts": len(row["chart_ids"]),
        "num_branches": num_branches,
        "num_base_branches": int(row["num_base_branches"]),
        "num_positive": int(row["num_positive"]),
        "num_neutral": int(row["num_neutral"]),
        "num_negative": int(row["num_negative"]),
        "num_expert": int(row["num_expert"]),
        "num_train_only_retrieval_rows": num_branches if split == "train" else 0,
        "num_eval_only_rows": num_branches if split != "train" else 0,
        "candidate_type_counts": dict(sorted(row["candidate_type_counts"].items())),
        "chart_hashes": sorted(_hash_id(value) for value in row["chart_ids"]),
        "state_hashes": sorted(_hash_id(value) for value in row["state_hashes"]),
    }


def _latex_table(rows: list[dict[str, Any]]) -> str:
    lines = [
        "% Auto-generated by scripts/build_data_accounting.py",
        "\\begin{tabular}{lrrrrrrrrrrr}",
        "\\toprule",
        "Split & Tasks & Seeds & Charts & Branches & Base & Positive & Neutral & Negative & Expert & TrainRows & EvalRows \\\\",
        "\\midrule",
    ]
    for row in rows:
        lines.append(
            f"{row['split']} & {row['num_tasks']} & {row['num_seeds']} & "
            f"{row['num_charts']} & {row['num_branches']} & {row['num_base_branches']} & "
            f"{row['num_positive']} & {row['num_neutral']} & {row['num_negative']} & "
            f"{row['num_expert']} & {row['num_train_only_retrieval_rows']} & "
            f"{row['num_eval_only_rows']} \\\\"
        )
    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)


def _markdown_report(payload: dict[str, Any]) -> str:
    lines = [
        "# CIL Chart Data Accounting",
        "",
        f"Dataset: `{payload['dataset']}`",
        f"Split seed: `{payload['split_seed']}`",
        f"Split hash: `{payload['split_hash']}`",
        "",
        "| Split | Tasks | Seeds | Charts | Branches | Base | Positive | Neutral | Negative | Expert | Train rows | Eval rows |",
        "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |",
    ]
    for row in payload["rows"]:
        lines.append(
            f"| {row['split']} | {row['num_tasks']} | {row['num_seeds']} | "
            f"{row['num_charts']} | {row['num_branches']} | {row['num_base_branches']} | "
            f"{row['num_positive']} | {row['num_neutral']} | {row['num_negative']} | "
            f"{row['num_expert']} | {row['num_train_only_retrieval_rows']} | "
            f"{row['num_eval_only_rows']} |"
        )
    if payload["skip_counts"]:
        lines.extend(["", "Skip counts:", ""])
        lines.extend(f"- `{key}`: {value}" for key, value in payload["skip_counts"].items())
    return "\n".join(lines)


def _parse_csv(value: str) -> tuple[str, ...]:
    return tuple(item.strip() for item in value.split(",") if item.strip())


def _parse_split_fractions(value: str) -> dict[str, float]:
    parts = [float(item.strip()) for item in value.split(",") if item.strip()]
    if len(parts) != 3:
        raise ValueError("--split-fractions must contain train,val,test values")
    total = sum(parts)
    if total <= 0.0 or any(part < 0.0 for part in parts):
        raise ValueError("split fractions must be non-negative with positive sum")
    train, val, test = [part / total for part in parts]
    return {"train": train, "val": val, "test": test}


def _assign_split(group_id: str, *, fractions: dict[str, float], seed: int) -> str:
    value = _stable_uniform(group_id, seed=seed)
    if value < fractions["train"]:
        return "train"
    if value < fractions["train"] + fractions["val"]:
        return "val"
    return "test"


def _stable_uniform(value: str, *, seed: int) -> float:
    digest = hashlib.sha256(f"{seed}:{value}".encode("utf-8")).digest()
    return int.from_bytes(digest[:8], "big") / float(2**64)


def _hash_id(value: str) -> str:
    return hashlib.sha256(str(value).encode()).hexdigest()


def _split_hash(rows: list[dict[str, Any]]) -> str:
    digest_rows = [
        {
            "split": row["split"],
            "chart_hashes": row["chart_hashes"],
            "state_hashes": row["state_hashes"],
        }
        for row in rows
    ]
    return hashlib.sha256(json.dumps(digest_rows, sort_keys=True).encode()).hexdigest()


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