import abc
import asyncio
import json
import os
import time
from http import HTTPStatus
from typing import Any, Generic, Literal, TypedDict, TypeVar, overload

from js import AbortController, setTimeout  # type: ignore
from pydantic import BaseModel
from pyodide.ffi import create_once_callable
from pyodide.http import pyfetch

from narada.actions.models import (
    AgentResponse,
    ExtensionActionRequest,
    GoToUrlRequest,
    GoToUrlResponse,
    PrintMessageRequest,
    PrintMessageResponse,
)
from narada.errors import NaradaError, NaradaTimeoutError
from narada.models import Agent, RemoteDispatchChatHistoryItem, UserResourceCredentials

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

_MaybeStructuredOutput = TypeVar("_MaybeStructuredOutput", bound=BaseModel | None)


class ResponseContent(TypedDict, Generic[_MaybeStructuredOutput]):
    text: str
    structuredOutput: _MaybeStructuredOutput


class Response(TypedDict, Generic[_MaybeStructuredOutput]):
    requestId: str
    status: Literal["success", "error"]
    response: ResponseContent[_MaybeStructuredOutput] | None
    createdAt: str
    completedAt: str | None


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


class BaseBrowserWindow(abc.ABC):
    api_key: str
    _browser_window_id: str

    def __init__(self, *, api_key: str, browser_window_id: str) -> None:
        self.api_key = api_key
        self._browser_window_id = browser_window_id

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

    @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,
        callback_url: str | None = None,
        callback_secret: str | None = None,
        timeout: int = 120,
    ) -> 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,
        callback_url: str | None = None,
        callback_secret: str | None = None,
        timeout: int = 120,
    ) -> 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,
        callback_url: str | None = None,
        callback_secret: str | None = None,
        timeout: int = 120,
    ) -> 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", "x-api-key": self.api_key}

        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,
        }
        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 callback_url is not None:
            body["callbackUrl"] = callback_url
        if callback_secret is not None:
            body["callbackSecret"] = callback_secret

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

            setTimeout(create_once_callable(controller.abort), timeout * 1000)
            fetch_response = await pyfetch(
                "https://api.narada.ai/fast/v2/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"https://api.narada.ai/fast/v2/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()

                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 NaradaTimeoutError

        except asyncio.TimeoutError:
            raise NaradaTimeoutError

    @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",
        timeout: int = 120,
    ) -> 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",
        timeout: int = 120,
    ) -> 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",
        timeout: int = 120,
    ) -> 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,
            timeout=timeout,
        )
        response_content = remote_dispatch_response["response"]
        assert response_content is not None

        return AgentResponse(
            status=remote_dispatch_response["status"],
            text=response_content["text"],
            structured_output=response_content.get("structuredOutput"),
        )

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

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

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

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

        fetch_response = await pyfetch(
            "https://api.narada.ai/fast/v2/extension-actions",
            method="POST",
            headers=headers,
            json=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}")

        return response_model.model_validate(await fetch_response.json())


class RemoteBrowserWindow(BaseBrowserWindow):
    def __init__(self, *, browser_window_id: str, api_key: str | None = None) -> None:
        api_key = api_key or os.environ["NARADA_API_KEY"]
        super().__init__(api_key=api_key, browser_window_id=browser_window_id)

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