"""Local test harness: runs the GAIA agent on all 20 questions with caching."""
import json
import os
import pathlib
import sys
import time

import requests

from gaia_agent import API_BASE

ROOT = pathlib.Path(__file__).parent
RESULTS = ROOT / "results"
RESULTS.mkdir(exist_ok=True)

PROVIDER = "gemini"
args = sys.argv[1:]
if args and args[0] == "--groq":
    PROVIDER = "groq"
    args = args[1:]

if PROVIDER == "groq":
    from groq_agent import GroqBrain, answer_question_groq
    key_file = ROOT / ".groq_api_key"
    brain = GroqBrain(api_key=key_file.read_text().strip())
    answer_question = answer_question_groq
else:
    from gaia_agent import GeminiBrain, answer_question
    key_file = ROOT / ".gemini_api_key"
    brain = GeminiBrain(api_key=key_file.read_text().strip())

CACHE = RESULTS / ("answers_groq.json" if PROVIDER == "groq" else "answers.json")
TRACE = RESULTS / "trace.json"


def fetch_questions():
    return requests.get(f"{API_BASE}/questions", timeout=30).json()


def main(targets=None):
    cache = json.loads(CACHE.read_text()) if CACHE.exists() else {}
    trace = json.loads(TRACE.read_text()) if TRACE.exists() else {}
    questions = fetch_questions()

    for i, q in enumerate(questions):
        tid = q["task_id"]
        short = tid[:8]
        if targets and str(i + 1) not in targets and short not in targets:
            continue
        if tid in cache and cache[tid].get("answer") and not str(cache[tid]["answer"]).startswith("ERROR"):
            print(f"Q{i+1} [{short}] cached: {cache[tid]['answer']}", flush=True)
            continue
        print(f"Q{i+1} [{short}] running... file={q.get('file_name') or 'none'}", flush=True)
        t0 = time.time()
        try:
            ans = answer_question(brain, q, log=lambda m: print(m))
        except Exception as e:
            ans = f"ERROR: {e}"
        cache[tid] = {"answer": ans, "question": q["question"][:100], "secs": round(time.time() - t0, 1)}
        trace[tid] = cache[tid]
        print(f"Q{i+1} -> {ans}  ({cache[tid]['secs']}s)", flush=True)
        CACHE.write_text(json.dumps(cache, indent=2, ensure_ascii=False))
        TRACE.write_text(json.dumps(trace, indent=2, ensure_ascii=False))
        time.sleep(8)

    print("\n===== SUMMARY =====")
    for i, q in enumerate(questions):
        tid = q["task_id"]
        a = cache.get(tid, {}).get("answer", "(not run)")
        print(f"Q{i+1}: {a}")


if __name__ == "__main__":
    print(f"provider={PROVIDER} cache={CACHE.name}", flush=True)
    main(args or None)