import json
import logging
import os

from tqdm import tqdm

from osworld.desktop_env.desktop_env import DesktopEnv
from osworld.lib_run_single import run_single_example
from osworld.mm_agents.agent import PromptAgent
from osworld.run_config import RunConfig

logger = logging.getLogger()


GOOGLE_DRIVE_TOKEN_PATH = (
    f"{os.path.dirname(__file__)}/../../../../../osworld/evaluation_examples/settings/googledrive/token.json"
)


def distribute_tasks(
    test_all_meta: dict[str, list[str]], num_envs: int, max_samples_to_test: int | None = None
) -> list[dict[str, list[str]]]:
    """Distribute tasks evenly across multiple environments.

    Args:
        test_all_meta: Dictionary mapping domains to lists of example IDs
        num_envs: Number of environments to distribute tasks across
        test_mode: If True, only distribute one task per environment

    Returns:
        List of dictionaries, each containing tasks for one environment
    """

    # Create flat list of (domain, example_id) pairs
    tasks = [(domain, example_id) for domain, examples in test_all_meta.items() for example_id in examples]
    if max_samples_to_test is not None:
        tasks = tasks[:max_samples_to_test]

    # Calculate base tasks per environment using divmod
    base_tasks_per_env, remaining_tasks = divmod(len(tasks), num_envs)

    # Distribute tasks
    distributed_tasks = []
    current_idx = 0

    for env_idx in range(num_envs):
        # Add one extra task if there are remaining tasks to distribute
        extra_task = 1 if env_idx < remaining_tasks else 0
        tasks_for_this_env = base_tasks_per_env + extra_task

        env_tasks: dict[str, list[str]] = {}
        for domain, example_id in tasks[current_idx : current_idx + tasks_for_this_env]:
            env_tasks.setdefault(domain, []).append(example_id)

        distributed_tasks.append(env_tasks)
        current_idx += tasks_for_this_env

    return distributed_tasks


def get_unfinished(config: RunConfig, total_file_json: dict[str, list[str]]) -> dict[str, list[str]]:
    target_dir = os.path.join(
        config.data.result_dir,
        config.environment.action_space,
        config.environment.observation_type,
        config.model.model_id,
    )

    if not os.path.exists(target_dir):
        return total_file_json

    finished: dict[str, list[str]] = {}
    for domain in os.listdir(target_dir):
        finished[domain] = []
        domain_path = os.path.join(target_dir, domain)
        if os.path.isdir(domain_path):
            for example_id in os.listdir(domain_path):
                if example_id == "onboard":
                    continue
                example_path = os.path.join(domain_path, example_id)
                if os.path.isdir(example_path):
                    if "result.txt" not in os.listdir(example_path):
                        # empty all files under example_id
                        for file in os.listdir(example_path):
                            os.remove(os.path.join(example_path, file))
                    else:
                        finished[domain].append(example_id)

    if not finished:
        return total_file_json

    for domain, examples in finished.items():
        if domain in total_file_json:
            total_file_json[domain] = [x for x in total_file_json[domain] if x not in examples]

    return total_file_json


def run_env_tasks(
    env_idx: int, env: DesktopEnv, agent: PromptAgent, env_tasks: dict, config: RunConfig, shared_results: dict
):
    """Run tasks for a single environment."""
    logger.info(f"Executing tasks in environment {env_idx + 1}/{config.num_envs}")

    total_missing = 0
    total_sample = 0
    thread_scores: list[float] = []

    for domain in tqdm(env_tasks, desc=f"Env{env_idx + 1}-Domain"):
        for example_id in tqdm(env_tasks[domain], desc="Example", leave=False):
            config_file = os.path.join(config.data.test_config_base_dir, f"examples/{domain}/{example_id}.json")
            with open(config_file, "r", encoding="utf-8") as f:
                example = json.load(f)

            logger.info(f"[Env {env_idx + 1}][Domain]: {domain}")
            logger.info(f"[Env {env_idx + 1}][Example ID]: {example_id}")
            logger.info(f"[Env {env_idx + 1}][Instruction]: {example['instruction']}")

            example_result_dir = os.path.join(
                config.data.result_dir,
                config.environment.action_space,
                config.environment.observation_type,
                config.model.model_id,
                domain,
                example_id,
            )
            os.makedirs(example_result_dir, exist_ok=True)

            try:
                score = run_single_example(
                    agent,
                    env,
                    example,
                    config.environment.max_steps,
                    example["instruction"],
                    config,
                    example_result_dir,
                    thread_scores,
                )
            except Exception as e:
                logger.error(f"Exception in Env{env_idx + 1} {domain}/{example_id}: {e}", exc_info=True)
                env.controller.end_recording(os.path.join(example_result_dir, "recording.mp4"))
                with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
                    f.write(json.dumps({"Error": f"Time limit exceeded in {domain}/{example_id}"}))
                    f.write("\n")
                total_missing += 1
            total_sample += 1

    env.close()

    # Store results in shared dictionary
    shared_results[env_idx] = (thread_scores, total_missing, total_sample)
