import asyncio
import json
import os
import time
from abc import ABC
from http import HTTPStatus
from typing import IO, TYPE_CHECKING, Any, Literal, Optional, TypeVar, cast, overload

from js import AbortController, setTimeout  # type: ignore
from narada_core.actions.models import (
    AgenticMouseAction,
    AgenticMouseActionRequest,
    AgenticSelectorAction,
    AgenticSelectorRequest,
    AgenticSelectorResponse,
    AgenticSelectors,
    AgentResponse,
    AgentUsage,
    CloseWindowRequest,
    ExtensionActionRequest,
    ExtensionActionResponse,
    GetFullHtmlRequest,
    GetFullHtmlResponse,
    GetScreenshotRequest,
    GetScreenshotResponse,
    GetSimplifiedHtmlRequest,
    GetSimplifiedHtmlResponse,
    GetUrlRequest,
    GetUrlResponse,
    GoToUrlRequest,
    PrintMessageRequest,
    ReadGoogleSheetRequest,
    ReadGoogleSheetResponse,
    RecordedClick,
    WriteGoogleSheetRequest,
    parse_action_trace,
)
from narada_core.errors import (
    NaradaAgentTimeoutError_INTERNAL_DO_NOT_USE,
    NaradaError,
    NaradaTimeoutError,
)
from narada_core.models import (
    Agent,
    File,
    McpServer,
    RemoteDispatchChatHistoryItem,
    Response,
    UserResourceCredentials,
)
from pydantic import BaseModel
from pyodide.ffi import JsProxy, create_once_callable
from pyodide.http import pyfetch

# Magic variable injected by the JavaScript harness that stores the IDs of the current runnables
# in the stack on the frontend.

_cached_parent_run_ids: list[str] | None = None


def _parent_run_ids() -> list[str]:
    # `_narada_parent_run_ids` is a Pyodide `JsProxy` object injected by the JavaScript harness.
    # Before we can use it as a regular Python list, we need to call `.to_py()` on it.
    global _cached_parent_run_ids
    if _cached_parent_run_ids is None:
        _cached_parent_run_ids = cast(
            JsProxy,
            _narada_parent_run_ids,  # noqa: F821  # pyright: ignore[reportUndefinedVariable]
        ).to_py()
    return _cached_parent_run_ids


if TYPE_CHECKING:
    # Magic function injected by the JavaScript harness to get the current user's ID token.
    async def _narada_get_id_token() -> str: ...


_StructuredOutput = TypeVar("_StructuredOutput", bound=BaseModel)

_ResponseModel = TypeVar("_ResponseModel", bound=BaseModel)


