# isort: off
import datetime
import logging as log
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any
from contextlib import contextmanager
import json

import numpy as np
from smolagents import Model
from smolagents.memory import ActionStep, TaskStep

from android_world.registry import TaskRegistry  # type: ignore
from android_world.suite_utils import create_suite  # type: ignore
from screensuite.agents.client.android_env_client import AndroidEnvClient
from screensuite.agents.vision_agents.android_agent import AndroidAgent
from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmark
from screensuite.benchmarks.multistep.android_world.config import AndroidWorldConfig
from screensuite.benchmarks.multistep.config import AgentRunResult, process_answer_to_string
from screensuite.benchmarks.multistep.generation import append_answer
from screensuite.registry import registry

# Set logging verbosity to ERROR to reduce noise
os.environ["GRPC_VERBOSITY"] = "ERROR"
os.environ["GRPC_TRACE"] = "none"
os.environ["ANDROID_LOG_LEVEL"] = "ERROR"
os.environ["ANDROID_ACCESSIBILITY_LOG_LEVEL"] = "ERROR"
os.environ["ANDROID_ACCESSIBILITY_TREE_LOG"] = "false"

logger = log.getLogger()


class MatchFunction:
    """Taken from https://github.com/iMeanAI/WebCanvas/blob/b9f289128614cd99b97abd0bb9bfc3a45f0847e0/evaluate/step_score_js.py#L193"""

    def __init__(self):
        pass

    @staticmethod
    def include_match(input_answer, reference_answer) -> int:
        return 1 if reference_answer in input_answer else 0

    @staticmethod
    def exact_match(input_answer, reference_answer) -> int:
        return 1 if input_answer == reference_answer else 0


@contextmanager
def task_cleanup(agent: AndroidAgent, task_type: str, task_idx: int):
    """Context manager to ensure proper task cleanup."""
    yield
    try:
        agent.env.tear_down_task(task_type, task_idx)
    except Exception as e:
        logger.error(f"Error tearing down task '{task_type} {task_idx}': {str(e)}")
        agent.reset(go_home=True, respawn=True)


def answer_single_question(
    agent: AndroidAgent,
    env: AndroidEnvClient,
    model: Model,
    answers_file: str,
    task_type: str,
    task_idx: int,
    max_steps: int,
    max_retries: int = 3,
) -> AgentRunResult:
    retry_count = 0
    while retry_count < max_retries:
        try:
            agent.env.initialize_task(task_type, task_idx)
            task_goal = env.get_task_goal(task_type, task_idx)
            task_template = env.get_task_template(task_type, task_idx)
            start_on_home_screen = env.start_on_home_screen(task_type, task_idx)

            augmented_question = f"""Each step consists of one code snippet. After running each snippet, the task will be evaluated.
Do not use while True, recursion, or any kind of infinite or long-running loop.
Each snippet must complete quickly and cleanly without requiring loop execution. Use single-pass logic only.
Here is the task:

{task_goal}
"""

            start_time = time.time()

            agent_messages = []

            agent.reset(go_home=start_on_home_screen)

            # Run agent 🚀
            final_response = agent.run(augmented_question, reset=True, max_steps=max_steps)

            processed_answer = process_answer_to_string(final_response)

            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)

            end_time = time.time()

            annotated_example = AgentRunResult(
                model_id=model.model_id,  # type: ignore
                question=augmented_question,
                original_question=task_goal,
                answer=processed_answer,
                reference_answer=task_template,
                intermediate_steps=agent_messages,  # type: ignore
                start_time=start_time,
                end_time=end_time,
                token_counts=token_count,
            )
            append_answer(annotated_example.model_dump(), answers_file)
            return annotated_example
        except Exception as e:
            logger.error(f"Error processing task '{task_goal}': {str(e)}")
            end_time = time.time()
            if "500 Server Error" in str(e) or "Max retries exceeded with url" in str(e):
                logger.error(
                    f"500 Server Error for task '{task_goal}', respawning android environment and retrying task"
                )
                agent.reset(go_home=start_on_home_screen, respawn=True)
                if retry_count >= max_retries:
                    logger.error(f"500 Server Error for task '{task_goal}', maximum retries exceeded")
                    break
            retry_count += 1

    return AgentRunResult(
        model_id=model.model_id,  # type: ignore
        question=augmented_question,
        original_question=task_goal,
        answer=None,
        reference_answer=task_template,
        intermediate_steps=agent_messages,  # type: ignore
        start_time=start_time,
        end_time=end_time,
        token_counts=None,
    )


