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

import argparse
import sys
from collections import Counter
from pathlib import Path

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.index import ShardIndex  # noqa: E402
from dovla_cil.data.schema import CILRecord
from dovla_cil.data.sharding import ShardReader, group_records, iter_cil_records


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Inspect a CIL dataset, metadata file, or shard.")
    parser.add_argument("path", type=Path)
    parser.add_argument("--group-id", default=None, help="Group to display. Defaults to first group.")
    parser.add_argument(
        "--max-rows",
        type=int,
        default=32,
        help="Maximum rows to print in the ranking table.",
    )
    args = parser.parse_args(argv)

    if _looks_like_dataset(args.path):
        _inspect_dataset(args.path, group_id=args.group_id, max_rows=args.max_rows)
    else:
        _inspect_single_shard(args.path, group_id=args.group_id, max_rows=args.max_rows)
    return 0


def _inspect_dataset(path: Path, *, group_id: str | None, max_rows: int) -> None:
    reader = ShardReader(path)
    index = reader.index
    _print_summary(index)
    entries = index.load_group_index()
    if not entries:
        print("group_index: missing or empty")
        return
    selected_group_id = group_id or str(entries[0]["group_id"])
    records = reader.load_group(selected_group_id)
    _print_group_index_entry(_entry_by_group(entries, selected_group_id))
    _print_ranking_table(records, max_rows=max_rows)


def _inspect_single_shard(path: Path, *, group_id: str | None, max_rows: int) -> None:
    records = list(iter_cil_records(path))
    groups = group_records(records)
    candidate_counts = Counter(record.candidate_type for record in records)
    print(f"shard: {path}")
    print(f"num_records: {len(records)}")
    print(f"num_groups: {len(groups)}")
    print(f"candidate_type_counts: {dict(sorted(candidate_counts.items()))}")
    if not groups:
        return
    selected_group_id = group_id or next(iter(groups))
    if selected_group_id not in groups:
        raise KeyError(f"Unknown group {selected_group_id!r} in shard {path}")
    _print_ranking_table(groups[selected_group_id], max_rows=max_rows)


def _print_summary(index: ShardIndex) -> None:
    metadata = index.metadata
    print(f"dataset: {metadata.get('dataset_name', index.dataset_dir.name)}")
    print(f"path: {index.dataset_dir}")
    print(f"version: {metadata.get('version', 'unknown')}")
    print(f"schema_version: {metadata.get('schema_version', 'unknown')}")
    print(f"backend: {metadata.get('backend', 'unknown')}")
    print(f"created_at: {metadata.get('created_at', 'unknown')}")
    print(f"num_groups: {index.group_count}")
    print(f"num_records: {index.record_count}")
    print(f"k: {metadata.get('k', 'unknown')}")
    print(f"task_count: {metadata.get('task_count', 'unknown')}")
    print(f"seed: {metadata.get('seed', 'unknown')}")
    print(f"shard_count: {len(index.shards)}")
    print(f"group_index: {index.group_index_path}")
    print(f"record_index: {index.record_index_path}")


def _print_group_index_entry(entry: dict[str, object]) -> None:
    print(f"sample_group: {entry.get('group_id')}")
    print(f"sample_group_task: {entry.get('task_id')}")
    print(f"sample_group_shard: {entry.get('shard_path')}")
    print(f"sample_group_records: {entry.get('num_records')}")
    print(f"sample_group_max_reward: {entry.get('max_reward')}")
    print(f"sample_group_success_count: {entry.get('success_count')}")
    print(f"sample_group_candidate_type_counts: {entry.get('candidate_type_counts')}")


def _print_ranking_table(records: list[CILRecord], *, max_rows: int) -> None:
    ordered = sorted(
        records,
        key=lambda record: (
            record.rank_within_group if record.rank_within_group is not None else 10**9,
            record.record_id,
        ),
    )
    print("ranking:")
    print("record_id\tcandidate_type\treward.progress\tsuccess\tregret\trank\tfailure.type")
    for record in ordered[:max_rows]:
        print(
            "\t".join(
                [
                    record.record_id,
                    record.candidate_type,
                    f"{record.reward.progress:.6g}",
                    str(record.reward.terminal_success),
                    "" if record.regret is None else f"{record.regret:.6g}",
                    ""
                    if record.rank_within_group is None
                    else str(record.rank_within_group),
                    record.failure.type if record.failure else "none",
                ]
            )
        )
    if len(ordered) > max_rows:
        print(f"... {len(ordered) - max_rows} more records")


def _entry_by_group(entries: list[dict[str, object]], group_id: str) -> dict[str, object]:
    for entry in entries:
        if entry.get("group_id") == group_id:
            return entry
    raise KeyError(f"Unknown group {group_id!r}")


def _looks_like_dataset(path: Path) -> bool:
    if path.is_dir():
        return (path / "metadata.json").exists() or (path / "manifest.json").exists()
    return path.name in {"metadata.json", "manifest.json"}


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