"""CLI for atomically evaluating canonical sample JSONL artifacts."""

from __future__ import annotations

import argparse
from dataclasses import asdict
import json
from pathlib import Path
from typing import Sequence

from dlb.io import atomic_json_write, sha256_file, validate_samples
from dlb.schema import SampleRecord

from .generative_perplexity import (
    compute_gen_ppl,
    load_offline_gpt2_large,
    resolve_gpt2_assets,
)
from .self_bleu import SelfBleuConfig, compute_self_bleu
from .unigram_entropy import mean_unigram_entropy


PRODUCTION_SAMPLE_COUNT = 1024
METRIC_NAMES = ("gen_ppl", "entropy", "self_bleu")
DATASET_PADDING_IDS = {"lm1b": frozenset({0}), "owt": frozenset()}


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--samples", type=Path, required=True)
    parser.add_argument("--metrics", default=",".join(METRIC_NAMES))
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument("--dataset", choices=tuple(DATASET_PADDING_IDS))
    parser.add_argument("--special-id", type=int, action="append", default=[])
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--context-length", type=int, default=1024)
    parser.add_argument("--device", default="cuda")
    parser.add_argument(
        "--allow-partial",
        action="store_true",
        help="allow a non-production sample count and label the artifact partial",
    )
    return parser.parse_args(argv)


def _metrics(value: str) -> tuple[str, ...]:
    names = tuple(part.strip() for part in value.split(",") if part.strip())
    if not names:
        raise ValueError("at least one metric is required")
    unknown = sorted(set(names) - set(METRIC_NAMES))
    if unknown:
        raise ValueError(f"unknown metrics: {', '.join(unknown)}")
    if len(set(names)) != len(names):
        raise ValueError("metric names must not be duplicated")
    return names


def _load_records(path: Path) -> list[SampleRecord]:
    records = []
    with path.open("r", encoding="utf-8") as source:
        for line in source:
            records.append(SampleRecord.model_validate(json.loads(line)))
    return records


def _display_path(path: Path, root: Path) -> str:
    try:
        return path.resolve().relative_to(root.resolve()).as_posix()
    except ValueError:
        return path.resolve().as_posix()


def evaluate(args: argparse.Namespace) -> dict[str, object]:
    root = args.root.resolve()
    metric_names = _metrics(args.metrics)
    expected = None if args.allow_partial else PRODUCTION_SAMPLE_COUNT
    sample_count = validate_samples(args.samples, expected=expected)
    if sample_count < 1:
        raise ValueError("sample artifact is empty")
    partial = sample_count != PRODUCTION_SAMPLE_COUNT
    if partial and not args.allow_partial:
        raise ValueError("partial input requires --allow-partial")
    if not partial and args.special_id:
        raise ValueError(
            "production evaluation forbids --special-id; dataset padding policy is fixed"
        )
    if not partial and "entropy" in metric_names and args.dataset is None:
        raise ValueError("production entropy evaluation requires --dataset")
    records = _load_records(args.samples)
    texts = [record.text for record in records]
    result_metrics: dict[str, object] = {}

    if "gen_ppl" in metric_names:
        assets = resolve_gpt2_assets(root)
        model, tokenizer = load_offline_gpt2_large(assets, device=args.device)
        ppl = compute_gen_ppl(
            texts,
            model,
            tokenizer,
            batch_size=args.batch_size,
            max_length=args.context_length,
            model_revision=assets.model_revision,
            tokenizer_revision=assets.tokenizer_revision,
            device=args.device,
        )
        result_metrics["generative_perplexity"] = {
            **asdict(ppl),
            "model_snapshot": _display_path(assets.model_path, root),
            "tokenizer_snapshot": _display_path(assets.tokenizer_path, root),
            "offline": True,
        }

    excluded = set()
    exclusion_source = "none_partial"
    special_token_policy = "partial_no_token_exclusions"
    if args.dataset is not None:
        excluded.update(DATASET_PADDING_IDS[args.dataset])
        exclusion_source = "dataset_contract"
        special_token_policy = "exclude_documented_padding_only_preserve_bos_eos"
    if partial and args.special_id:
        excluded.update(args.special_id)
        exclusion_source = "custom_partial_cli"
        special_token_policy = "custom_partial_not_baseline_comparable"
    if "entropy" in metric_names:
        entropy = mean_unigram_entropy(records, excluded)
        result_metrics["unigram_entropy"] = {
            **asdict(entropy),
            "special_token_policy": special_token_policy,
            "exclusion_source": exclusion_source,
        }

    if "self_bleu" in metric_names:
        result_metrics["self_bleu"] = asdict(
            compute_self_bleu(texts, SelfBleuConfig())
        )

    return {
        "schema_version": 1,
        "samples": _display_path(args.samples, root),
        "samples_sha256": sha256_file(args.samples),
        "sample_count": sample_count,
        "production_sample_count": PRODUCTION_SAMPLE_COUNT,
        "partial": partial,
        "dataset": args.dataset,
        "requested_metrics": list(metric_names),
        "metrics": result_metrics,
    }


def main(argv: Sequence[str] | None = None) -> int:
    args = parse_args(argv)
    document = evaluate(args)
    atomic_json_write(args.output, document)
    print(
        f"WROTE {args.output} samples={document['sample_count']} "
        f"partial={str(document['partial']).lower()}"
    )
    return 0


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