import copy
import logging
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Sequence

from datasets import Dataset
from smolagents import Model, Tool
from smolagents.memory import ActionStep, TaskStep
from smolagents.models import MessageRole
from smolagents.monitoring import LogLevel
from tqdm import tqdm

from screensuite.agents.vision_agents.desktop_agent import DesktopAgent
from screensuite.agents.vision_agents.e2b_agent import E2BVisionAgent
from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmarks.multistep.config import AgentRunResult
from screensuite.response_generation import append_answer, select_examples_to_test

logger = logging.getLogger(__name__)


def get_current_firefox_url(agent: E2BVisionAgent | DesktopAgent) -> str:
    # Command below sets the display, then searches for the firefox window, focuses on it, then copies the url to clipboard to return it
    if isinstance(agent, E2BVisionAgent):
        url_extraction_command = """export DISPLAY=:0.0
        export XAUTHORITY=$HOME/.Xauthority
        xdotool search --onlyvisible --class "firefox" windowfocus && xdotool key ctrl+l && xdotool key ctrl+c && sleep 0.5 && xclip -o -selection clipboard"""
    elif isinstance(agent, DesktopAgent):
        url_extraction_command = """export DISPLAY=:0.0
        export XAUTHORITY=$HOME/.Xauthority
        wmctrl -a Google Chrome && python3 -c 'import pyautogui, time, tkinter as tk; pyautogui.hotkey("ctrl", "l"); time.sleep(0.1); pyautogui.hotkey("ctrl", "c"); time.sleep(0.5); root = tk.Tk(); text = root.clipboard_get(); root.destroy(); print(text, end="")'"""
    else:
        raise ValueError(f"Unsupported agent type: {type(agent)}")
    return agent.run_arbitrary_command(url_extraction_command)


class ModifiedFinalAnswerTool(Tool):
    name = "final_answer"
    description = "Tests a function: if correct, returns it as the final answer; Else returns the test case that errored for solving."
    inputs = {"answer": {"type": "any", "description": "The final answer to the question"}}
    output_type = "string"

    def __init__(self, agent: E2BVisionAgent | DesktopAgent):
        super().__init__()
        self.agent = agent

    def forward(self, answer) -> str:
        try:
            return get_current_firefox_url(self.agent)
        except Exception as e:
            return "Could not extract the url from the browser: make sure that the browser is open!"


def answer_single_question(
    example: dict[str, Any],
    model: Model,
    answers_folder: str,
    reformulate_responses: bool = False,
    return_browser_url: bool = False,
    e2b_provider: bool = False,
) -> AgentRunResult:
    # Initialize agent

    augmented_question = example["question"]
    question_answers_path = Path(answers_folder) / augmented_question.lower().replace(" ", "_").replace(",", "")[:20]
    question_answers_path.mkdir(parents=True, exist_ok=True)

    if e2b_provider:
        api_key = os.getenv("E2B_API_KEY")
        assert api_key is not None, "E2B_API_KEY is not set"
        agent = E2BVisionAgent(
            model=model,
            e2b_api_key=api_key,
            data_dir=question_answers_path,
            tools=[],
            max_steps=30,
            additional_authorized_imports=["numpy"],
            verbosity_level=LogLevel.DEBUG,
        )

    else:
        agent = DesktopAgent(
            model=model,
            data_dir=question_answers_path,
            tools=[],
            max_steps=30,
            verbosity_level=LogLevel.DEBUG,
            additional_authorized_imports=["numpy"],
        )
    if return_browser_url:
        agent.tools["final_answer"] = ModifiedFinalAnswerTool(agent)

    start_time = time.time()
    agent_messages = []

    try:
        # Run agent 🚀
        final_response = agent.run(augmented_question)
        token_count = agent.monitor.get_total_token_counts()
        for memory_step in agent.memory.steps:
            if isinstance(memory_step, ActionStep):
                if hasattr(memory_step, "observations_images"):
                    memory_step.observations_images = None
            elif isinstance(memory_step, TaskStep):
                if hasattr(memory_step, "task_images"):
                    memory_step.task_images = None

        agent_messages = agent.write_memory_to_messages(summary_mode=True)

        if reformulate_responses:
            reformulated_answer = prepare_response(augmented_question, agent_messages, reformulation_model=model)
        else:
            reformulated_answer = final_response

        end_time = time.time()
        annotated_example = AgentRunResult(
            model_id=model.model_id,  # type: ignore
            question=augmented_question,
            original_question=example["question"],
            answer=reformulated_answer,  # type: ignore
            reference_answer=example["reference_answer"],
            intermediate_steps=agent_messages,
            start_time=start_time,
            end_time=end_time,
            token_counts=token_count,
        )
        append_answer(annotated_example.model_dump(), Path(answers_folder) / "answers.jsonl")
        return annotated_example
    except Exception as e:
        print("Error on ", augmented_question, e)
        end_time = time.time()
        return AgentRunResult(
            model_id=model.model_id,  # type: ignore
            question=augmented_question,
            original_question=example["question"],
            answer=None,
            reference_answer=example["reference_answer"],
            intermediate_steps=agent_messages,
            start_time=start_time,
            end_time=end_time,
            token_counts=None,
        )
    finally:
        # Clean up
        agent.close()