def process_single_task(
    agent: AndroidAgent,
    model: Model,
    file_name: str,
    task_type: str,
    task_idx: int,
    data_dir: str,
    thread_id: int,
    max_steps: int,
) -> tuple[float | None, AgentRunResult]:
    """
    Process a single task and return the score and result.

    Returns:
        Tuple of (score, agent_run_result) where score is None if task failed
    """
    agent.data_dir = os.path.join(data_dir, task_type, str(task_idx))

    with task_cleanup(agent, task_type, task_idx):
        try:
            agent_run_result = answer_single_question(
                agent, agent.env, model, file_name, task_type, task_idx, max_steps
            )

            try:
                score = agent.env.get_task_score(task_type, task_idx)
                return score, agent_run_result
            except Exception as e:
                logger.error(
                    f"Thread {thread_id} - Error getting score for task '{agent_run_result.original_question}': {str(e)}"
                )
                return 0.0, agent_run_result

        except Exception as e:
            print(f"Thread {thread_id} - Error processing task '{task_type} {task_idx}': {str(e)}")
            # Create a minimal result for failed tasks
            start_time = time.time()
            end_time = start_time
            failed_result = AgentRunResult(
                model_id=model.model_id,  # type: ignore
                question="",
                original_question=f"{task_type} {task_idx}",
                answer=None,
                reference_answer="",
                intermediate_steps=[],
                start_time=start_time,
                end_time=end_time,
                token_counts=None,
            )
            return None, failed_result


def process_task_chunk(
    model: Model, task_chunk: list[tuple[str, int, int]], file_name: str, max_steps: int, thread_id: int
) -> tuple[list[float], list[AgentRunResult], int]:
    """Process a chunk of tasks in a single thread.

    Args:
        model: The model to use
        task_chunk: List of (task_type, task_idx) tuples to process
        file_name: Output file name
        max_steps: Maximum steps per task
        thread_id: Thread identifier for logging

    Returns:
        Tuple of (accuracies, run_results, count_missing, count_sample)
    """
    accuracies: list[float] = []
    run_results: list[AgentRunResult] = []
    count_missing = 0
    data_dir = f"./output/android_world/thread_{thread_id}"

    # Create agent for this thread
    agent = AndroidAgent(
        model=model,
        data_dir=data_dir,
        max_steps=max_steps,
        planning_interval=10,
        transition_pause=None,
    )

    try:
        for task_type, task_idx_len, max_steps in task_chunk:
            for task_idx in range(task_idx_len):
                score, agent_run_result = process_single_task(
                    agent, model, file_name, task_type, task_idx, data_dir, thread_id, max_steps
                )

                run_results.append(agent_run_result)

                if isinstance(score, float):
                    accuracies.append(score)
                else:
                    count_missing += 1
    finally:
        agent.close()

    return accuracies, run_results, count_missing


