"""
GAIA Benchmark Agent — Claude + web search
Produces a JSONL file ready to submit to:
  https://huggingface.co/spaces/gaia-benchmark/leaderboard

Requirements:
    pip install anthropic datasets huggingface_hub

Usage:
    export ANTHROPIC_API_KEY="sk-ant-..."
    huggingface-cli login          # needed to access the gated GAIA dataset
    python gaia_agent.py

Optional flags:
    --split       test | validation          (default: test)
    --level       1 | 2 | 3 | all           (default: all)
    --concurrency number of parallel calls   (default: 3)
    --no-search   disable web search tool
    --output      path to output JSONL       (default: submission.jsonl)
    --limit       max questions to run       (default: all)
"""

import argparse
import json
import os
import re
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import anthropic
from datasets import load_dataset
from huggingface_hub import snapshot_download

# ── Configuration ────────────────────────────────────────────────────────────

MODEL = "claude-sonnet-4-20250514"
MAX_TOKENS = 2048

SYSTEM_PROMPT = (
    "You are a general AI assistant. I will ask you a question. "
    "Report your thoughts, and finish your answer with the following template: "
    "FINAL ANSWER: [YOUR FINAL ANSWER]. "
    "YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma "
    "separated list of numbers and/or strings. "
    "If you are asked for a number, don't use comma to write your number neither use "
    "units such as $ or percent sign unless specified otherwise. "
    "If you are asked for a string, don't use articles, neither abbreviations "
    "(e.g. for cities), and write the digits in plain text unless specified otherwise. "
    "If you are asked for a comma separated list, apply the above rules depending of "
    "whether the element to be put in the list is a number or a string."
)

WEB_SEARCH_TOOL = {
    "type": "web_search_20250305",
    "name": "web_search",
}

# ── Helpers ──────────────────────────────────────────────────────────────────

def extract_final_answer(text: str) -> str:
    """Pull the text after 'FINAL ANSWER:' from the model response."""
    match = re.search(r"FINAL ANSWER:\s*(.+)", text, re.IGNORECASE)
    if match:
        return match.group(1).strip()
    # Fallback: last non-empty line
    lines = [l.strip() for l in text.strip().splitlines() if l.strip()]
    return lines[-1] if lines else text.strip()


def build_question_content(example: dict, data_dir: str) -> list:
    """
    Build the message content for a question.
    Attaches any associated file (PDF or image) as a base64 document/image block.
    """
    content = []

    file_path = example.get("file_path", "")
    if file_path:
        full_path = Path(data_dir) / file_path
        if full_path.exists():
            suffix = full_path.suffix.lower()
            try:
                with open(full_path, "rb") as f:
                    import base64
                    b64 = base64.standard_b64encode(f.read()).decode()

                if suffix == ".pdf":
                    content.append({
                        "type": "document",
                        "source": {
                            "type": "base64",
                            "media_type": "application/pdf",
                            "data": b64,
                        },
                    })
                elif suffix in {".png", ".jpg", ".jpeg", ".gif", ".webp"}:
                    media_map = {
                        ".png": "image/png", ".jpg": "image/jpeg",
                        ".jpeg": "image/jpeg", ".gif": "image/gif",
                        ".webp": "image/webp",
                    }
                    content.append({
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": media_map[suffix],
                            "data": b64,
                        },
                    })
                else:
                    # Plain text / CSV / other — embed as a document
                    text_data = full_path.read_text(errors="replace")
                    content.append({
                        "type": "text",
                        "text": f"[Attached file: {full_path.name}]\n{text_data}",
                    })
            except Exception as e:
                print(f"  Warning: could not read attachment {full_path}: {e}")

    content.append({"type": "text", "text": example["Question"]})
    return content


def run_single(
    client: anthropic.Anthropic,
    example: dict,
    data_dir: str,
    use_search: bool,
    retries: int = 3,
) -> dict:
    """Call the Claude API for one GAIA question and return a result dict."""
    task_id = example["task_id"]
    content = build_question_content(example, data_dir)

    kwargs = dict(
        model=MODEL,
        max_tokens=MAX_TOKENS,
        system=SYSTEM_PROMPT,
        messages=[{"role": "user", "content": content}],
    )
    if use_search:
        kwargs["tools"] = [WEB_SEARCH_TOOL]

    for attempt in range(1, retries + 1):
        try:
            response = client.messages.create(**kwargs)
            full_text = "".join(
                block.text for block in response.content if hasattr(block, "text")
            )
            answer = extract_final_answer(full_text)
            return {
                "task_id": task_id,
                "model_answer": answer,
                "reasoning_trace": full_text,
            }
        except anthropic.RateLimitError:
            wait = 2 ** attempt
            print(f"  Rate limit on {task_id}, waiting {wait}s…")
            time.sleep(wait)
        except Exception as e:
            if attempt == retries:
                print(f"  Failed {task_id} after {retries} attempts: {e}")
                return {
                    "task_id": task_id,
                    "model_answer": "",
                    "reasoning_trace": f"ERROR: {e}",
                }
            time.sleep(1)

# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="GAIA benchmark agent using Claude")
    parser.add_argument("--split",       default="test",
                        choices=["test", "validation"])
    parser.add_argument("--level",       default="all",
                        choices=["1", "2", "3", "all"])
    parser.add_argument("--concurrency", default=3, type=int)
    parser.add_argument("--no-search",   action="store_true")
    parser.add_argument("--output",      default="submission.jsonl")
    parser.add_argument("--limit",       default=None, type=int,
                        help="Run only the first N questions (useful for testing)")
    args = parser.parse_args()

    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        sys.exit("Error: ANTHROPIC_API_KEY environment variable not set.")

    client = anthropic.Anthropic(api_key=api_key)
    use_search = not args.no_search

    # ── Download GAIA dataset ────────────────────────────────────────────────
    print("Downloading GAIA dataset from Hugging Face…")
    print("(Make sure you have run 'huggingface-cli login' and accepted dataset terms)")
    try:
        data_dir = snapshot_download(
            repo_id="gaia-benchmark/GAIA",
            repo_type="dataset",
        )
    except Exception as e:
        sys.exit(f"Dataset download failed: {e}\n"
                 "Run 'huggingface-cli login' and accept the dataset terms at "
                 "https://huggingface.co/datasets/gaia-benchmark/GAIA")

    # ── Load split(s) ────────────────────────────────────────────────────────
    config_map = {
        "all": ["2023_level1", "2023_level2", "2023_level3"],
        "1":   ["2023_level1"],
        "2":   ["2023_level2"],
        "3":   ["2023_level3"],
    }
    configs = config_map[args.level]
    examples = []
    for cfg in configs:
        try:
            ds = load_dataset(data_dir, cfg, split=args.split)
            examples.extend(list(ds))
            print(f"  Loaded {len(ds)} questions from {cfg}/{args.split}")
        except Exception as e:
            print(f"  Warning: could not load {cfg}/{args.split}: {e}")

    if not examples:
        sys.exit("No questions loaded. Exiting.")

    if args.limit:
        examples = examples[: args.limit]

    print(f"\nRunning {len(examples)} questions | "
          f"model={MODEL} | concurrency={args.concurrency} | "
          f"web_search={use_search}\n")

    # ── Resume from existing output ──────────────────────────────────────────
    done_ids = set()
    output_path = Path(args.output)
    if output_path.exists():
        with open(output_path) as f:
            for line in f:
                try:
                    done_ids.add(json.loads(line)["task_id"])
                except Exception:
                    pass
        print(f"Resuming — {len(done_ids)} questions already answered.\n")

    pending = [ex for ex in examples if ex["task_id"] not in done_ids]
    if not pending:
        print("All questions already answered. Nothing to do.")
        return

    # ── Run agent ────────────────────────────────────────────────────────────
    total = len(pending)
    completed = 0
    errors = 0

    with open(output_path, "a") as out_f:
        with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
            futures = {
                executor.submit(run_single, client, ex, data_dir, use_search): ex
                for ex in pending
            }
            for future in as_completed(futures):
                ex = futures[future]
                completed += 1
                try:
                    result = future.result()
                    if result["model_answer"]:
                        status = "✓"
                    else:
                        status = "✗"
                        errors += 1
                    print(
                        f"[{completed}/{total}] {status} {result['task_id']} "
                        f"→ {result['model_answer'][:60]}"
                    )
                    out_f.write(json.dumps(result) + "\n")
                    out_f.flush()
                except Exception as e:
                    errors += 1
                    print(f"[{completed}/{total}] ✗ {ex['task_id']} — unexpected error: {e}")

    print(f"\nDone. {completed - errors}/{total} answered successfully.")
    print(f"Submission file: {output_path.resolve()}")
    print("\nNext step: upload to https://huggingface.co/spaces/gaia-benchmark/leaderboard")


if __name__ == "__main__":
    main()