import json
import logging
from multiprocessing import Manager, Process

from smolagents import Model

from osworld.desktop_env.desktop_env import DesktopEnv
from osworld.mm_agents.agent import PromptAgent
from osworld.run_config import DataConfig, EnvironmentConfig, ModelConfig, RunConfig
from screensuite.basebenchmark import BaseBenchmark, EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.multistep.osworld.config import (
    OSWorldConfig,
    OSWorldEnvironmentConfig,
)
from screensuite.benchmarks.multistep.osworld.utils import (
    distribute_tasks,
    get_unfinished,
    run_env_tasks,
)
from screensuite.registry import registry

logger = logging.getLogger()


class OSWorldBenchmark(BaseBenchmark[OSWorldConfig, OSWorldEnvironmentConfig]):
    def __init__(self, name: str, config: OSWorldConfig, tags: list[str]):
        super().__init__(name, config, tags=tags)
        self.run_config: RunConfig | None = None
        self.dataset: dict[str, list[str]] | None = None

    def load(self) -> None:
        """
        Load the dataset from Hugging Face

        Returns:
            The loaded dataset
        """
        if self.run_config is None:
            raise ValueError("Run config is not set")

        with open(
            self.run_config.data.test_all_meta_path,
            "r",
            encoding="utf-8",
        ) as f:
            test_all_meta = json.load(f)

        if self.run_config.data.domain != "all":
            test_all_meta = {self.run_config.data.domain: test_all_meta[self.run_config.data.domain]}

        self.dataset = get_unfinished(self.run_config, test_all_meta)

        left_info = ""
        for domain in self.dataset.keys():
            left_info += f"{domain}: {len(self.dataset[domain])}\n"
        logger.info(f"Left tasks:\n{left_info}")

    def evaluate(
        self,
        model: Model,
        evaluation_config: EvaluationConfig,
        env_config: OSWorldEnvironmentConfig = OSWorldEnvironmentConfig(),
    ) -> BenchmarkResult:
        """
        Evaluate the model on the benchmark

        Args:
            model: The model to evaluate
            parallel_workers: The number of processes to run in parallel

        Returns:
            Evaluation results
        """

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

        self.run_config = RunConfig(
            environment=EnvironmentConfig(
                provider_name=env_config.provider_name,
                path_to_vm=env_config.path_to_vm,
                region=env_config.region,
                os_type=env_config.os_type,
                headless=env_config.headless,
                action_space=env_config.action_space,
                observation_type=env_config.observation_type,
                screen_width=env_config.screen_width,
                screen_height=env_config.screen_height,
                sleep_after_execution=env_config.sleep_after_execution,
                max_steps=env_config.max_steps,
                max_trajectory_length=env_config.max_trajectory_length,
                num_envs=evaluation_config.parallel_workers,
            ),
            model=ModelConfig(
                _model=model,
                model_id=model.model_id,  # type: ignore
                temperature=self.config.temperature,
                top_p=self.config.top_p,
                max_tokens=self.config.max_tokens,
                stop_token=self.config.stop_token,
            ),
            data=DataConfig(
                test_config_base_dir=self.config.test_config_base_dir,
                test_all_meta_path=self.config.test_all_meta_path,
                result_dir=self.config.result_dir,
                domain=self.config.domain,
            ),
        )

        logger.info("Env Config: %s", env_config)

        if self.dataset is None:
            self.load()

        if self.run_config is None or self.dataset is None:
            raise ValueError("Fatal error: Run config or dataset is not set")

        distributed_tasks = distribute_tasks(
            self.dataset, evaluation_config.parallel_workers, evaluation_config.max_samples_to_test
        )

        # First, set up all environments
        logger.info("Setting up all environments...")
        envs = []
        agents = []

        for env_idx in range(evaluation_config.parallel_workers):
            logger.info(f"Setting up environment {env_idx + 1}/{evaluation_config.parallel_workers}")

            agent = PromptAgent(
                model=model,
                max_tokens=self.run_config.model.max_tokens,
                top_p=self.run_config.model.top_p,
                temperature=self.run_config.model.temperature,
                action_space=self.run_config.environment.action_space,
                observation_type=self.run_config.environment.observation_type,
                max_trajectory_length=self.run_config.environment.max_trajectory_length,
            )
            agents.append(agent)

            env = DesktopEnv(
                provider_name=self.run_config.environment.provider_name,
                path_to_vm=self.run_config.environment.path_to_vm,
                os_type=self.run_config.environment.os_type,
                region=self.run_config.environment.region,
                action_space=agent.action_space,
                screen_size=self.run_config.screen_size,
                headless=self.run_config.environment.headless,
                require_a11y_tree=self.run_config.environment.observation_type
                in [
                    "a11y_tree",
                    "screenshot_a11y_tree",
                    "som",
                ],
            )
            envs.append(env)

        logger.info("All environments are ready. Starting parallel task execution...")

        # Create a shared dictionary for results across processes
        with Manager() as manager:
            shared_results = manager.dict()

            processes = []
            for env_idx, (env, agent, env_tasks) in enumerate(zip(envs, agents, distributed_tasks)):
                p = Process(
                    target=run_env_tasks, args=(env_idx, env, agent, env_tasks, self.run_config, shared_results)
                )
                processes.append(p)
                p.start()

            for p in processes:
                p.join()

            results = dict(shared_results)

        all_scores = []
        total_missing = 0
        total_sample = 0

        for env_idx in range(len(processes)):
            if env_idx in results:
                thread_scores, thread_missing, thread_sample = results[env_idx]
                avg_thread_score = sum(thread_scores) / len(thread_scores) if thread_scores else 0.0
                logger.info(f"Thread {env_idx} results:")
                logger.info(f"  - Average score: {avg_thread_score:.3f}")
                logger.info(f"  - Missing/Failed tasks: {thread_missing}/{thread_sample}")

                all_scores.extend(thread_scores)
                total_missing += thread_missing
            else:
                logger.warning(f"Thread {env_idx} results not found in shared dictionary")
                total_missing += len(distributed_tasks[env_idx])
            total_sample += thread_sample

        avg_accuracy = sum(all_scores) / len(all_scores) if all_scores else 0.0
        logger.info(f"Overall average score: {avg_accuracy:.3f}")
        logger.info(f"Total missing/failed tasks: {total_missing}/{total_sample}")

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


osworld_benchmark = OSWorldBenchmark(
    name="osworld",
    config=OSWorldConfig.create(),
    tags=["osworld", "multistep", "online", "software", "web", "to_evaluate"],
)

if __name__ == "__main__":
    from smolagents import InferenceClientModel

    # model = OpenAIServerModel(model_id="gpt-4o")
    evaluation_config = EvaluationConfig(parallel_workers=4, run_name=None)
    # env_config = OSWorldEnvironmentConfig(provider_name="aws", region="us-east-1")
    model = InferenceClientModel(
        model_id="Qwen/Qwen2.5-VL-32B-Instruct",
        provider="fireworks-ai",
        max_tokens=4096,
    )
    env_config = OSWorldEnvironmentConfig(provider_name="docker")
    benchmark_result = osworld_benchmark.evaluate(model, evaluation_config, env_config=env_config)
    print(benchmark_result)
else:
    registry.register(osworld_benchmark)
