# Copyright 2025 Daytona Platforms Inc.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import json
import re

import httpx
import httpx_ws
from wsproto.events import BytesMessage, CloseConnection, TextMessage

from daytona_toolbox_api_client import CreateContextRequest, InterpreterApi, InterpreterContext

from .._utils.errors import intercept_errors
from ..common.code_interpreter import ExecutionError, ExecutionResult, OutputMessage
from ..common.errors import DaytonaConnectionError, DaytonaTimeoutError
from ..common.process import OutputHandler

WEBSOCKET_TIMEOUT_CODE = 4008


class CodeInterpreter:
    """Handles code interpretation and execution within a Sandbox. Currently supports only Python.

    This class provides methods to execute code in isolated interpreter contexts,
    manage contexts, and stream execution output via callbacks. If subsequent code executions
    are performed in the same context, the variables, imports, and functions defined in
    the previous execution will be available.

    For other languages, use the `code_run` method from the `Process` interface,
    or execute the appropriate command directly in the sandbox terminal.
    """

    def __init__(
        self,
        api_client: InterpreterApi,
        http_client: httpx.Client,
    ):
        """Initialize a new CodeInterpreter instance.

        Args:
            api_client: API client for interpreter operations.
            http_client: Shared httpx.Client whose connection pool the WS upgrade reuses.
        """
        self._api_client: InterpreterApi = api_client
        self._http_client: httpx.Client = http_client

    @intercept_errors(message_prefix="Failed to run code: ")
    def run_code(
        self,
        code: str,
        *,
        context: InterpreterContext | None = None,
        on_stdout: OutputHandler[OutputMessage] | None = None,
        on_stderr: OutputHandler[OutputMessage] | None = None,
        on_error: OutputHandler[ExecutionError] | None = None,
        envs: dict[str, str] | None = None,
        timeout: int | None = None,
    ) -> ExecutionResult:
        """Execute Python code in the sandbox.

        By default, code runs in the default shared context which persists variables,
        imports, and functions across executions. To run in an isolated context,
        create a new context with `create_context()` and pass it as `context` argument.

        Args:
            code (str): Code to execute.
            context (InterpreterContext | None): Context to run code in. If not provided, uses default context.
            on_stdout (OutputHandler[OutputMessage] | None): Callback for stdout messages.
            on_stderr (OutputHandler[OutputMessage] | None): Callback for stderr messages.
            on_error (OutputHandler[ExecutionError] | None): Callback for execution errors
                (e.g., syntax errors, runtime errors).
            envs (dict[str, str] | None): Environment variables for this execution.
            timeout (int | None): Timeout in seconds. 0 means no timeout. Default is 10 minutes.

        Returns:
            ExecutionResult: Result object containing stdout, stderr and error if any.

        Raises:
            DaytonaTimeoutError: If execution times out.
            DaytonaError: If execution fails due to communication or other SDK errors.

        Examples:
            ```python
            def handle_stdout(msg: OutputMessage):
                print(f"STDOUT: {msg.output}", end="")

            def handle_stderr(msg: OutputMessage):
                print(f"STDERR: {msg.output}", end="")

            def handle_error(err: ExecutionError):
                print(f"ERROR: {err.name}: {err.value}")

            code = '''
            import sys
            import time
            for i in range(5):
                print(i)
                time.sleep(1)
            sys.stderr.write("Counting done!")
            '''
            result = sandbox.code_interpreter.run_code(
                code=code,
                on_stdout=handle_stdout,
                on_stderr=handle_stderr,
                on_error=handle_error,
                timeout=10
            )
            ```
        """
        _, url, headers, *_ = self._api_client._execute_interpreter_code_serialize(
            _request_auth=None,
            _content_type=None,
            _headers=None,
            _host_index=None,
        )
        url = re.sub(r"^http", "ws", url)

        result = ExecutionResult()

        with httpx_ws.connect_ws(url, self._http_client, headers=headers) as ws:
            request: dict[str, str | int | dict[str, str]] = {"code": code}
            if context is not None:
                request["contextId"] = context.id
            if envs is not None:
                request["envs"] = envs
            if timeout is not None:
                request["timeout"] = timeout

            ws.send_text(json.dumps(request))

            while True:
                try:
                    event = ws.receive()
                except httpx_ws.WebSocketDisconnect as e:
                    self._maybe_raise_from_close(e.code, e.reason)
                    break

                if isinstance(event, TextMessage):
                    chunk = json.loads(event.data)
                elif isinstance(event, BytesMessage):
                    chunk = json.loads(event.data.decode("utf-8"))
                elif isinstance(event, CloseConnection):
                    self._maybe_raise_from_close(event.code, event.reason)
                    break
                else:
                    # Ping/Pong/etc. — wsproto auto-replies to Ping; just continue.
                    continue

                chunk_type = chunk.get("type")
                if chunk_type == "stdout":
                    stdout = chunk.get("text", "")
                    if on_stdout:
                        _ = on_stdout(OutputMessage(output=stdout))
                    result.stdout += stdout
                elif chunk_type == "stderr":
                    stderr = chunk.get("text", "")
                    if on_stderr:
                        _ = on_stderr(OutputMessage(output=stderr))
                    result.stderr += stderr
                elif chunk_type == "error":
                    error = ExecutionError(
                        name=chunk.get("name", ""),
                        value=chunk.get("value", ""),
                        traceback=chunk.get("traceback", ""),
                    )
                    result.error = error
                    if on_error:
                        _ = on_error(error)

        return result

    @intercept_errors(message_prefix="Failed to create interpreter context: ")
    def create_context(
        self,
        cwd: str | None = None,
    ) -> InterpreterContext:
        """Create a new isolated interpreter context.

        Contexts provide isolated execution environments with their own global namespace.
        Variables, imports, and functions defined in one context don't affect others.

        Args:
            cwd (str | None): Working directory for the context. If not specified, uses sandbox working directory.

        Returns:
            InterpreterContext: The created context with its ID and metadata.

        Raises:
            DaytonaError: If context creation fails.

        Examples:
            ```python
            # Create isolated context
            ctx = sandbox.code_interpreter.create_context()

            # Execute code in this context
            sandbox.code_interpreter.run_code("x = 100", context=ctx)

            # Variable only exists in this context
            result = sandbox.code_interpreter.run_code("print(x)", context=ctx)  # OK

            # Won't see the variable in default context
            result = sandbox.code_interpreter.run_code("print(x)")  # NameError

            # Clean up
            sandbox.code_interpreter.delete_context(ctx)
            ```
        """
        return self._api_client.create_interpreter_context(request=CreateContextRequest(cwd=cwd))

    @intercept_errors(message_prefix="Failed to list interpreter contexts: ")
    def list_contexts(self) -> list[InterpreterContext]:
        """List all user-created interpreter contexts.

        The default context is not included in this list. Only contexts created
        via `create_context()` are returned.

        Returns:
            list[InterpreterContext]: List of context objects.

        Raises:
            DaytonaError: If listing fails.

        Examples:
            ```python
            contexts = sandbox.code_interpreter.list_contexts()
            for ctx in contexts:
                print(f"Context {ctx.id}: {ctx.language} at {ctx.cwd}")
            ```
        """
        return (self._api_client.list_interpreter_contexts()).contexts or []

    @intercept_errors(message_prefix="Failed to delete interpreter context: ")
    def delete_context(self, context: InterpreterContext) -> None:
        """Delete an interpreter context and shut down all associated processes.

        This permanently removes the context and all its state (variables, imports, etc.).
        The default context cannot be deleted.

        Args:
            context (InterpreterContext): Context to delete.

        Raises:
            DaytonaError: If deletion fails or context not found.

        Examples:
            ```python
            ctx = sandbox.code_interpreter.create_context()
            # ... use context ...
            sandbox.code_interpreter.delete_context(ctx)
            ```
        """
        _ = self._api_client.delete_interpreter_context(id=context.id)

    def _maybe_raise_from_close(self, close_code: int | None, close_reason: str | None) -> None:
        """Translate a websocket close into a Daytona error when it indicates failure.

        Mirrors the async interpreter: code ``4008`` from the daemon means timeout,
        non-1000 codes are connection errors, code 1000 / unknown means clean completion.
        """
        if close_code == WEBSOCKET_TIMEOUT_CODE:
            raise DaytonaTimeoutError(
                "Execution timed out: operation exceeded the configured `timeout`. Provide a larger value if needed."
            )
        if close_code in (None, 1000):
            return
        detail = close_reason or "WebSocket connection closed unexpectedly"
        detail = f"{detail} (close code {close_code})"
        raise DaytonaConnectionError(detail)
