import asyncio
import json
import sys
import tempfile
import time
import traceback
import warnings
from collections.abc import Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Any, Unpack, overload

# import websockets
from halo import Halo  # pyright: ignore[reportMissingTypeStubs]
from notte_core.agent_types import AgentCompletion
from notte_core.common.logging import logger
from notte_core.common.notifier import BaseNotifier
from notte_core.common.telemetry import track_usage
from notte_core.utils.webp_replay import WebpReplay
from pydantic import BaseModel, Field, ValidationError
from typing_extensions import final

from notte_sdk.endpoints.base import BaseClient, NotteEndpoint
from notte_sdk.endpoints.functions import NotteFunction
from notte_sdk.endpoints.personas import NottePersona
from notte_sdk.endpoints.sessions import RemoteSession
from notte_sdk.endpoints.vaults import NotteVault
from notte_sdk.types import (
    AgentCreateRequestDict,
    AgentFunctionCodeResponse,
    AgentListRequest,
    AgentListRequestDict,
    AgentResponse,
    AgentRunRequest,
    AgentRunRequestDict,
    AgentStatus,
    AgentStatusRequest,
    AgentStatusResponse,
    AgentWorkflowCodeRequest,
    GetFunctionResponse,
    ReplayResponse,
    SdkAgentCreateRequest,
    SdkAgentStartRequestDict,
)

# Conditional imports for Pyodide vs native Python
RUNNING_IN_PYODIDE = "pyodide" in sys.modules

if RUNNING_IN_PYODIDE:
    import js  # pyright: ignore[reportMissingImports]
    from pyodide.ffi import (  # pyright: ignore[reportMissingImports]
        create_proxy,  # pyright: ignore[reportUnknownVariableType]
    )
else:
    from websockets.sync import client as sync_client


if TYPE_CHECKING:
    from notte_sdk.client import NotteClient


class SdkAgentStartRequest(SdkAgentCreateRequest, AgentRunRequest):
    pass


class LegacyAgentStatusResponse(AgentStatusResponse):
    """
    This class is used to handle the legacy agent status response.
    The rationale is that we are likely to change the `AgentStepResponse` in the future and we want to be able to handle the legacy response.
    This is a temporary solution to avoid breaking changes.
    """

    steps: list[dict[str, Any]] = Field(default_factory=list)