class AndroidWorldBenchmark(HubBaseBenchmark[AndroidWorldConfig, None]):
    def evaluate(
        self,
        model: Model,
        evaluation_config: EvaluationConfig,
        env_config: None = None,
    ) -> Any:
        """
        Evaluate the model on the benchmark

        Args:
            model: The model to evaluate
            evaluation_config: Configuration for inference
            env_config: Environment configuration (unused)

        Returns:
            Evaluation results
        """
        date = datetime.date.today().isoformat()
        task_combinations = 2

        if evaluation_config.test_mode:
            evaluation_config.max_samples_to_test = 1
            evaluation_config.parallel_workers = 1

        if hasattr(model, "model_id"):
            run_name = f"{model.model_id.replace('/', '__')}__{date}"  # type: ignore
        else:
            run_name = f"{str(model)}__{date}"

        task_registry = TaskRegistry()
        aw_registry = task_registry.get_registry(task_registry.ANDROID_WORLD_FAMILY)
        suite = create_suite(
            task_registry=aw_registry,
            n_task_combinations=task_combinations,
            seed=42,  # Optional: for reproducibility
        )
        all_tasks = list(suite.keys())
        if evaluation_config.max_samples_to_test:
            all_tasks = all_tasks[: evaluation_config.max_samples_to_test]

        num_threads = (
            evaluation_config.parallel_workers
            if evaluation_config.parallel_workers <= len(all_tasks)
            else len(all_tasks)
        )

        with open(f"{os.path.dirname(__file__)}/task_metadata.json", "r") as f:
            metadata_file = json.load(f)
        optimal_steps: dict[str, int] = {}
        for task in metadata_file:
            optimal_steps[task["task_name"]] = int(task["optimal_steps"]) + 5

        base_chunk_size, remaining_tasks = divmod(len(all_tasks), num_threads)

        task_chunks = []
        current_idx = 0

        for thread_idx in range(num_threads):
            extra_task = 1 if thread_idx < remaining_tasks else 0
            chunk_size = base_chunk_size + extra_task

            chunk_tasks: list[str] = all_tasks[current_idx : current_idx + chunk_size]

            chunk_tuples = []
            for task_type in chunk_tasks:
                chunk_tuples.append((task_type, len(suite[task_type]), optimal_steps[task_type]))

            task_chunks.append(chunk_tuples)
            current_idx += chunk_size

        file_name = f"./output/android_world/{run_name}.jsonl"
        print(f"Starting processing with {num_threads} threads, writing output to '{file_name}'")

        # Process chunks in parallel
        all_accuracies: list[float] = []
        all_run_results: list[AgentRunResult] = []
        total_missing = 0
        total_sample = 0

        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            future_to_chunk = {
                executor.submit(process_task_chunk, model, chunk, file_name, self.config.max_steps, i): i
                for i, chunk in enumerate(task_chunks)
            }

            for future in as_completed(future_to_chunk):
                thread_id = future_to_chunk[future]
                try:
                    accuracies, run_results, count_missing = future.result()
                    all_accuracies.extend(accuracies)
                    all_run_results.extend(run_results)
                    total_missing += count_missing
                except Exception as e:
                    logger.error(f"Thread {thread_id} generated an exception: {str(e)}")
                    total_missing += len(task_chunks[thread_id]) * task_combinations
                finally:
                    total_sample += len(task_chunks[thread_id]) * task_combinations

        avg_accuracy = np.mean(all_accuracies) if all_accuracies else 0.0

        evaluation_results = {
            "avg_accuracy": float(avg_accuracy),
            "proportion_missing": float(total_missing / total_sample) if total_sample > 0 else 0.0,
            "count_sample": total_sample,
        }

        return BenchmarkResult(evaluation_results, "avg_accuracy")


android_world: AndroidWorldBenchmark = AndroidWorldBenchmark(
    name="android_world",
    config=AndroidWorldConfig(),
    tags=["android_world", "multistep", "hf_dataset", "online", "mobile", "android", "to_evaluate"],
)

if __name__ == "__main__":
    # test run
    from smolagents import OpenAIServerModel

    model = OpenAIServerModel(
        model_id="gpt-4o",
    )
    score = android_world.evaluate(
        model=model, evaluation_config=EvaluationConfig(parallel_workers=1, run_name="test_run", max_samples_to_test=1)
    )
    logger.info(f"Score: {score}")
else:
    registry.register(android_world)
