"""Convert auto-labelled hallucination rows into baseline evaluator input."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys

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

from baselines import AutoLabelledDatasetAdapter, AutoLabelledRow


def parse_args():
    parser = argparse.ArgumentParser(description="Prepare auto_labelled.jsonl for token-baseline evaluation.")
    parser.add_argument("--input", required=True, help="Path to auto_labelled.jsonl.")
    parser.add_argument("--output", required=True, help="Output JSONL path for scripts/eval_baselines_dgx.py.")
    parser.add_argument("--limit", type=int, default=None, help="Optional maximum number of converted rows.")
    parser.add_argument("--source", action="append", help="Optional source filter; may be provided more than once.")
    parser.add_argument(
        "--unknown-policy",
        choices=["skip", "grounded", "error"],
        default="skip",
        help="How to handle rows containing U labels. Default skips them for clean binary evaluation.",
    )
    return parser.parse_args()


def increment_count(counts, key):
    for idx, item in enumerate(counts):
        existing_key, count = item
        if existing_key == key:
            counts[idx] = (existing_key, count + 1)
            return
    counts.append((key, 1))


def counts_to_json_dict(counts):
    return {str(key): count for key, count in counts}


def main():
    args = parse_args()
    adapter = AutoLabelledDatasetAdapter(
        unknown_policy=args.unknown_policy,
        source_filter=args.source,
    )
    input_path = Path(args.input)
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    kept = 0
    skipped = []
    source_counts = []
    span_rows = 0
    span_count = 0

    with input_path.open("r", encoding="utf-8") as f_in, output_path.open("w", encoding="utf-8") as f_out:
        for line_no, line in enumerate(f_in, start=1):
            if not line.strip():
                continue
            row = AutoLabelledRow.from_json_dict(json.loads(line))
            skip_reason = adapter.skip_reason(row)
            if skip_reason:
                increment_count(skipped, skip_reason)
                continue

            converted = adapter.convert_row(row)
            spans = converted.hallucinated_spans
            if spans:
                span_rows += 1
                span_count += len(spans)
            f_out.write(json.dumps(converted.to_json_dict()) + "\n")
            kept += 1
            increment_count(source_counts, converted.metadata.source)

            if args.limit is not None and kept >= args.limit:
                break

    print(f"wrote {kept} rows: {output_path}")
    print(f"rows with hallucinated spans: {span_rows}")
    print(f"hallucinated spans: {span_count}")
    print(f"sources: {counts_to_json_dict(source_counts)}")
    print(f"skipped: {counts_to_json_dict(skipped)}")


if __name__ == "__main__":
    main()