@final
class AgentsClient(BaseClient):
    """
    Client for the Notte API.

    Note: this client is only able to handle one session at a time.
    If you need to handle multiple sessions, you need to create a new client for each session.
    """

    # Session
    AGENT_START = "start"
    AGENT_START_CUSTOM = "start/custom"
    AGENT_STOP = "{agent_id}/stop?session_id={session_id}"
    AGENT_STATUS = "{agent_id}"
    AGENT_FUNCTION = "{agent_id}/workflow/code"
    AGENT_LIST = ""
    AGENT_LOGS_WS = "{agent_id}/debug/logs?token={token}&session_id={session_id}"

    def __init__(
        self,
        root_client: "NotteClient",
        api_key: str | None = None,
        server_url: str | None = None,
        verbose: bool = False,
    ):
        """
        Initialize an AgentsClient instance.

        Configures the client to use the "agents" endpoint path and sets optional API key and server URL for authentication and server configuration. The initial state has no recorded agent response.

        Args:
            api_key: Optional API key for authenticating requests.
        """
        super().__init__(
            root_client=root_client,
            base_endpoint_path="agents",
            server_url=server_url,
            api_key=api_key,
            verbose=verbose,
        )

    @staticmethod
    def _agent_start_endpoint() -> NotteEndpoint[AgentResponse]:
        """
        Returns an endpoint for running an agent.

        Creates a NotteEndpoint configured with the AGENT_START path, a POST method, and an expected AgentResponse.
        """
        return NotteEndpoint(path=AgentsClient.AGENT_START, response=AgentResponse, method="POST")

    @staticmethod
    def _agent_start_custom_endpoint() -> NotteEndpoint[AgentResponse]:
        """
        Returns an endpoint for running an agent.
        """
        return NotteEndpoint(path=AgentsClient.AGENT_START_CUSTOM, response=AgentResponse, method="POST")

    @staticmethod
    def _agent_stop_endpoint(
        agent_id: str | None = None, session_id: str | None = None
    ) -> NotteEndpoint[AgentResponse]:
        """
        Constructs a DELETE endpoint for stopping an agent.

        If an agent ID is provided, it is inserted into the endpoint URL. The returned
        endpoint is configured with the DELETE HTTP method and expects an AgentStatusResponse.

        Args:
            agent_id (str, optional): The identifier of the agent to stop. If omitted,
                the URL template will remain unformatted.

        Returns:
            NotteEndpoint[AgentResponse]: The endpoint object for stopping the agent.
        """
        path = AgentsClient.AGENT_STOP
        if agent_id is not None:
            path = path.format(agent_id=agent_id, session_id=session_id)
        return NotteEndpoint(path=path, response=AgentStatusResponse, method="DELETE")

    @staticmethod
    def _agent_status_endpoint(agent_id: str | None = None) -> NotteEndpoint[LegacyAgentStatusResponse]:
        """
        Creates an endpoint for retrieving an agent's status.

        If an agent ID is provided, formats the endpoint path to target that specific agent.

        Args:
            agent_id: Optional identifier of the agent; if specified, the endpoint path will include this ID.

        Returns:
            NotteEndpoint configured with the GET method and AgentStatusResponse as the expected response.
        """
        path = AgentsClient.AGENT_STATUS
        if agent_id is not None:
            path = path.format(agent_id=agent_id)
        return NotteEndpoint(path=path, response=LegacyAgentStatusResponse, method="GET")

    @staticmethod
    def _agent_function_endpoint(agent_id: str | None = None) -> NotteEndpoint[AgentFunctionCodeResponse]:
        """
        Creates an endpoint for retrieving an agent's script.

        If an agent ID is provided, formats the endpoint path to target that specific agent.

        Args:
            agent_id: Optional identifier of the agent; if specified, the endpoint path will include this ID.

        Returns:
            NotteEndpoint configured with the GET method and AgentFunctionCodeResponse as the expected response.
        """
        path = AgentsClient.AGENT_FUNCTION
        if agent_id is not None:
            path = path.format(agent_id=agent_id)
        return NotteEndpoint(path=path, response=AgentFunctionCodeResponse, method="GET")

    @staticmethod
    def _agent_list_endpoint(params: AgentListRequest | None = None) -> NotteEndpoint[AgentResponse]:
        """
        Creates a NotteEndpoint for listing agents.

        Returns an endpoint configured with the agent listing path and a GET method.
        The optional params argument provides filtering or pagination details for the request.
        """
        return NotteEndpoint(
            path=AgentsClient.AGENT_LIST,
            response=AgentResponse,
            method="GET",
            request=None,
            params=params,
        )

    def start(self, **data: Unpack[SdkAgentStartRequestDict]) -> AgentResponse:
        """
        Start an agent with the specified request parameters.

        Validates the provided data using the AgentRunRequest model, sends a run request through the
        designated endpoint, updates the last agent response, and returns the resulting AgentResponse.

        Args:
            **data: Keyword arguments representing the fields of an AgentRunRequest.

        Returns:
            AgentResponse: The response obtained from the agent run request.
        """
        request = SdkAgentStartRequest.model_validate(data)
        response = self.request(AgentsClient._agent_start_endpoint().with_request(request))
        return response

    def wait(
        self,
        agent_id: str,
        polling_interval_seconds: int = 10,
        max_attempts: int = 30,
    ) -> AgentStatusResponse:
        """
        Waits for the specified agent to complete.

        Args:
            agent_id: The identifier of the agent to wait for.
            polling_interval_seconds: The interval between status checks.
            max_attempts: The maximum number of attempts to check the agent's status.

        Returns:
            AgentStatusResponse: The response from the agent status check.
        """
        last_step = 0
        for _ in range(max_attempts):
            response = self.status(agent_id=agent_id)
            if len(response.steps) > last_step:
                for _step in response.steps[last_step:]:
                    step = AgentCompletion.model_validate(_step)
                    step.live_log_state()
                    if step.is_completed():
                        logger.info(f"Agent {agent_id} completed in {len(response.steps)} steps")
                        return response

                last_step = len(response.steps)

            if response.status == AgentStatus.closed:
                return response

            spinner = None
            try:
                if not WebpReplay.in_notebook():
                    spinner = Halo(
                        text=f"Waiting {polling_interval_seconds} seconds for agent to complete (current step: {last_step})...",
                    )
                time.sleep(polling_interval_seconds)

            finally:
                if spinner is not None:
                    _ = spinner.succeed()  #  pyright: ignore[reportUnknownMemberType]

        raise TimeoutError("Agent did not complete in time")

    def _process_ws_message(
        self,
        message: str,
        agent_id: str,
        log: bool,
        counter: list[int],
    ) -> tuple[AgentCompletion | AgentStatusResponse | None, bool]:
        """
        Process a websocket message. Returns (response, should_stop).

        Args:
            message: The raw websocket message string.
            agent_id: The agent identifier for logging.
            log: Whether to log the agent steps.
            counter: A mutable list containing [step_count] to track step number.

        Returns:
            Tuple of (response, should_stop) where response is the parsed message
            and should_stop indicates if the agent has completed.
        """
        try:
            dic = json.loads(message)
            response = None

            # output from validator
            if isinstance(dic, dict) and "validation" in dic:
                logger.opt(colors=True).info("<g>{message}</g>", message=dic["validation"])

            # termination message
            elif isinstance(dic, dict) and "status" in dic:
                if dic["status"] == "agent_stop":
                    # Parse the agent status response from the message
                    if "agent" in dic:
                        agent_status = AgentStatusResponse.model_validate(dic["agent"])
                        return (agent_status, True)
                    # Fallback: no agent field, this shouldn't happen but handle gracefully
                    return (None, True)

            # actual step
            else:
                if isinstance(dic, dict):
                    response = AgentCompletion.model_validate(dic)
                else:
                    # Unexpected: log and skip
                    logger.warning(f"Expected dict, got {type(dic).__name__}: {message[:200]}")
                    return (None, False)
                if log:
                    logger.opt(colors=True).info(
                        "✨ <r>Step {counter}</r> <y>(agent: {agent_id})</y>",
                        counter=(counter[0] + 1),
                        agent_id=agent_id,
                    )
                    response.live_log_state()
                counter[0] += 1

            return (response, False)

        except (json.JSONDecodeError, ValidationError, TypeError, ValueError) as e:
            if "error" in message and "last action failed with error" not in message:
                logger.error(f"Error in agent logs: {e} {agent_id} {message}")
            elif agent_id in message and "agent_id" in message:
                logger.error(f"Error parsing AgentStatusResponse for message: {message}: {e}")
            else:
                logger.error(f"Error parsing agent logs for message: {message}: {e}")
            return (None, False)

    def watch_logs(self, agent_id: str, session_id: str, log: bool = True) -> AgentStatusResponse | None:
        """
        Watch the logs of the specified agent.
        """
        endpoint = NotteEndpoint(path=AgentsClient.AGENT_LOGS_WS, response=BaseModel, method="GET")
        wss_url = self.request_path(endpoint).format(agent_id=agent_id, token=self.token, session_id=session_id)
        wss_url = wss_url.replace("https://", "wss://").replace("http://", "ws://")

        counter = [0]  # mutable container for step count

        if RUNNING_IN_PYODIDE:
            raise NotImplementedError(
                "Synchronous watch_logs is not supported in Pyodide. Use async_watch_logs() or async_watch_logs_and_wait() instead."
            )

        # Use native Python sync websockets library
        try:
            with sync_client.connect(  # pyright: ignore[reportPossiblyUnboundVariable]
                uri=wss_url,
                open_timeout=30,
                ping_interval=5,
                ping_timeout=40,
                close_timeout=5,
                max_size=5 * (2**20),  # 5MB max size
            ) as websocket:
                for message in websocket:
                    if not isinstance(message, str):
                        logger.warning(f"Expected str message, got {type(message).__name__}. Skipping.")
                        continue
                    response, should_stop = self._process_ws_message(message, agent_id, log, counter)

                    if should_stop:
                        # If we got an AgentStatusResponse, return it; otherwise return None (failure)
                        if isinstance(response, AgentStatusResponse):
                            return response
                        return None
        except ConnectionError as e:
            logger.error(f"Connection error: {agent_id} {e}")
            return None
        except Exception as e:
            logger.error(f"Unexpected websocket processing error: {agent_id} {e} {traceback.format_exc()}")
            raise

        return None

    def watch_logs_and_wait(self, agent_id: str, session_id: str, log: bool = True) -> AgentStatusResponse:
        """
        Execute a task with the agent and wait for completion.

        Args:
            agent_id (str): The agent identifier.
            session_id (str): The session identifier.
            log (bool): Whether to log the agent steps.

        Returns:
            AgentStatusResponse: The response from the completed agent execution.
        """
        # In Pyodide, sync WebSocket connections aren't supported and there's always a running event loop
        if RUNNING_IN_PYODIDE:
            raise RuntimeError(
                "watch_logs_and_wait() cannot be used in Pyodide. Use `await async_watch_logs_and_wait(...)` instead."
            )

        try:
            response = self.watch_logs(agent_id=agent_id, session_id=session_id, log=log)
            if response is not None:
                return response
            # If we didn't get a response, poll status until agent is closed
            logger.warning(f"[Agent] {agent_id} did not return status response. Polling status until closed.")
            max_wait_secs = 300
            waited = 0
            while waited < max_wait_secs:
                status = self.status(agent_id=agent_id)
                if status.status == AgentStatus.closed:
                    return status
                time.sleep(1)
                waited += 1
            raise TimeoutError(f"Agent {agent_id} did not reach a terminal state within {max_wait_secs}s")

        except KeyboardInterrupt:
            status = self.status(agent_id=agent_id)
            if status.status != AgentStatus.closed:
                _ = self.stop(agent_id=agent_id, session_id=session_id)
            raise

    async def async_watch_logs(self, agent_id: str, session_id: str, log: bool = True) -> AgentStatusResponse | None:
        """
        Watch the logs of the specified agent asynchronously.

        This method is required for Pyodide environments where synchronous WebSocket
        connections are not supported.

        Args:
            agent_id (str): The agent identifier.
            session_id (str): The session identifier.
            log (bool): Whether to log the agent steps.

        Returns:
            AgentStatusResponse | None: The final agent status, or None if failed.
        """
        if not RUNNING_IN_PYODIDE:
            raise NotImplementedError("async_watch_logs is only supported in Pyodide. Use watch_logs instead.")

        endpoint = NotteEndpoint(path=AgentsClient.AGENT_LOGS_WS, response=BaseModel, method="GET")
        wss_url = self.request_path(endpoint).format(agent_id=agent_id, token=self.token, session_id=session_id)
        wss_url = wss_url.replace("https://", "wss://").replace("http://", "ws://")

        counter = [0]  # mutable container for step count

        # Use JavaScript WebSocket API via Pyodide FFI
        ws = js.WebSocket.new(wss_url)  # pyright: ignore[reportPossiblyUnboundVariable, reportUnknownMemberType, reportUnknownVariableType]
        message_queue: asyncio.Queue[str | None] = asyncio.Queue()

        def on_message(event: Any) -> None:
            message_queue.put_nowait(str(event.data))

        def on_error(_event: Any) -> None:
            logger.error("WebSocket error occurred")
            message_queue.put_nowait(None)  # Signal consumer to stop on error

        def on_close(_event: Any) -> None:
            message_queue.put_nowait(None)

        # Initialize proxies to None for cleanup handling
        on_message_proxy: Any = None
        on_error_proxy: Any = None
        on_close_proxy: Any = None

        # Helper to clean up WebSocket resources (tolerates partially initialized state)
        def cleanup_ws() -> None:
            cleanup_errors: list[str] = []
            for cleanup in (  # pyright: ignore[reportUnknownVariableType]
                lambda: ws.removeEventListener("message", on_message_proxy),  # pyright: ignore[reportUnknownMemberType, reportUnknownLambdaType]
                lambda: ws.removeEventListener("error", on_error_proxy),  # pyright: ignore[reportUnknownMemberType, reportUnknownLambdaType]
                lambda: ws.removeEventListener("close", on_close_proxy),  # pyright: ignore[reportUnknownMemberType, reportUnknownLambdaType]
                lambda: on_message_proxy.destroy() if on_message_proxy is not None else None,
                lambda: on_error_proxy.destroy() if on_error_proxy is not None else None,
                lambda: on_close_proxy.destroy() if on_close_proxy is not None else None,
                ws.close,  # pyright: ignore[reportUnknownMemberType]
            ):
                try:
                    _ = cleanup()  # pyright: ignore[reportUnknownVariableType]
                except Exception as e:
                    cleanup_errors.append(str(e))
            if cleanup_errors:
                logger.debug(f"Failed to clean up WebSocket resources: {'; '.join(cleanup_errors)}")

        # Wait for connection with timeout
        connect_timeout = 30.0
        connect_waited = 0.0

        try:
            # Create proxies and register event listeners inside try block for proper cleanup
            on_message_proxy = create_proxy(on_message)  # pyright: ignore[reportPossiblyUnboundVariable]
            on_error_proxy = create_proxy(on_error)  # pyright: ignore[reportPossiblyUnboundVariable]
            on_close_proxy = create_proxy(on_close)  # pyright: ignore[reportPossiblyUnboundVariable]

            ws.addEventListener("message", on_message_proxy)  # pyright: ignore[reportUnknownMemberType]
            ws.addEventListener("error", on_error_proxy)  # pyright: ignore[reportUnknownMemberType]
            ws.addEventListener("close", on_close_proxy)  # pyright: ignore[reportUnknownMemberType]

            while ws.readyState == 0:  # CONNECTING  # pyright: ignore[reportUnknownMemberType]
                if connect_waited >= connect_timeout:
                    logger.error(f"[Agent] {agent_id} websocket connection timed out after {connect_timeout}s")
                    return None
                await asyncio.sleep(0.1)
                connect_waited += 0.1
            while True:
                message = await message_queue.get()
                if message is None:
                    break

                response, should_stop = self._process_ws_message(message, agent_id, log, counter)

                if should_stop:
                    if isinstance(response, AgentStatusResponse):
                        return response
                    return None

        except ConnectionError as e:
            logger.error(f"Connection error: {agent_id} {e}")
            return None
        except Exception as e:
            logger.error(f"Unexpected websocket processing error: {agent_id} {e} {traceback.format_exc()}")
            raise
        finally:
            cleanup_ws()

        return None

    async def async_watch_logs_and_wait(self, agent_id: str, session_id: str, log: bool = True) -> AgentStatusResponse:
        """
        Execute a task with the agent and wait for completion asynchronously.

        This method is required for Pyodide environments where synchronous WebSocket
        connections are not supported.

        Args:
            agent_id (str): The agent identifier.
            session_id (str): The session identifier.
            log (bool): Whether to log the agent steps.

        Returns:
            AgentStatusResponse: The response from the completed agent execution.
        """
        if not RUNNING_IN_PYODIDE:
            raise NotImplementedError(
                "async_watch_logs_and_wait is only supported in Pyodide. Use watch_logs_and_wait instead."
            )

        try:
            response = await self.async_watch_logs(agent_id=agent_id, session_id=session_id, log=log)
            if response is not None:
                return response
            # If we didn't get a response, poll status until agent is closed
            logger.warning(f"[Agent] {agent_id} did not return status response. Polling status until closed.")
            max_wait_secs = 300
            waited = 0
            while waited < max_wait_secs:
                status = self.status(agent_id=agent_id)
                if status.status == AgentStatus.closed:
                    return status
                await asyncio.sleep(1)
                waited += 1
            raise TimeoutError(f"Agent {agent_id} did not reach a terminal state within {max_wait_secs}s")

        except asyncio.CancelledError:
            # Best-effort cleanup: don't let HTTP failures mask the CancelledError
            try:
                status = self.status(agent_id=agent_id)
                if status.status != AgentStatus.closed:
                    _ = self.stop(agent_id=agent_id, session_id=session_id)
            except Exception:
                pass
            raise

    def stop(self, agent_id: str, session_id: str) -> AgentResponse:
        """
        Stops the specified agent and clears the last agent response.

        Retrieves a valid agent identifier using the provided value or the last stored
        response, sends a stop request to the API, resets the internal agent response,
        and returns the resulting AgentResponse.

        Args:
            agent_id: The identifier of the agent to stop.

        Returns:
            AgentResponse: The response from the stop operation.

        Raises:
            ValueError: If a valid agent identifier cannot be determined.
        """
        logger.info(f"[Agent] {agent_id} is stopping")
        endpoint = AgentsClient._agent_stop_endpoint(agent_id=agent_id, session_id=session_id)
        response = self.request(endpoint)
        logger.info(f"[Agent] {agent_id} stopped")
        return response

    def run(self, **data: Unpack[SdkAgentStartRequestDict]) -> AgentStatusResponse:
        """
        Run an agent with the specified request parameters.
        and wait for completion

        ```python
        with notte.Session() as session:
            agent = notte.Agent(session=session)
            agent.run(task="go to notte.cc and explain what their product is")
        ```

        This function is synchronous and will block the main thread until the agent is completed.

        > Websockets are used to stream the agent logs to the standard output to provide live logs to the user.
        """
        response = self.start(**data)

        if RUNNING_IN_PYODIDE:
            # In Pyodide, use the async path - asyncio.run works with Pyodide's WebLoop
            return asyncio.run(
                self.async_watch_logs_and_wait(
                    agent_id=response.agent_id,
                    session_id=response.session_id,
                )
            )

        return self.watch_logs_and_wait(
            agent_id=response.agent_id,
            session_id=response.session_id,
        )

    def function_code(self, agent_id: str, as_workflow: bool = True) -> AgentFunctionCodeResponse:
        """
        Retrieves a script that reproduces the steps of the specified agent.

        Queries the API using a validated agent ID.
        The provided ID is confirmed (or obtained from the last response if needed), and the
        resulting script is stored internally before being returned.

        Args:
            agent_id: Unique identifier of the agent to check.

        Returns:
            AgentFunctionCodeResponse: The script that reproduces the steps of the specified agent

        Raises:
            ValueError: If no valid agent ID can be determined.
        """
        request = AgentWorkflowCodeRequest(as_workflow=as_workflow)
        endpoint = AgentsClient._agent_function_endpoint(agent_id=agent_id).with_params(request)
        response = self.request(endpoint)
        return response

    def create_function(self, agent_id: str) -> GetFunctionResponse:
        """
        Creates a function that reproduces the steps of the specified agent.

        Queries the API using a validated agent ID.
        The provided ID is confirmed (or obtained from the last response if needed), and the
        resulting script is stored internally before being returned.

        Args:
            agent_id: Unique identifier of the agent to check.

        Returns:
            GetFunctionResponse: The workflow that reproduces the steps of the agent

        Raises:
            ValueError: If no valid agent ID can be determined.
        """
        script = self.function_code(agent_id, as_workflow=True)
        with tempfile.TemporaryDirectory() as tmp_dir:
            filename = Path(tmp_dir) / "code.py"
            with open(filename, "w") as f:
                _ = f.write(script.python_script)

            return self.root_client.functions.create(path=str(filename))

    def status(self, agent_id: str) -> LegacyAgentStatusResponse:
        """
        Retrieves the status of the specified agent.

        Queries the API for the current status of an agent using a validated agent ID.
        The provided ID is confirmed (or obtained from the last response if needed), and the
        resulting status is stored internally before being returned.

        Args:
            agent_id: Unique identifier of the agent to check.

        Returns:
            AgentResponse: The current status information of the specified agent.

        Raises:
            ValueError: If no valid agent ID can be determined.
        """
        request = AgentStatusRequest(agent_id=agent_id)
        endpoint = AgentsClient._agent_status_endpoint(agent_id=agent_id).with_params(request)
        response = self.request(endpoint)
        return response

    def list(self, **data: Unpack[AgentListRequestDict]) -> Sequence[AgentResponse]:
        """
        Lists agents matching specified criteria.

        Validates the keyword arguments using the AgentListRequest model, constructs
        the corresponding endpoint for listing agents, and returns a sequence of agent
        responses.

        Args:
            data: Arbitrary keyword arguments representing filter criteria for agents.

        Returns:
            A sequence of AgentResponse objects.
        """
        params = AgentListRequest.model_validate(data)
        endpoint = AgentsClient._agent_list_endpoint(params=params)
        return self.request_list(endpoint)

    def run_custom(self, request: BaseModel, viewer: bool = False) -> AgentStatusResponse:
        """
        Run a custom agent with the specified request parameters and wait for completion.

        Note: not all servers support custom agents.
        """
        if not self.is_custom_endpoint_available():
            raise ValueError(f"Custom endpoint is not available for this server: {self.server_url}")

        response = self.request(AgentsClient._agent_start_custom_endpoint().with_request(request))

        if viewer:
            self.root_client.sessions.viewer(response.session_id)

        if RUNNING_IN_PYODIDE:
            # In Pyodide, use the async path - asyncio.run works with Pyodide's WebLoop
            return asyncio.run(
                self.async_watch_logs_and_wait(
                    agent_id=response.agent_id,
                    session_id=response.session_id,
                    log=True,
                )
            )

        return self.watch_logs_and_wait(
            agent_id=response.agent_id,
            session_id=response.session_id,
            log=True,
        )