class BaseBrowserWindow(ABC):
    _api_key: str | None
    _base_url: str
    _user_id: str | None
    _env: Literal["prod", "dev", None]
    _browser_window_id: str

    def __init__(
        self,
        *,
        api_key: str | None,
        base_url: str,
        user_id: str | None,
        env: Literal["prod", "dev", None] = "prod",
        browser_window_id: str,
    ) -> None:
        if api_key is None and (user_id is None or env is None):
            raise ValueError(
                "Either `api_key` or all of `user_id`, `user_id_token`, and `env` must be provided"
            )

        self._api_key = api_key
        self._base_url = base_url
        self._user_id = user_id
        self._env = env
        self._browser_window_id = browser_window_id

    @property
    def browser_window_id(self) -> str:
        return self._browser_window_id

    async def upload_file(self, *, file: IO) -> File:
        """Uploads a file that can be used as an attachment in a subsequent `agent` request.

        The file is temporarily saved in Narada cloud and expires after 1 day. It can only be
        accessed by the user who uploaded it.
        """
        raise NotImplementedError(
            "Uploading files is not supported in the browser environment"
        )

    @overload
    async def dispatch_request(
        self,
        *,
        prompt: str,
        agent: Agent | str = Agent.OPERATOR,
        clear_chat: bool | None = None,
        generate_gif: bool | None = None,
        output_schema: None = None,
        previous_request_id: str | None = None,
        chat_history: list[RemoteDispatchChatHistoryItem] | None = None,
        additional_context: dict[str, str] | None = None,
        time_zone: str = "America/Los_Angeles",
        user_resource_credentials: UserResourceCredentials | None = None,
        mcp_servers: list[McpServer] | None = None,
        variables: dict[str, str] | None = None,
        callback_url: str | None = None,
        callback_secret: str | None = None,
        callback_headers: dict[str, Any] | None = None,
        timeout: int = 1000,
    ) -> Response[None]: ...

    @overload
    async def dispatch_request(
        self,
        *,
        prompt: str,
        agent: Agent | str = Agent.OPERATOR,
        clear_chat: bool | None = None,
        generate_gif: bool | None = None,
        output_schema: type[_StructuredOutput],
        previous_request_id: str | None = None,
        chat_history: list[RemoteDispatchChatHistoryItem] | None = None,
        additional_context: dict[str, str] | None = None,
        time_zone: str = "America/Los_Angeles",
        user_resource_credentials: UserResourceCredentials | None = None,
        mcp_servers: list[McpServer] | None = None,
        variables: dict[str, str] | None = None,
        callback_url: str | None = None,
        callback_secret: str | None = None,
        callback_headers: dict[str, Any] | None = None,
        timeout: int = 1000,
    ) -> Response[_StructuredOutput]: ...

    async def dispatch_request(
        self,
        *,
        prompt: str,
        agent: Agent | str = Agent.OPERATOR,
        clear_chat: bool | None = None,
        generate_gif: bool | None = None,
        output_schema: type[BaseModel] | None = None,
        previous_request_id: str | None = None,
        chat_history: list[RemoteDispatchChatHistoryItem] | None = None,
        additional_context: dict[str, str] | None = None,
        time_zone: str = "America/Los_Angeles",
        user_resource_credentials: UserResourceCredentials | None = None,
        mcp_servers: list[McpServer] | None = None,
        variables: dict[str, str] | None = None,
        callback_url: str | None = None,
        callback_secret: str | None = None,
        callback_headers: dict[str, Any] | None = None,
        timeout: int = 1000,
    ) -> Response:
        """Low-level API for invoking an agent in the Narada extension side panel chat.

        The higher-level `agent` method should be preferred for most use cases.
        """
        deadline = time.monotonic() + timeout

        headers = {"Content-Type": "application/json"}
        if self._api_key is not None:
            headers["x-api-key"] = self._api_key
        else:
            assert self._user_id is not None
            assert self._env is not None

            headers["Authorization"] = f"Bearer {await _narada_get_id_token()}"
            headers["X-Narada-User-ID"] = self._user_id
            headers["X-Narada-Env"] = self._env

        agent_prefix = (
            agent.prompt_prefix() if isinstance(agent, Agent) else f"{agent} "
        )
        body: dict[str, Any] = {
            "prompt": agent_prefix + prompt,
            "browserWindowId": self.browser_window_id,
            "timeZone": time_zone,
            "parentRunIds": _parent_run_ids(),
        }
        if clear_chat is not None:
            body["clearChat"] = clear_chat
        if generate_gif is not None:
            body["saveScreenshots"] = generate_gif
        if output_schema is not None:
            body["responseFormat"] = {
                "type": "jsonSchema",
                "jsonSchema": output_schema.model_json_schema(),
            }
        if previous_request_id is not None:
            body["previousRequestId"] = previous_request_id
        if chat_history is not None:
            body["chatHistory"] = chat_history
        if additional_context is not None:
            body["additionalContext"] = additional_context
        if user_resource_credentials is not None:
            body["userResourceCredentials"] = user_resource_credentials
        if mcp_servers is not None:
            body["mcpServers"] = [
                server.model_dump(mode="json") for server in mcp_servers
            ]
        if variables is not None:
            body["variables"] = variables
        if callback_url is not None:
            body["callbackUrl"] = callback_url
        if callback_secret is not None:
            body["callbackSecret"] = callback_secret
        if callback_headers is not None:
            body["callbackHeaders"] = callback_headers

        try:
            controller = AbortController.new()
            signal = controller.signal

            setTimeout(create_once_callable(controller.abort), timeout * 1000)
            fetch_response = await pyfetch(
                f"{self._base_url}/remote-dispatch",
                method="POST",
                headers=headers,
                body=json.dumps(body),
                signal=signal,
            )

            if not fetch_response.ok:
                status = fetch_response.status
                text = await fetch_response.text()
                raise NaradaError(f"Failed to dispatch request: {status} {text}")

            request_id = (await fetch_response.json())["requestId"]

            while (now := time.monotonic()) < deadline:
                abort_controller = AbortController.new()
                signal = abort_controller.signal

                setTimeout(
                    create_once_callable(abort_controller.abort),
                    (deadline - now) * 1000,
                )
                fetch_response = await pyfetch(
                    f"{self._base_url}/remote-dispatch/responses/{request_id}",
                    headers=headers,
                    signal=signal,
                )

                if not fetch_response.ok:
                    status = fetch_response.status
                    text = await fetch_response.text()
                    raise NaradaError(f"Failed to poll for response: {status} {text}")

                response = await fetch_response.json()
                response["requestId"] = request_id

                if response["status"] != "pending":
                    response_content = response["response"]
                    if response_content is not None:
                        # Populate the `structuredOutput` field. This is a client-side field
                        # that's not directly returned by the API.
                        if output_schema is None:
                            response_content["structuredOutput"] = None
                        else:
                            structured_output = output_schema.model_validate_json(
                                response_content["text"]
                            )
                            response_content["structuredOutput"] = structured_output

                    return response

                # Poll every 3 seconds.
                await asyncio.sleep(3)
            else:
                raise NaradaAgentTimeoutError_INTERNAL_DO_NOT_USE(timeout)

        except asyncio.TimeoutError:
            raise NaradaAgentTimeoutError_INTERNAL_DO_NOT_USE(timeout)

    @overload
    async def agent(
        self,
        *,
        prompt: str,
        agent: Agent | str = Agent.OPERATOR,
        clear_chat: bool | None = None,
        generate_gif: bool | None = None,
        output_schema: None = None,
        time_zone: str = "America/Los_Angeles",
        mcp_servers: list[McpServer] | None = None,
        variables: dict[str, str] | None = None,
        timeout: int = 1000,
    ) -> AgentResponse[None]: ...

    @overload
    async def agent(
        self,
        *,
        prompt: str,
        agent: Agent | str = Agent.OPERATOR,
        clear_chat: bool | None = None,
        generate_gif: bool | None = None,
        output_schema: type[_StructuredOutput],
        time_zone: str = "America/Los_Angeles",
        mcp_servers: list[McpServer] | None = None,
        variables: dict[str, str] | None = None,
        timeout: int = 1000,
    ) -> AgentResponse[_StructuredOutput]: ...

    async def agent(
        self,
        *,
        prompt: str,
        agent: Agent | str = Agent.OPERATOR,
        clear_chat: bool | None = None,
        generate_gif: bool | None = None,
        output_schema: type[BaseModel] | None = None,
        time_zone: str = "America/Los_Angeles",
        mcp_servers: list[McpServer] | None = None,
        variables: dict[str, str] | None = None,
        timeout: int = 1000,
    ) -> AgentResponse:
        """Invokes an agent in the Narada extension side panel chat."""
        remote_dispatch_response = await self.dispatch_request(
            prompt=prompt,
            agent=agent,
            clear_chat=clear_chat,
            generate_gif=generate_gif,
            output_schema=output_schema,
            time_zone=time_zone,
            mcp_servers=mcp_servers,
            variables=variables,
            timeout=timeout,
        )
        response_content = remote_dispatch_response["response"]
        assert response_content is not None

        action_trace_raw = response_content.get("actionTrace")
        action_trace = (
            parse_action_trace(action_trace_raw)
            if action_trace_raw is not None
            else None
        )

        return AgentResponse(
            request_id=remote_dispatch_response["requestId"],
            status=remote_dispatch_response["status"],
            text=response_content["text"],
            structured_output=response_content.get("structuredOutput"),
            usage=AgentUsage.model_validate(remote_dispatch_response["usage"]),
            action_trace=action_trace,
        )

    async def agentic_selector(
        self,
        *,
        action: AgenticSelectorAction,
        selectors: AgenticSelectors,
        fallback_operator_query: str,
        # Larger default timeout because Operator can take a bit to run.
        timeout: int | None = 60,
    ) -> AgenticSelectorResponse:
        """Performs an action on an element specified by the given selectors, falling back to using
        the Operator agent if the selectors fail to match a unique element.

        Returns AgenticSelectorResponse with the value for 'get_text' and 'get_property' actions,
        otherwise returns None.
        """
        response_model = (
            AgenticSelectorResponse
            if action["type"] in {"get_text", "get_property"}
            else None
        )

        result = await self._run_extension_action(
            AgenticSelectorRequest(
                action=action,
                selectors=selectors,
                fallback_operator_query=fallback_operator_query,
            ),
            response_model,
            timeout=timeout,
        )

        if result is None:
            return AgenticSelectorResponse(value=None)

        return result

    async def agentic_mouse_action(
        self,
        *,
        action: AgenticMouseAction,
        recorded_click: RecordedClick,
        resize_window: Optional[bool] = True,
        fallback_operator_query: str,
        timeout: int | None = 60,
    ) -> None:
        """Performs a mouse action at the specified click coordinates, falling back to using
        the Operator agent if the click fails.
        """
        return await self._run_extension_action(
            AgenticMouseActionRequest(
                action=action,
                recorded_click=recorded_click,
                resize_window=resize_window or True,
                fallback_operator_query=fallback_operator_query,
            ),
            timeout=timeout,
        )

    async def close(self, *, timeout: int | None = None) -> None:
        """Gracefully closes the current browser window."""
        return await self._run_extension_action(CloseWindowRequest(), timeout=timeout)

    async def go_to_url(
        self, *, url: str, new_tab: bool = False, timeout: int | None = None
    ) -> None:
        """Navigates the active page in this window to the given URL."""
        return await self._run_extension_action(
            GoToUrlRequest(url=url, new_tab=new_tab), timeout=timeout
        )

    async def get_url(self, *, timeout: int | None = None) -> str:
        """Gets the URL of the current active page."""
        result = await self._run_extension_action(
            GetUrlRequest(),
            GetUrlResponse,
            timeout=timeout,
        )
        return result.url

    async def print_message(self, *, message: str, timeout: int | None = None) -> None:
        """Prints a message in the Narada extension side panel chat."""
        return await self._run_extension_action(
            PrintMessageRequest(message=message), timeout=timeout
        )

    async def read_google_sheet(
        self,
        *,
        spreadsheet_id: str,
        range: str,
        timeout: int | None = None,
    ) -> ReadGoogleSheetResponse:
        """Reads a range of cells from a Google Sheet."""
        return await self._run_extension_action(
            ReadGoogleSheetRequest(spreadsheet_id=spreadsheet_id, range=range),
            ReadGoogleSheetResponse,
            timeout=timeout,
        )

    async def write_google_sheet(
        self,
        *,
        spreadsheet_id: str,
        range: str,
        values: list[list[str]],
        timeout: int | None = None,
    ) -> None:
        """Writes a range of cells to a Google Sheet."""
        return await self._run_extension_action(
            WriteGoogleSheetRequest(
                spreadsheet_id=spreadsheet_id, range=range, values=values
            ),
            timeout=timeout,
        )

    async def get_full_html(self, *, timeout: int | None = None) -> GetFullHtmlResponse:
        """Gets the full HTML content of the current page."""
        return await self._run_extension_action(
            GetFullHtmlRequest(),
            GetFullHtmlResponse,
            timeout=timeout,
        )

    async def get_simplified_html(
        self, *, timeout: int | None = None
    ) -> GetSimplifiedHtmlResponse:
        """Gets the simplified HTML content of the current page."""
        return await self._run_extension_action(
            GetSimplifiedHtmlRequest(),
            GetSimplifiedHtmlResponse,
            timeout=timeout,
        )

    async def get_screenshot(
        self, *, timeout: int | None = None
    ) -> GetScreenshotResponse:
        """Takes a screenshot of the current browser window."""
        return await self._run_extension_action(
            GetScreenshotRequest(),
            GetScreenshotResponse,
            timeout=timeout,
        )

    @overload
    async def _run_extension_action(
        self,
        request: ExtensionActionRequest,
        response_model: None = None,
        *,
        timeout: int | None = None,
    ) -> None: ...

    @overload
    async def _run_extension_action(
        self,
        request: ExtensionActionRequest,
        response_model: type[_ResponseModel],
        *,
        timeout: int | None = None,
    ) -> _ResponseModel: ...

    async def _run_extension_action(
        self,
        request: ExtensionActionRequest,
        response_model: type[_ResponseModel] | None = None,
        *,
        timeout: int | None = None,
    ) -> _ResponseModel | None:
        headers = {"Content-Type": "application/json"}
        if self._api_key is not None:
            headers["x-api-key"] = self._api_key
        else:
            assert self._user_id is not None
            assert self._env is not None

            headers["Authorization"] = f"Bearer {await _narada_get_id_token()}"
            headers["X-Narada-User-ID"] = self._user_id
            headers["X-Narada-Env"] = self._env

        body = {
            "action": request.model_dump(),
            "browserWindowId": self.browser_window_id,
            "parentRunIds": _parent_run_ids(),
        }
        if timeout is not None:
            body["timeout"] = timeout

        fetch_response = await pyfetch(
            f"{self._base_url}/extension-actions",
            method="POST",
            headers=headers,
            body=json.dumps(body),
            # Don't specify `timeout` here as the (soft) timeout is handled by the server.
        )

        if fetch_response.status == HTTPStatus.GATEWAY_TIMEOUT:
            raise NaradaTimeoutError
        elif not fetch_response.ok:
            status = fetch_response.status
            text = await fetch_response.text()
            raise NaradaError(f"Failed to run extension action: {status} {text}")

        resp_json = await fetch_response.json()

        response = ExtensionActionResponse.model_validate(resp_json)
        if response.status == "error":
            raise NaradaError(response.error)

        if response_model is None:
            return None

        assert response.data is not None
        return response_model.model_validate_json(response.data)


