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 screensuite.response_generation import append_answer
from smolagents import (
    AgentError,
    CodeAgent,
    GoogleSearchTool,
    HfApiModel,
    LiteLLMModel,
    VisitWebpageTool,
)
from smolagents.agents import ActionStep
from tqdm import tqdm

load_dotenv()

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(
        "--name",
        type=str,
        default=None,
        help="A custom name the evaluation. If not provided, left empty",
    )
    parser.add_argument(
        "--mode",
        type=str,
        default="generation",
        choices=["generation", "eval"],
    )
    # 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-type",
        type=str,
        default="HfApiModel",
        choices=["LiteLLMModel", "HfApiModel"],
        help="The model type to use (LiteLLMModel or HfApiModel)",
    )
    parser.add_argument(
        "--provider",
        type=str,
        default="together",
        help="The provider for inference",
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="Qwen/Qwen2.5-Coder-32B-Instruct",
        help="The model ID to use for the specified model type",
    )
    parser.add_argument(
        "--parallel-workers",
        type=int,
        default=32,
        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):
    agent = CodeAgent(
        tools=[GoogleSearchTool(provider="serper"), VisitWebpageTool()],
        model=model,
        additional_authorized_imports=["numpy", "sympy"],
        max_steps=10,
    )

    augmented_question = example["question"]
    if example["source"] == "SimpleQA":
        augmented_question += " Answer with only the final number."
    if example["source"] == "MATH":
        augmented_question += " Write code, not latex."

    start_time = time.time()

    try:
        # Run agent 🚀
        answer = str(agent.run(augmented_question))
        token_counts = agent.monitor.get_total_token_counts()
        # 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 = agent.write_memory_to_messages()
        end_time = time.time()
    except Exception as e:
        print("Error on ", augmented_question, e)
        intermediate_steps = []
        answer = "error"
        token_counts = {"input": 0, "output": 0}
    end_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    annotated_example = {
        "model_id": model.model_id,
        "question": augmented_question,
        "original_question": example["question"],
        "prediction": answer,
        "true_answer": example["true_answer"],
        "source": example["source"],
        "intermediate_steps": intermediate_steps,
        "start_time": start_time,
        "end_time": end_time,
        "token_counts": token_counts,
    }
    append_answer(annotated_example, answers_file)


def answer_questions(
    eval_ds,
    model,
    name: str = None,
    output_dir: str = "output",
    parallel_workers: int = 32,
):
    name = f"_{name}" if name else ""
    model_id = model.model_id

    os.makedirs(output_dir, exist_ok=True)

    file_name = f"{output_dir}/{model_id.replace('/', '__')}__{datetime.date.today().isoformat()}{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()


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

    if args.mode == "eval":
        gaia = datasets.load_dataset("smolagents/benchmark-v1", "gaia")
        math = datasets.load_dataset("smolagents/benchmark-v1", "math")
        simple_qa = datasets.load_dataset("smolagents/benchmark-v1", "simpleqa")
        eval_ds = datasets.concatenate_datasets([gaia["test"], math["test"], simple_qa["test"]])
    else:
        eval_ds = datasets.load_dataset("smolagents/trace-generation-tasks")["train"]

    if args.model_type == "LiteLLMModel":
        model = LiteLLMModel(
            model_id=args.model_id,
            max_completion_tokens=8192,
        )
    else:
        model = HfApiModel(
            model_id=args.model_id,
            provider=args.provider if not args.model_id.startswith("https") else None,
            max_tokens=8192,
        )

    answer_questions(
        eval_ds,
        model,
        args.name,
        parallel_workers=args.parallel_workers,
    )