class RemoteAgent:
    """
    A remote agent that can execute tasks through the Notte API.

    This class provides an interface for running tasks, checking status, and managing
    agent executions. It maintains state about the current agent execution and provides
    methods to interact with the agent through an AgentsClient.

    The agent can be started, monitored, and controlled through various methods. It supports
    both synchronous and asynchronous execution modes.

    Key Features:
    - Start and stop agent execution
    - Monitor agent status and progress
    - Wait for task completion with progress updates
    - Support for both sync and async execution

    Attributes:
        request (AgentCreateRequest): The configuration request used to create this agent.
        client (AgentsClient): The client used to communicate with the Notte API.
        response (AgentResponse | None): The latest response from the agent execution.

    Note: This class is designed to work with a single agent instance at a time. For multiple
    concurrent agents, create separate RemoteAgent instances.
    """

    class AgentWorkflow:
        def __init__(
            self,
            client: AgentsClient,
            agent_id: str,
        ):
            self.client: AgentsClient = client
            self.agent_id: str = agent_id

        def code(self, as_workflow: bool = True) -> AgentFunctionCodeResponse:
            """
            Retrieves a script that reproduces the steps of the specified agent.

            Queries the API using a validated agent ID.
            The provided ID is confirmed (or obtained from the last response if needed), and the
            resulting script is stored internally before being returned.

            Args:
                as_workflow: Whether to return a full standalone workflow script or just the relevant steps

            Returns:
                AgentFunctionCodeResponse: The script that reproduces the steps of the specified agent

            Raises:
                ValueError: If no valid agent ID can be determined.
            """
            return self.client.function_code(self.agent_id, as_workflow=as_workflow)

        def create_function(self) -> NotteFunction:
            """
            Creates a function that reproduces the steps of the specified agent.

            Queries the API using a validated agent ID.
            The provided ID is confirmed (or obtained from the last response if needed), and the
            resulting script is stored internally before being returned.

            Returns:
                NotteFunction: The workflow that reproduces the steps of the agent

            Raises:
                ValueError: If no valid agent ID can be determined.
            """

            function_resp = self.client.create_function(self.agent_id)
            return NotteFunction(function_id=function_resp.function_id, _client=self.client.root_client)

    @overload
    def __init__(
        self,
        session: RemoteSession,
        *,
        vault: NotteVault | None = None,
        notifier: BaseNotifier | None = None,
        persona: NottePersona | None = None,
        _client: AgentsClient | None = None,
        agent_id: str | None = None,
        **data: Unpack[AgentCreateRequestDict],
    ) -> None: ...

    @overload
    def __init__(self, *, agent_id: str, _client: AgentsClient | None = None) -> None: ...

    def __init__(
        self,
        session: RemoteSession | None = None,
        vault: NotteVault | None = None,
        notifier: BaseNotifier | None = None,
        persona: NottePersona | None = None,
        _client: AgentsClient | None = None,
        agent_id: str | None = None,
        **data: Unpack[AgentCreateRequestDict],
    ) -> None:
        """
        Create a new RemoteAgent instance with the specified configuration.

        This method validates the agent creation request and sets up the appropriate
        connections with the provided vault and session if specified.

        Args:
            vault: A notte vault instance, if the agent requires authentication
            session: The session to connect to. The session's `open_viewer` parameter controls
                whether to display a live viewer (browsers are always headless).
            notifier: A notifier (for example, email), which will get called upon task completion.
            session_id: (deprecated) use session instead
            **data: Additional keyword arguments for the agent creation request.

        Returns:
            RemoteAgent: A new RemoteAgent instance configured with the specified parameters.
        """
        if _client is None:
            raise ValueError("NotteClient is required")

        if session is None and agent_id is None:
            raise ValueError(
                "Either session (for running a new agent) or agent_id (for accessing an existing agent) have to be provided"
            )

        if session is not None and agent_id is not None:
            raise ValueError(
                "Either session (for running a new agent) or agent_id (for accessing an existing agent) have to be provided, not both"
            )

        existing_agent: bool = agent_id is not None
        self.existing_agent: bool = existing_agent
        self.client: AgentsClient = _client

        if existing_agent:
            self.response = _client.status(agent_id=agent_id)
            return

        if session is None:
            raise ValueError("Session is required for running a new agent")

        data["session_id"] = session.session_id  # pyright: ignore[reportGeneralTypeIssues]
        request = SdkAgentCreateRequest.model_validate(data)
        if notifier is not None:
            notifier_config = notifier.model_dump()
            request.notifier_config = notifier_config

        # #########################################################
        # ###################### Vault checks #####################
        # #########################################################

        if vault is not None:
            if len(vault.vault_id) == 0:
                raise ValueError("Vault ID cannot be empty")
            request.vault_id = vault.vault_id

        if persona is not None:
            if len(persona.persona_id) == 0:
                raise ValueError("Persona ID cannot be empty")
            request.persona_id = persona.persona_id

        # #########################################################
        # #################### Session checks #####################
        # #########################################################

        if not isinstance(session, RemoteSession):  # pyright: ignore[reportUnnecessaryIsInstance]
            raise ValueError(
                "You are trying to use a local session with a remote agent. This is not supported. Use `notte.Agent(session=session)` instead."
            )  # pyright: ignore[reportUnreachable]
        if len(session.session_id) == 0:
            raise ValueError("Session ID cannot be empty")
        request.session_id = session.session_id

        self.request: SdkAgentCreateRequest = request
        self.response: AgentResponse | None = None

    @property
    def agent_id(self) -> str:
        """
        Get the ID of the current agent execution.

        Returns:
            str: The unique identifier of the current agent execution.

        Raises:
            ValueError: If the agent hasn't been run yet (no response available).
        """
        if self.response is None:
            raise ValueError("You need to run the agent first to get the agent id")
        return self.response.agent_id

    @property
    def session_id(self) -> str:
        """
        Get the ID of the current session.
        """
        if self.response is None:
            raise ValueError("You need to run the agent first to get the session id")
        return self.response.session_id

    @track_usage("cloud.agent.start")
    def start(self, **data: Unpack[AgentRunRequestDict]) -> AgentResponse:
        """
        Start the agent with the specified request parameters.

        This method initiates the agent execution with the provided configuration.
        The agent will begin processing the task immediately after starting.

        Args:
            **data: Keyword arguments representing the fields of an AgentStartRequest.

        Returns:
            AgentResponse: The initial response from starting the agent.
        """
        if self.existing_agent:
            raise ValueError("You cannot call run() on an agent instantiated from agent id")

        self.response = self.client.start(**self.request.model_dump(), **data)
        return self.response

    def wait(self) -> AgentStatusResponse:
        """
        Wait for the agent to complete its current task.

        This method polls the agent's status at regular intervals until completion.
        During waiting, it displays progress updates and a spinner (in non-notebook environments).
        The polling continues until either the agent completes or a timeout is reached.

        Returns:
            AgentStatusResponse: The final status response after completion.

        Raises:
            TimeoutError: If the agent doesn't complete within the maximum allowed attempts.
        """
        if self.existing_agent:
            raise ValueError("You cannot call wait() on an agent instantiated from agent id")

        return self.client.wait(agent_id=self.agent_id)

    def watch_logs(self, log: bool = False) -> AgentStatusResponse | None:
        """
        Watch the logs of the agent.
        """
        if self.existing_agent:
            raise ValueError("You cannot call watch_logs() on an agent instantiated from agent id")

        return self.client.watch_logs(agent_id=self.agent_id, session_id=self.session_id, log=log)

    def watch_logs_and_wait(self, log: bool = True) -> AgentStatusResponse:
        """
        Watch the logs of the agent and wait for completion.
        """
        if self.existing_agent:
            raise ValueError("You cannot call watch_logs_and_wait() on an agent instantiated from agent id")

        return self.client.watch_logs_and_wait(agent_id=self.agent_id, session_id=self.session_id, log=log)

    async def async_watch_logs_and_wait(self, log: bool = True) -> AgentStatusResponse:
        """
        Watch the logs of the agent and wait for completion asynchronously.

        In Pyodide (WebAssembly), this delegates to the client's async method directly
        since asyncio.to_thread is not supported in single-threaded environments.
        In native Python, this runs the synchronous watch_logs_and_wait in a thread pool
        to avoid blocking the event loop.

        Note: When cancelled (e.g., via asyncio.timeout), this method stops the agent
        gracefully. However, in native Python the underlying thread cannot be interrupted
        immediately - it will continue until the server processes the stop signal and
        closes the WebSocket connection. This is not a leak, but cancellation may not
        be instantaneous under high load.
        """
        if self.existing_agent:
            raise ValueError("You cannot call async_watch_logs_and_wait() on an agent instantiated from agent id")

        if RUNNING_IN_PYODIDE:
            return await self.client.async_watch_logs_and_wait(
                agent_id=self.agent_id,
                session_id=self.session_id,
                log=log,
            )

        try:
            return await asyncio.to_thread(
                self.client.watch_logs_and_wait,
                agent_id=self.agent_id,
                session_id=self.session_id,
                log=log,
            )
        except asyncio.CancelledError:
            # Gracefully stop the agent on cancellation (mirrors KeyboardInterrupt handling in sync version)
            # Best-effort cleanup: don't let HTTP failures mask the CancelledError
            try:
                status = self.client.status(agent_id=self.agent_id)
                if status.status != AgentStatus.closed:
                    _ = self.client.stop(agent_id=self.agent_id, session_id=self.session_id)
            except Exception:
                pass
            raise

    @track_usage("cloud.agent.stop")
    def stop(self) -> AgentResponse:
        """
        Stop the currently running agent.

        This method sends a stop request to the agent, terminating its current execution.
        The agent will stop processing its current task immediately.

        Returns:
            AgentResponse: The response from the stop operation.

        Raises:
            ValueError: If the agent hasn't been run yet (no agent_id available).
        """
        if self.existing_agent:
            raise ValueError("You cannot call stop() on an agent instantiated from agent id")

        return self.client.stop(agent_id=self.agent_id, session_id=self.session_id)

    @track_usage("cloud.agent.run")
    def run(self, **data: Unpack[AgentRunRequestDict]) -> AgentStatusResponse:
        """
        Run an agent with the specified request parameters and wait for completion.

        ```python
        with notte.Session() as session:
            agent = notte.Agent(session=session)
            agent.run(task="go to notte.cc and explain what their product is")
        ```

        This function is synchronous and will block the main thread until the agent is completed.

        > Websockets are used to stream the agent logs to the standard output to provide live logs to the user.

        Args:
            **data: Keyword arguments representing the fields of an AgentRunRequest.

        Returns:
            AgentStatusResponse: The final status response after task completion.

        Raises:
            TimeoutError: If the agent doesn't complete within the maximum allowed attempts.
        """
        if self.existing_agent:
            raise ValueError("You cannot call run() on an agent instantiated from agent id")

        self.response = self.start(**data)
        logger.info(f"[Agent] {self.agent_id} started with model: {self.request.reasoning_model}")

        if RUNNING_IN_PYODIDE:
            # In Pyodide, use the async path - asyncio.run works with Pyodide's WebLoop
            status_response = asyncio.run(self.async_watch_logs_and_wait())
        else:
            status_response = self.watch_logs_and_wait()
        prefix = "✅ Agent returned with success:" if status_response.success else "❌ Agent returned with failure:"
        logger.info(f"{prefix} {status_response.answer}")
        return status_response

    @track_usage("cloud.agent.status")
    def status(self) -> LegacyAgentStatusResponse:
        """
        Get the current status of the agent.

        This method retrieves the current state of the agent, including its progress,
        actions taken, and any errors or messages.

        ```python
        status = agent.status()
        ```


        Returns:
            LegacyAgentStatusResponse: The current status of the agent execution.

        Raises:
            ValueError: If the agent hasn't been run yet (no agent_id available).
        """
        return self.client.status(agent_id=self.agent_id)

    def replay(
        self,
        wait: bool = True,
        timeout: float = 240.0,
        poll_interval: float = 5.0,
    ) -> ReplayResponse:
        """
        Get the replay for the agent's session.

        .. deprecated::
            Use ``session.replay()`` instead. Agent replay is deprecated
            in favor of session-level replay with presigned URLs.

        Args:
            wait: If True (default), poll until the replay is ready.
            timeout: Maximum seconds to wait (default 120).
            poll_interval: Seconds between polling attempts (default 2).

        Returns:
            ReplayResponse: Presigned URLs for HLS playlist and MP4 download.
        """
        warnings.warn(
            "agent.replay() is deprecated. Use session.replay() instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.client.root_client.sessions.replay(
            session_id=self.session_id,
            wait=wait,
            timeout=timeout,
            poll_interval=poll_interval,
        )

    @property
    @track_usage("cloud.agent.workflow")
    def workflow(self) -> AgentWorkflow:
        """
        Get the workflow from the completed steps of the agent.

        ```python
        agent.run(task="...")
        workflow = agent.workflow
        ```

        Returns:
            AgentWorkflow: The agent workflow that replicates the agent steps

        Raises:
            ValueError: If the agent hasn't been run yet (no agent_id available).
        """
        return RemoteAgent.AgentWorkflow(self.client, self.agent_id)
