import logging
import os
from typing import Literal

from pydantic import BaseModel, model_validator

from osworld.desktop_env.providers.aws.manager import IMAGE_ID_MAP
from screensuite.agents.client.osworld_vm_downloader import get_osworld_vm_path

logger = logging.getLogger()


class OSWorldEnvironmentConfig(BaseModel):
    """OSWorld environment configuration"""

    provider_name: str = "docker"
    _path_to_vm: str | None = None  # path local path to the osworld vm
    region: str | None = None
    os_type: Literal["Ubuntu", "Windows"] = "Ubuntu"
    headless: bool = False
    action_space: str = "pyautogui"
    observation_type: Literal["screenshot", "a11y_tree", "screenshot_a11y_tree", "som"] = "screenshot"
    screen_width: int = 1920
    screen_height: int = 1080
    sleep_after_execution: int = 0
    max_trajectory_length: int = 3
    max_steps: int = 15
    result_dir: str = "./results"

    @model_validator(mode="after")
    def validate_provider(self):
        if self.provider_name == "aws":
            region = self.region
            if region is None:
                logger.info("Region is not set, using us-east-1")
            if region not in IMAGE_ID_MAP.keys():
                raise ValueError(
                    f"Region {region} is not supported. "
                    "Please set the AMI ID for the region in the AWSVMManager file "
                    "(screensuite/src/screensuite/agents/osworld/desktop_env/providers/aws/manager.py)."
                )
            if os.environ.get("AWS_SUBNET_ID") is None:
                raise ValueError("AWS_SUBNET_ID is not set. Please set the AWS_SUBNET_ID environment variable.")
            if os.environ.get("AWS_SECURITY_GROUP_ID") is None:
                raise ValueError(
                    "AWS_SECURITY_GROUP_ID is not set. Please set the AWS_SECURITY_GROUP_ID environment variable."
                )
        return self

    @property
    def path_to_vm(self):
        if self._path_to_vm is None and self.provider_name == "docker":
            self._path_to_vm = get_osworld_vm_path(self.os_type)
        return self._path_to_vm


OSWORLD_DIR = f"{os.path.dirname(__file__)}/../../../../../osworld"


class OSWorldConfig(BaseModel):
    """OSWorld benchmark config"""

    test_config_base_dir: str = f"{OSWORLD_DIR}/evaluation_examples"
    domain: str = "all"
    test_all_meta_path: str = f"{OSWORLD_DIR}/evaluation_examples/test_all.json"
    result_dir: str = f"{OSWORLD_DIR}/results"
    temperature: float = 1.0
    top_p: float = 0.9
    max_tokens: int = 1500
    stop_token: str | None = None

    @classmethod
    def create(cls) -> "OSWorldConfig":
        return cls()
