#!/usr/bin/env pytho

"""
GAIA Benchmark Runner — CLI
Usage:
    python run_benchmark.py --level 1 --split validation --max 20
    python run_benchmark.py --level all --split validation --max 0   # run all
"""

from agent import gaia_benchmark_iter
import argparse, os, json, re, string, warnings
import numpy as np


# ── GAIA official scorer ────────────────────────────────────────────────────

def normalize_number_str(number_str: str) -> float:
    for char in ["$", "%", ","]:
        number_str = number_str.replace(char, "")
    try:
        return float(number_str)
    except ValueError:
        return float("inf")


def split_string(s: str, char_list: list[str] = [",", ";"]) -> list[str]:
    pattern = f"[{''.join(re.escape(c) for c in char_list)}]"
    return [x.strip() for x in re.split(pattern, s)]


def normalize_str(input_str: str, remove_articles: bool = True) -> str:
    input_str = str(input_str).strip()
    if remove_articles:
        input_str = re.sub(r"\b(a|an|the)\b", "", input_str, flags=re.IGNORECASE)
    return " ".join(input_str.lower().translate(str.maketrans("", "", string.punctuation)).split())


def question_scorer(model_answer: str, ground_truth: str) -> bool:
    def is_float(s):
        try:
            float(s)
            return True
        except ValueError:
            return False

    if is_float(ground_truth):
        try:
            return abs(normalize_number_str(model_answer) - float(ground_truth)) <= 1e-6
        except Exception:
            return False

    if any(c in ground_truth for c in [",", ";"]):
        gt_parts = split_string(ground_truth)
        ma_parts = split_string(model_answer)
        if len(gt_parts) != len(ma_parts):
            return False
        return all(
            question_scorer(ma.strip(), gt.strip())
            for ma, gt in zip(ma_parts, gt_parts)
        )

    return normalize_str(model_answer) == normalize_str(ground_truth)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--level", default="1", choices=["1", "2", "3", "all"])
    parser.add_argument("--split", default="validation", choices=["validation", "test"])
    parser.add_argument("--max", type=int, default=20)
    args = parser.parse_args()

    hf_token = os.environ.get("HF_TOKEN", "")
    if not hf_token:
        print("ERROR: HF_TOKEN not set.  export HF_TOKEN=hf_...")
        return

    serper = os.environ.get("SERPER_API_KEY") or os.environ.get("SERPER_API", "")
    if not serper:
        print("⚠  WARNING: SERPER_API not set — search_tool will fail. Set it first:")
        print("   export SERPER_API=your_key")
        print("   Continuing anyway...\n")

    level_filter = "All" if args.level == "all" else f"Level {args.level}"
    print(
        f"\nStarting GAIA benchmark  |  Level: {level_filter}  |  Split: {args.split}  |  Max: {args.max or 'all'}\n"
    )

    results = []
    for log_text, payload in gaia_benchmark_iter(
        level_filter, args.split, args.max, hf_token
    ):
        last_line = log_text.strip().split("\n")[-1]
        print(last_line)
        if isinstance(payload, list):  # final yield — payload is results list
            results = payload

    if not results:
        print("No results — check errors above.")
        return

    # ── Save submission JSONL (leaderboard format) ──────────────────────────
    sub_file = f"gaia_submission_{args.split}.jsonl"
    with open(sub_file, "w", encoding="utf-8") as f:
        for r in results:
            f.write(
                json.dumps(
                    {
                        "task_id": r["task_id"],
                        "model_answer": r["model_answer"],
                        "reasoning_trace": r["reasoning_trace"],
                    },
                    ensure_ascii=False,
                )
                + "\n"
            )

    # ── Save eval JSONL with ground truth (validation only) ─────────────────
    eval_file = None
    if args.split == "validation":
        correct, total = 0, 0
        eval_file = f"gaia_eval_{args.split}.jsonl"
        with open(eval_file, "w", encoding="utf-8") as f:
            for r in results:
                gt = r.get("ground_truth", "")
                is_correct = question_scorer(r["model_answer"], gt) if gt else None
                if is_correct is not None:
                    total += 1
                    correct += is_correct
                f.write(
                    json.dumps(
                        {
                            "task_id": r["task_id"],
                            "level": r["level"],
                            "question": r["question"][:120],
                            "model_answer": r["model_answer"],
                            "ground_truth": gt,
                            "correct": is_correct,
                        },
                        ensure_ascii=False,
                    )
                    + "\n"
                )

        score = correct / total * 100 if total else 0
        print(f"\n{'─' * 50}")
        print(f"Score: {correct}/{total}  ({score:.1f}%)")
        print(f"{'─' * 50}")

    print(f"\nSubmission → {sub_file}")
    if eval_file:
        print(f"Eval+GT    → {eval_file}")


if __name__ == "__main__":
    main()