from __future__ import annotations

import argparse
from pathlib import Path
from typing import Any

from . import _json as json


def _load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
        stripped = line.strip()
        if not stripped:
            continue
        try:
            payload = json.loads(stripped)
        except json.JSONDecodeError as exc:
            raise ValueError(f"{path}:{line_no}: invalid JSON: {exc.msg}") from exc
        if not isinstance(payload, dict):
            raise ValueError(f"{path}:{line_no}: expected JSON object row")
        rows.append(payload)
    return rows


def _coerce_attempt(value: Any) -> int:
    try:
        return max(0, int(value))
    except (TypeError, ValueError):
        return 0


def _terminal_by_case(
    *,
    rows: list[dict[str, Any]],
    condition: str,
) -> dict[str, dict[str, Any]]:
    terminal: dict[str, dict[str, Any]] = {}
    for row in rows:
        if str(row.get("condition") or "") != condition:
            continue
        case_id = str(row.get("case_id") or "").strip()
        if not case_id:
            continue
        attempt = _coerce_attempt(row.get("attempt"))
        previous = terminal.get(case_id)
        if previous is None:
            terminal[case_id] = row
            continue
        prev_attempt = _coerce_attempt(previous.get("attempt"))
        if attempt > prev_attempt or attempt == prev_attempt:
            # Last-write-wins for duplicate rows at the same attempt number.
            terminal[case_id] = row
    return terminal


def _select_case_ids(
    *,
    terminal_rows: dict[str, dict[str, Any]],
    mode: str,
    attempt_threshold: int,
) -> list[str]:
    selected: list[str] = []
    threshold = max(1, int(attempt_threshold))
    for case_id, row in terminal_rows.items():
        success = bool(row.get("success"))
        attempt = _coerce_attempt(row.get("attempt"))
        is_failed = not success
        is_attempt_ge = attempt >= threshold
        include = False
        if mode == "failed":
            include = is_failed
        elif mode == "attempt-ge":
            include = is_attempt_ge
        elif mode == "failed-or-attempt-ge":
            include = is_failed or is_attempt_ge
        if include:
            selected.append(case_id)
    return sorted(selected)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Select rerun case IDs from eval JSONL outputs."
    )
    parser.add_argument(
        "--results-jsonl",
        required=True,
        help="Path to eval JSONL file (attempt-level records).",
    )
    parser.add_argument(
        "--condition",
        default="traceback_only",
        help="Condition to read terminal rows from (default: traceback_only).",
    )
    parser.add_argument(
        "--mode",
        choices=["failed", "attempt-ge", "failed-or-attempt-ge"],
        default="failed",
        help="Selection mode (default: failed).",
    )
    parser.add_argument(
        "--attempt-threshold",
        type=int,
        default=2,
        help="Attempt threshold for attempt-ge modes (default: 2).",
    )
    parser.add_argument(
        "--out",
        required=True,
        help="Output path for newline-delimited case IDs.",
    )
    args = parser.parse_args(argv)

    input_path = Path(str(args.results_jsonl)).expanduser()
    if not input_path.exists() or not input_path.is_file():
        print(f"Results JSONL not found or not a file: {input_path}")
        return 2

    output_path = Path(str(args.out)).expanduser()
    try:
        rows = _load_jsonl(input_path)
    except ValueError as exc:
        print(str(exc))
        return 2

    condition = str(args.condition).strip()
    if not condition:
        print("Invalid --condition: must be a non-empty string")
        return 2

    terminal_rows = _terminal_by_case(rows=rows, condition=condition)
    if not terminal_rows:
        print(f"No rows found for condition={condition!r} in {input_path}")
        return 2

    selected_case_ids = _select_case_ids(
        terminal_rows=terminal_rows,
        mode=str(args.mode),
        attempt_threshold=int(args.attempt_threshold),
    )
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        "".join(f"{case_id}\n" for case_id in selected_case_ids),
        encoding="utf-8",
    )

    success_count = sum(1 for row in terminal_rows.values() if bool(row.get("success")))
    print(
        "Selected "
        f"{len(selected_case_ids)} case(s) from {len(terminal_rows)} terminal row(s) "
        f"for condition={condition!r}, mode={args.mode!r}, "
        f"attempt_threshold={int(args.attempt_threshold)} (terminal successes={success_count})."
    )
    print(f"Wrote case list: {output_path}")
    return 0


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