import os
from typing import Any, Callable, Literal, Optional

from pydantic import BaseModel, ConfigDict, Field, field_validator

from stagehand.schemas import AvailableModel


class StagehandConfig(BaseModel):
    """
    Configuration for the Stagehand client.

    Attributes:
        env (str): Environment type. 'BROWSERBASE' for remote usage
        api_key (Optional[str]): BrowserbaseAPI key for authentication.
        project_id (Optional[str]): Browserbase Project identifier.
        api_url (Optional[str]): Stagehand API URL.
        browserbase_session_create_params (Optional[BrowserbaseSessionCreateParams]): Browserbase session create params.
        browserbase_session_id (Optional[str]): Session ID for resuming Browserbase sessions.
        model_name (Optional[str]): Name of the model to use.
        model_api_key (Optional[str]): Model API key.
        model_client_options (Optional[dict[str, Any]]): Options for the model client.
        logger (Optional[Callable[[Any], None]]): Custom logging function.
        verbose (Optional[int]): Verbosity level for logs (1=minimal, 2=medium, 3=detailed).
        use_rich_logging (bool): Whether to use Rich for colorized logging.
        dom_settle_timeout_ms (Optional[int]): Timeout for DOM to settle (in milliseconds).
        enable_caching (Optional[bool]): Enable caching functionality.
        self_heal (Optional[bool]): Enable self-healing functionality.
        wait_for_captcha_solves (Optional[bool]): Whether to wait for CAPTCHA to be solved.
        act_timeout_ms (Optional[int]): Timeout for act commands (in milliseconds).
        headless (bool): Run browser in headless mode
        system_prompt (Optional[str]): System prompt to use for LLM interactions.
        local_browser_launch_options (Optional[dict[str, Any]]): Local browser launch options.
        use_api (bool): Whether to use API mode.
        experimental (bool): Enable experimental features.
    """

    env: Literal["BROWSERBASE", "LOCAL"] = "BROWSERBASE"
    api_key: Optional[str] = Field(
        None, alias="apiKey", description="Browserbase API key for authentication"
    )
    project_id: Optional[str] = Field(
        None, alias="projectId", description="Browserbase project ID"
    )
    api_url: Optional[str] = Field(
        os.environ.get("STAGEHAND_API_URL", "https://api.stagehand.browserbase.com/v1"),
        alias="apiUrl",
        description="Stagehand API URL",
    )
    model_api_key: Optional[str] = Field(
        None, alias="modelApiKey", description="Model API key"
    )
    model_client_options: Optional[dict[str, Any]] = Field(
        None,
        alias="modelClientOptions",
        description="Configuration options for the language model client (i.e. api_base)",
    )
    verbose: Optional[int] = Field(
        1,
        description="Verbosity level for logs: 0=minimal (ERROR), 1=medium (INFO), 2=detailed (DEBUG)",
    )
    logger: Optional[Callable[[Any], None]] = Field(
        None, description="Custom logging function"
    )
    use_rich_logging: Optional[bool] = Field(
        True, description="Whether to use Rich for colorized logging"
    )
    dom_settle_timeout_ms: Optional[int] = Field(
        3000,
        alias="domSettleTimeoutMs",
        description="Timeout for DOM to settle (in ms)",
    )
    headless: bool = Field(False, description="Run browser in headless mode")
    browserbase_session_create_params: Optional[dict[str, Any]] = Field(
        None,
        alias="browserbaseSessionCreateParams",
        description="Browserbase session create params",
    )
    enable_caching: Optional[bool] = Field(
        False, alias="enableCaching", description="Enable caching functionality"
    )
    browserbase_session_id: Optional[str] = Field(
        None,
        alias="browserbaseSessionID",
        description="Session ID for resuming Browserbase sessions",
    )
    model_name: Optional[str] = Field(
        AvailableModel.GPT_4_1_MINI,
        alias="modelName",
        description="Name of the model to use",
    )
    self_heal: Optional[bool] = Field(
        True, alias="selfHeal", description="Enable self-healing functionality"
    )
    wait_for_captcha_solves: Optional[bool] = Field(
        False,
        alias="waitForCaptchaSolves",
        description="Whether to wait for CAPTCHA to be solved",
    )
    system_prompt: Optional[str] = Field(
        None,
        alias="systemPrompt",
        description="System prompt to use for LLM interactions",
    )
    local_browser_launch_options: Optional[dict[str, Any]] = Field(
        {},
        alias="localBrowserLaunchOptions",
        description="Local browser launch options",
    )
    use_api: Optional[bool] = Field(
        True,
        alias=None,
        description="Whether to use the Stagehand API",
    )
    experimental: Optional[bool] = Field(
        False,
        alias=None,
        description="Whether to use experimental features",
    )

    model_config = ConfigDict(populate_by_name=True)

    @field_validator("browserbase_session_create_params", mode="before")
    @classmethod
    def validate_browserbase_params(cls, v, info):
        """Validate and convert browserbase session create params."""
        if isinstance(v, dict):
            # Make a copy to avoid mutating the original
            v = dict(v)

            # Add project_id if not present
            if "project_id" not in v:
                values = info.data
                project_id = values.get("project_id") or values.get("projectId")
                if project_id:
                    v["project_id"] = project_id

            # Handle browser_settings
            if "browser_settings" in v:
                # Make a copy of browser_settings to avoid mutating
                v["browser_settings"] = dict(v["browser_settings"])

                # Normalize 'os' key to lowercase
                if "os" in v["browser_settings"]:
                    v["browser_settings"]["os"] = v["browser_settings"]["os"].lower()
        return v

    def with_overrides(self, **overrides) -> "StagehandConfig":
        """
        Create a new config instance with the specified overrides.

        Args:
            **overrides: Key-value pairs to override in the config

        Returns:
            StagehandConfig: New config instance with overrides applied
        """
        config_dict = self.model_dump()
        config_dict.update(overrides)
        return StagehandConfig(**config_dict)


# Default configuration instance
default_config = StagehandConfig()
