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

import argparse
import json
from pathlib import Path


def load_jsonl(path: Path) -> list[dict]:
    rows = []
    with path.open("r", encoding="utf-8", errors="replace") as handle:
        for raw_line in handle:
            line = raw_line.strip()
            if not line:
                continue
            try:
                payload = json.loads(line)
            except json.JSONDecodeError:
                continue
            if isinstance(payload, dict):
                rows.append(payload)
    return rows


def as_int(value) -> int:
    if isinstance(value, bool):
        return int(value)
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        return int(value)
    if isinstance(value, str) and value.strip():
        try:
            return int(float(value))
        except ValueError:
            return 0
    return 0


def summarize(task_name: str, path: Path) -> dict:
    rows = load_jsonl(path)
    prompt_vals = [as_int(row.get("prompt_tokens")) for row in rows]
    completion_vals = [as_int(row.get("completion_tokens")) for row in rows]
    total_vals = [as_int(row.get("total_tokens")) for row in rows]
    call_vals = [as_int(row.get("llm_call_count")) for row in rows]
    with_usage = sum(1 for row in rows if any(as_int(row.get(key)) > 0 for key in ("prompt_tokens", "completion_tokens", "total_tokens")))
    n = len(rows)
    return {
        "task": task_name,
        "cases": n,
        "cases_with_usage": with_usage,
        "total_prompt_tokens": sum(prompt_vals),
        "total_completion_tokens": sum(completion_vals),
        "total_tokens": sum(total_vals),
        "avg_prompt_tokens_per_case": round(sum(prompt_vals) / n, 2) if n else 0.0,
        "avg_completion_tokens_per_case": round(sum(completion_vals) / n, 2) if n else 0.0,
        "avg_total_tokens_per_case": round(sum(total_vals) / n, 2) if n else 0.0,
        "avg_llm_calls_per_case": round(sum(call_vals) / n, 2) if n else 0.0,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Summarize real token usage from LAB-Bench compact result JSONL files.")
    parser.add_argument("--results-dir", type=Path, required=True)
    args = parser.parse_args()

    mapping = {
        "DbQA": [
            args.results_dir / "dbqa_results.jsonl",
            args.results_dir / "dbqa_batch" / "dbqa_results.jsonl",
        ],
        "SeqQA": [
            args.results_dir / "seqqa_results.jsonl",
            args.results_dir / "seqqa_batch" / "seqqa_results.jsonl",
        ],
    }
    summaries = []
    for task, candidates in mapping.items():
        path = next((candidate for candidate in candidates if candidate.exists()), None)
        if path is None:
            continue
        summaries.append(summarize(task, path))

    print(json.dumps(summaries, indent=2, ensure_ascii=False))
    return 0


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