class LocalBrowserWindow(BaseBrowserWindow):
    def __init__(self) -> None:
        env = os.environ.get("NARADA_ENV")
        if env is not None and env not in ("prod", "dev"):
            raise ValueError(f"Invalid environment: {env!r}")

        super().__init__(
            api_key=os.environ.get("NARADA_API_KEY"),
            base_url=os.getenv("NARADA_API_BASE_URL", "https://api.narada.ai/fast/v2"),
            user_id=os.environ.get("NARADA_USER_ID"),
            env=env,
            browser_window_id=os.environ["NARADA_BROWSER_WINDOW_ID"],
        )

    def __str__(self) -> str:
        return f"LocalBrowserWindow(browser_window_id={self.browser_window_id})"


class RemoteBrowserWindow(BaseBrowserWindow):
    def __init__(self, *, browser_window_id: str, api_key: str | None = None) -> None:
        super().__init__(
            api_key=api_key or os.environ["NARADA_API_KEY"],
            base_url=os.getenv("NARADA_API_BASE_URL", "https://api.narada.ai/fast/v2"),
            user_id=None,
            env=None,
            browser_window_id=browser_window_id,
        )

    def __str__(self) -> str:
        return f"RemoteBrowserWindow(browser_window_id={self.browser_window_id})"
