import argparse
import datetime
import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import datasets
from dotenv import load_dotenv
from reformulator import prepare_response
from screensuite.agents.vision_agents.desktop_agent import DesktopAgent
from screensuite.agents.vision_agents.e2b_agent import E2BVisionAgent
from screensuite.response_generation import append_answer
from smolagents import AgentError, GoogleSearchTool, HfApiModel
from smolagents.agents import ActionStep
from tqdm import tqdm

load_dotenv(override=True)
os.makedirs("output", exist_ok=True)

assert os.getenv("HF_TOKEN") is not None, "Need to set HF_TOKEN"

APPEND_ANSWER_LOCK = threading.Lock()


def parse_arguments():
    parser = argparse.ArgumentParser(description="Runs an agent powered by the given model on smolagent benchmark.")
    parser.add_argument(
        "--date",
        type=str,
        default=None,
        help="The date for the evaluation.",
    )
    parser.add_argument(
        "--run-name",
        type=str,
        default=None,
        help="The name for your run.",
    )
    parser.add_argument(
        "--eval-dataset",
        type=str,
        default="GeekAgents/GAIA-web",
    )
    # The eval dataset is gated, so you must first visit its page to request access: https://huggingface.co/datasets/smolagents-benchmark/benchmark-v1
    parser.add_argument(
        "--model-id",
        type=str,
        default="Qwen/Qwen2.5-VL-7B-Instruct",
        help="The model ID to use for the specified model type",
    )
    parser.add_argument(
        "--parallel-workers",
        type=int,
        default=4,
        help="The number of processes to run in parallel",
    )
    return parser.parse_args()


def serialize_agent_error(obj):
    if isinstance(obj, AgentError):
        return {"error_type": obj.__class__.__name__, "message": obj.message}
    else:
        return str(obj)


def answer_single_question(example, model, answers_file, e2b_provider: bool = False):
    # Initialize agent
    print("Using api key: ", os.getenv("E2B_API_KEY"))
    if e2b_provider:
        agent = E2BVisionAgent(
            model=model,
            data_dir="./output/gaiaweb",
            e2b_api_key=os.getenv("E2B_API_KEY"),
            tools=[GoogleSearchTool(provider="serper")],
            resolution=(1024, 768),
            max_steps=50,
            additional_authorized_imports=["numpy", "sympy"],
        )
    else:
        agent = DesktopAgent(
            model=model,
            data_dir="./output/gaiaweb",
            tools=[GoogleSearchTool(provider="serper")],
            max_steps=50,
            additional_authorized_imports=["numpy", "sympy"],
        )

    augmented_question = example["Question"]
    import random

    time.sleep(1.0 + random.random() * 2)

    start_time = time.time()

    try:
        # Run agent 🚀
        agent.run(augmented_question)
        token_count = agent.monitor.get_total_token_counts()

        agent_messages = agent.write_memory_to_messages(summary_mode=True)

        reformulated_answer = prepare_response(augmented_question, agent_messages, reformulation_model=model)

        # Remove memory from logs to make them more compact.
        for step in agent.memory.steps:
            if isinstance(step, ActionStep):
                step.agent_memory = None
        intermediate_steps = str(agent.memory.steps)
        end_time = time.time()
        annotated_example = {
            "model_id": model.model_id,
            "question": augmented_question,
            "original_question": example["Question"],
            "answer": reformulated_answer,
            "true_answer": example["Final answer"],
            "intermediate_steps": intermediate_steps,
            "start_time": start_time,
            "end_time": end_time,
            "token_counts": token_count,
            "level": example["Level"],
        }
        append_answer(annotated_example, answers_file)
    except Exception as e:
        print("Error on ", augmented_question, e)
    finally:
        # Clean up
        agent.close()


def answer_questions(
    eval_ds,
    model,
    date,
    output_dir: str = "output",
    push_answers_to_hub: bool = False,
    parallel_workers: int = 32,
    run_name: str = None,
):
    date = date or datetime.date.today().isoformat()
    model_id = model.model_id

    if run_name is None:
        run_name = f"{model_id.replace('/', '__')}__{date}"

    file_name = f"{output_dir}/{run_name}.jsonl"
    print(f"Starting processing and writing output to '{file_name}'")
    answered_questions = []
    if os.path.exists(file_name):
        with open(file_name, "r") as f:
            for line in f:
                answered_questions.append(json.loads(line)["original_question"])
    examples_todo = [example for example in eval_ds if example["Question"] not in answered_questions]

    print(f"Launching {parallel_workers} parallel workers.")
    with ThreadPoolExecutor(max_workers=parallel_workers) as exe:
        futures = [exe.submit(answer_single_question, example, model, file_name) for example in examples_todo]
        for f in tqdm(as_completed(futures), total=len(examples_todo), desc="Processing tasks"):
            f.result()

    # Use below instead of ThreadPoolExecutor for debugging
    # for example in tqdm(examples_todo):
    #     answer_single_question(example, model, file_name)

    print("All tasks processed.")


if __name__ == "__main__":
    args = parse_arguments()

    if args.parallel_workers > 20:
        raise ValueError("E2B free tier is limited to 20 parallel workers.")

    eval_ds = datasets.load_dataset(args.eval_dataset)["validation"]
    eval_ds = eval_ds.filter(lambda row: int(row["Level"]) == 1)
    print("Running model on ", len(eval_ds), "examples")

    from smolagents import HfApiModel

    model = HfApiModel("Qwen/Qwen2.5-VL-72B-Instruct", provider="nebius", max_tokens=8192)

    # model = LiteLLMModel("gpt-4o")

    answer_questions(
        eval_ds,
        model,
        args.date,
        parallel_workers=args.parallel_workers,
        run_name=args.run_name,
    )
