"""
Generate answers locally and save them to answers.json (nothing is submitted here).

  python run_agent.py                 # run all questions not answered yet (resumes if interrupted)
  python run_agent.py --only 3 5      # (re)run just question indexes 3 and 5
  python run_agent.py --retry-empty   # re-run questions whose answer is empty/errored
  python run_agent.py --fresh         # ignore existing answers.json and start over

answers.json is plain JSON - open it and fix any answer by hand before submitting.
"""
import argparse
import json
from pathlib import Path

import requests

from agent import BasicAgent, download_task_file

API_URL = "https://agents-course-unit4-scoring.hf.space"
OUT = Path("answers.json")


def save(data):
    OUT.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--only", type=int, nargs="*", help="question indexes to (re)run")
    ap.add_argument("--retry-empty", action="store_true")
    ap.add_argument("--fresh", action="store_true")
    args = ap.parse_args()

    r = requests.get(f"{API_URL}/questions", timeout=30)
    r.raise_for_status()
    questions = r.json()
    print(f"Fetched {len(questions)} questions")

    data = {} if args.fresh or not OUT.exists() else json.loads(OUT.read_text(encoding="utf-8"))
    agent = BasicAgent()

    for i, q in enumerate(questions):
        tid = q["task_id"]
        if args.only is not None and args.only and i not in args.only:
            continue
        prev = data.get(tid, {}).get("answer", "")
        if not args.only:
            if prev and not data[tid].get("error"):
                continue  # already done
            if tid in data and not args.retry_empty and not prev and not data[tid].get("error"):
                continue

        print("\n" + "=" * 70)
        print(f"[{i}] {q['question'][:300]}")
        file_path = None
        if q.get("file_name"):
            file_path = download_task_file(API_URL, tid, q["file_name"])
            print("Attached file:", file_path)

        try:
            answer, error = agent(q["question"], file_path=file_path), None
        except Exception as e:  # noqa: BLE001
            answer, error = "", str(e)
            print("ERROR:", e)

        print(f"--> {answer!r}")
        data[tid] = {
            "index": i,
            "question": q["question"],
            "file_name": q.get("file_name") or "",
            "answer": answer,
            **({"error": error} if error else {}),
        }
        save(data)  # save after every question so a crash loses nothing

    done = sum(1 for v in data.values() if v.get("answer"))
    print(f"\nSaved {OUT} - {done}/{len(questions)} answered.")


if __name__ == "__main__":
    main()