"""
This module defines Pydantic models for OSWorld `Example`, `ExampleCorpusModel` concepts (and their dependent types).

The models are designed to handle different types of configuration actions, such as downloading files,
launching commands, interacting with Google Drive, and manipulating Chrome tabs. Additionally, it
includes models for evaluating conditions and handling process results.
"""

import json
from typing import Any, Literal

from pydantic import BaseModel, HttpUrl, RootModel, field_validator

# OSWorld uses both JSON bools and "true"/"false" strings
LenientBool = bool | Literal["true", "false"]


def _maybe_bool_from_lenient(input_bool: LenientBool | None) -> bool | None:
    return (
        None
        if input_bool is None
        else input_bool
        if isinstance(input_bool, bool)
        else True
        if input_bool == "true"
        else False
    )


class ConfigDownloadFile(BaseModel, frozen=True):
    url: str
    path: str


class ConfigDownloadParameters(BaseModel, frozen=True):
    files: list[ConfigDownloadFile]


class ConfigDownload(BaseModel, frozen=True):
    type: Literal["download"]
    parameters: ConfigDownloadParameters


class ConfigGoogleDriveParameters(BaseModel, frozen=True):
    settings_file: str
    operation: list[str]
    args: list[dict[str, Any]]


class ConfigGoogleDrive(BaseModel, frozen=True):
    type: Literal["googledrive"]
    parameters: ConfigGoogleDriveParameters


class ConfigLaunchParameters(BaseModel, frozen=True):
    command: str | list[str]
    shell: bool | None = None

    @field_validator("shell", mode="before")
    def transform_shell(cls, v: LenientBool | None) -> bool | None:
        return _maybe_bool_from_lenient(v)


class ConfigLaunch(BaseModel, frozen=True):
    type: Literal["launch"]
    parameters: ConfigLaunchParameters


class ConfigExecuteParameters(BaseModel, frozen=True):
    command: str | list[str]
    shell: bool | None = None
    stdout: str | None = None
    stderr: str | None = None

    @field_validator("shell", mode="before")
    def transform_shell(cls, v: LenientBool | None) -> bool | None:
        return _maybe_bool_from_lenient(v)


class ConfigExecuteOrCommand(BaseModel, frozen=True):
    type: Literal["execute", "command"]
    parameters: ConfigExecuteParameters


class ConfigOpenParameters(BaseModel, frozen=True):
    path: str


class ConfigOpen(BaseModel, frozen=True):
    type: Literal["open"]
    parameters: ConfigOpenParameters


class ConfigActivateWindowParameters(BaseModel, frozen=True):
    window_name: str
    strict: bool | None = None
    by_class: bool | None = None


class ConfigActivateWindow(BaseModel, frozen=True):
    type: Literal["activate_window"]
    parameters: ConfigActivateWindowParameters


class ConfigCloseWindowParameters(BaseModel, frozen=True):
    window_name: str
    strict: bool | None = None
    by_class: bool | None = None


class ConfigCloseWindow(BaseModel, frozen=True):
    type: Literal["close_window"]
    parameters: ConfigCloseWindowParameters


class ConfigSleepParameters(BaseModel, frozen=True):
    seconds: float


class ConfigSleep(BaseModel, frozen=True):
    type: Literal["sleep"]
    parameters: ConfigSleepParameters


class ConfigLogin(BaseModel, frozen=True):
    type: Literal["login"]
    parameters: dict[str, Any]


class ConfigChromeOpenTabsParameters(BaseModel, frozen=True):
    urls_to_open: list[HttpUrl]


class ConfigChromeOpenTabs(BaseModel, frozen=True):
    type: Literal["chrome_open_tabs"]
    parameters: ConfigChromeOpenTabsParameters


class ConfigChromeCloseTabsParameters(BaseModel, frozen=True):
    urls_to_close: list[HttpUrl]


class ConfigChromeCloseTabs(BaseModel, frozen=True):
    type: Literal["chrome_close_tabs"]
    parameters: ConfigChromeCloseTabsParameters


class ConfigUpdateBrowseHistoryEntry(BaseModel, frozen=True):
    url: HttpUrl
    title: str
    visit_time_from_now_in_seconds: int


class ConfigUpdateBrowseHistoryParameters(BaseModel, frozen=True):
    history: list[ConfigUpdateBrowseHistoryEntry]


class ConfigUpdateBrowseHistory(BaseModel, frozen=True):
    type: Literal["update_browse_history"]
    parameters: ConfigUpdateBrowseHistoryParameters