def prepare_response(original_task: str, inner_messages, reformulation_model: Model) -> str:
    """Shamelessly stolen from Microsoft Autogen team: thanks to them for this great resource!
    https://github.com/microsoft/autogen/blob/gaia_multiagent_v01_march_1st/autogen/browser_utils.py"""
    messages = [
        {
            "role": MessageRole.SYSTEM,
            "content": [
                {
                    "type": "text",
                    "text": f"""Earlier you were asked the following:

{original_task}

Your team then worked diligently to address that request. Read below a transcript of that conversation:""",
                }
            ],
        }
    ]

    try:
        for message in inner_messages:
            if not message.get("content"):
                continue
            message = copy.deepcopy(message)
            message["role"] = MessageRole.USER
            messages.append(message)
    except Exception:
        messages += [{"role": MessageRole.ASSISTANT, "content": str(inner_messages)}]

    # ask for the final answer
    messages.append(
        {
            "role": MessageRole.USER,
            "content": [
                {
                    "type": "text",
                    "text": f"""
Read the above conversation and output a FINAL ANSWER to the question. The question is repeated here for convenience:

{original_task}

To output the final answer, use 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.
ADDITIONALLY, your FINAL ANSWER MUST adhere to any formatting instructions specified in the original question (e.g., alphabetization, sequencing, units, rounding, decimal places, etc.)
If you are asked for a number, express it numerically (i.e., with digits rather than words), don't use commas, and DO NOT INCLUDE UNITS such as $ or USD or percent signs unless specified otherwise.
If you are asked for a string, don't use articles or abbreviations (e.g. for cities), unless specified otherwise. Don't output any final sentence punctuation such as '.', '!', or '?'.
If you are asked for a comma separated list, apply the above rules depending on whether the elements are numbers or strings.
If you are unable to determine the final answer, output 'FINAL ANSWER: Unable to determine'
""",
                }
            ],
        }
    )

    response = reformulation_model(messages).content or "Error in generation"

    final_answer = response.split("FINAL ANSWER: ")[-1].strip()
    print("> Reformulated answer: ", final_answer)
    return final_answer


def get_agent_responses(
    dataset: Sequence | Dataset,
    model: Model,
    evaluation_config: EvaluationConfig,
    reformulate_responses: bool = False,
    return_browser_url: bool = False,
    run_storage_folder: str = "./output/run",
) -> list[AgentRunResult | None]:
    """
    Get the responses from the model
    """
    responses: list[AgentRunResult | None] = []
    logger.warning(f"Starting processing and writing output to '{run_storage_folder}'")

    dataset_to_test = select_examples_to_test(dataset, evaluation_config)

    if evaluation_config.run_name is not None:
        print(f"Answers will be logged to {run_storage_folder}")
        answers_file_jsonl = Path(run_storage_folder) / "answers.jsonl"
        if not os.path.exists(answers_file_jsonl):
            os.makedirs(run_storage_folder, exist_ok=True)
            answers_file_jsonl.touch()

    # answered_questions: list[str] = []
    annotated_answered_questions: list[AgentRunResult] = []
    # 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"])
    #             annotated_answered_questions.append(json.loads(line))
    # NOTE: checkpointing is deactivated here
    examples_todo = [example for example in dataset_to_test]  # type: ignore

    if evaluation_config.test_mode:
        examples_todo = examples_todo[:1]  # Keep only the first example for testing
        for example in examples_todo:
            answer = answer_single_question(
                example, model, run_storage_folder, reformulate_responses, return_browser_url
            )  # type: ignore
            responses.append(answer)
        return responses

    with ThreadPoolExecutor(max_workers=evaluation_config.parallel_workers) as exe:
        futures = [
            exe.submit(
                answer_single_question,
                example,  # type: ignore
                model,
                run_storage_folder,
                reformulate_responses,
                return_browser_url,
            )
            for example in examples_todo
        ]
        for f in tqdm(as_completed(futures), total=len(examples_todo), desc="Processing tasks"):
            try:
                result = f.result()  # type: ignore
                responses.append(result)
                append_answer(result.model_dump(), answers_file_jsonl)
                annotated_answered_questions.append(result)
            except Exception as e:
                logger.error(f"Error processing task: {str(e)}", exc_info=True)
                responses.append(None)
    return responses