Config = (
    ConfigDownload
    | ConfigLaunch
    | ConfigExecuteOrCommand
    | ConfigOpen
    | ConfigActivateWindow
    | ConfigCloseWindow
    | ConfigGoogleDrive
    | ConfigLogin
    | ConfigChromeOpenTabs
    | ConfigChromeCloseTabs
    | ConfigUpdateBrowseHistory
    | ConfigSleep
)


class GetterVMCommandLineOrError(BaseModel, frozen=True):
    type: Literal["vm_command_line", "vm_command_error"]
    command: str | list[str]
    shell: bool | None = None

    @field_validator("shell", mode="before")
    def transform_shell(cls, v: LenientBool | None) -> bool | None:
        return _maybe_bool_from_lenient(v)


class GetterCloudFileSingle(BaseModel, frozen=True):
    type: Literal["cloud_file"]
    path: str
    dest: str
    multi: Literal[False] | None = None

    class Config:
        extra = "allow"


class GetterCloudFileMulti(BaseModel, frozen=True):
    type: Literal["cloud_file"]
    path: list[str]
    dest: list[str]
    multi: Literal[True]
    gives: list[int] | None = None

    @field_validator("multi", mode="before")
    def transform_shell(cls, v: LenientBool | None) -> bool | None:
        return _maybe_bool_from_lenient(v)

    class Config:
        extra = "allow"


class GetterBuiltInFunction(BaseModel, frozen=True):
    type: Literal[
        "accessibility_tree",
        "active_tab_html_parse",
        "active_tab_info",
        "active_tab_url_parse",
        "active_url_from_accessTree",
        "audio_in_slide",
        "background_image_in_slide",
        "bookmarks",
        "cache_file",
        "chrome_font_size",
        "chrome_language",
        "chrome_saved_address",
        "conference_city_in_order",
        "content_from_vm_file",
        "cookie_data",
        "data_delete_automacally",
        "default_search_engine",
        "default_video_player",
        "enable_do_not_track",
        "enable_enhanced_safety_browsing",
        "enabled_experiments",
        "extensions_installed_from_shop",
        "find_installed_extension_name",
        "find_unpacked_extension_path",
        "gimp_config_file",
        "googledrive_file",
        "gotoRecreationPage_and_get_html_content",
        "history",
        "info_from_website",
        "list_directory",
        "new_startup_page",
        "number_of_search_results",
        "open_tabs_info",
        "page_info",
        "pdf_from_url",
        "profile_name",
        "replay",
        "rule_relativeTime",
        "rule",
        "shortcuts_on_desktop",
        "single_file(_query, _path):",
        "time_diff_range",
        "url_dashPart",
        "vlc_config",
        "vlc_playing_info",
        "vm_file",
        "vm_screen_size",
        "vm_terminal_output",
        "vm_wallpaper",
        "vm_window_size",
        "vscode_config",
    ]

    class Config:
        extra = "allow"


Getter = (
    GetterVMCommandLineOrError
    | GetterCloudFileSingle
    | GetterCloudFileMulti
    | GetterBuiltInFunction
)


class Evaluator(BaseModel, frozen=True):
    postconfig: list[Config] | None = None
    func: str | list[str]
    conj: Literal["and", "or"] | None = None
    expected: Getter | list[Getter] | None = None
    result: Getter | list[Getter] | None = None
    options: dict[str, Any] | list[dict[str, Any]] | None = None


class Example(BaseModel, frozen=True):
    id: str
    snapshot: str
    source: str | list[str]
    instruction: str
    config: list[Config] | None = None
    related_apps: list[str]
    evaluator: Evaluator


def example_from_file(file_path: str) -> Example:
    with open(file_path, "r", encoding="utf-8") as file:
        return Example(**json.loads(file.read()))


class ExampleCorpusModel(RootModel[dict[str, list[str]]]):
    root: dict[str, list[str]]


class SubscriptableModel(BaseModel, frozen=True):
    def __getitem__(self, key: str) -> Any:
        # Check if key is a field and return its value
        if key in self.model_fields:  # Use `self.__fields__` in v1
            return getattr(self, key)
        raise KeyError(key)


class OSWorldProcessSuccessResult(SubscriptableModel, frozen=True):
    status: Literal["success"]
    output: str
    error: str
    returncode: int


class OSWorldProcessErrorResult(SubscriptableModel, frozen=True):
    status: Literal["error"]
    message: str


OSWorldProcessResult = OSWorldProcessSuccessResult | OSWorldProcessErrorResult
