import json
import time

import grpc
from typing import Any, Dict, List, Optional, Union
from urllib.parse import urlparse

from .proto import sandbox_pb2
from .proto import sandbox_pb2_grpc
from .http_client import _resolve_id, _prepare_command, Command, ExecResult, SandboxID


class _OCFProxyInterceptor(
    grpc.UnaryUnaryClientInterceptor,
    grpc.UnaryStreamClientInterceptor,
    grpc.StreamUnaryClientInterceptor,
    grpc.StreamStreamClientInterceptor,
):
    """Interceptor that prepends an OCF proxy path prefix to gRPC method paths.

    Python's grpc library does NOT include the URL path portion of the channel
    target in the HTTP/2 :path pseudo-header.  It always sends
    :path = /package.Service/Method.  The OCF proxy needs the full path
    (e.g. /v1/service/sandbox/package.Service/Method) to route the request.
    This interceptor rewrites every outgoing RPC's method string to include
    the prefix.
    """

    def __init__(self, path_prefix: str):
        # Ensure the prefix starts with / and does not end with /
        self._prefix = "/" + path_prefix.strip("/")

    def _rewrite(self, client_call_details):
        """Return a new ClientCallDetails with the rewritten method path."""
        new_method = self._prefix + client_call_details.method
        return _new_client_call_details(client_call_details, new_method)

    def intercept_unary_unary(self, continuation, client_call_details, request):
        return continuation(self._rewrite(client_call_details), request)

    def intercept_unary_stream(self, continuation, client_call_details, request):
        return continuation(self._rewrite(client_call_details), request)

    def intercept_stream_unary(self, continuation, client_call_details, request_iterator):
        return continuation(self._rewrite(client_call_details), request_iterator)

    def intercept_stream_stream(self, continuation, client_call_details, request_iterator):
        return continuation(self._rewrite(client_call_details), request_iterator)


class _ClientCallDetails(
    grpc.ClientCallDetails,
):
    """Concrete implementation of ClientCallDetails for rewriting."""

    def __init__(self, method, timeout, metadata, credentials, wait_for_ready, compression):
        self.method = method
        self.timeout = timeout
        self.metadata = metadata
        self.credentials = credentials
        self.wait_for_ready = wait_for_ready
        self.compression = compression


def _new_client_call_details(original, new_method):
    return _ClientCallDetails(
        method=new_method,
        timeout=original.timeout,
        metadata=original.metadata,
        credentials=original.credentials,
        wait_for_ready=original.wait_for_ready,
        compression=original.compression,
    )


class SandboxClient:
    """Client for interacting with the Agent Sandbox Orchestrator via gRPC."""

    def __init__(self, host: Optional[str] = None, port: int = 50051, address: Optional[str] = None):
        """
        Initialize the SandboxClient.

        Args:
            host: The hostname or IP address of the orchestrator.
            port: The port number of the orchestrator gRPC service.
            address: Full address for proxy-based routing (e.g.
                     ``localhost:8092/v1/service/sandbox``).  The path portion
                     is extracted and prepended to each gRPC method via an
                     interceptor so the OCF proxy can route the request.
                     Overrides *host* and *port*.
        """
        if address:
            # Strip scheme if the user passed one — gRPC channels don't use schemes.
            addr = address
            for scheme in ("http://", "https://"):
                if addr.startswith(scheme):
                    addr = addr[len(scheme):]
                    break
            # Split "host:port/path..." into the channel target and the proxy path.
            # urlparse needs a scheme, so we add one temporarily.
            parsed = urlparse(f"http://{addr}")
            channel_target = f"{parsed.hostname}:{parsed.port or 80}"
            proxy_path = parsed.path  # e.g. "/v1/service/sandbox"
        elif host:
            channel_target = f"{host}:{port}"
            proxy_path = ""
        else:
            raise ValueError("Must provide either an address or a host")

        self.address = channel_target
        self.channel = grpc.insecure_channel(channel_target)

        if proxy_path and proxy_path != "/":
            interceptor = _OCFProxyInterceptor(proxy_path)
            self.channel = grpc.intercept_channel(self.channel, interceptor)

        self.stub = sandbox_pb2_grpc.SandboxServiceStub(self.channel)

    def close(self):
        """Close the gRPC channel."""
        self.channel.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def start_sandbox(
        self,
        type: str,
        image: str,
        command: Optional[List[str]] = None,
        memory_mb: int = 512,
        cpu_cores: float = 1.0,
    ) -> str:
        """
        Start a new sandbox.

        Args:
            type: The type of sandbox (e.g., "docker", "firecracker").
            image: The image to use for the sandbox.
            command: The command to run in the sandbox.
            memory_mb: Memory limit in MB.
            cpu_cores: CPU cores limit.

        Returns:
            The ID of the started sandbox.
        """
        request = sandbox_pb2.StartRequest(
            type=type,
            image=image,
            command=command or [],
            memory_mb=memory_mb,
            cpu_cores=cpu_cores,
        )
        response = self.stub.StartSandbox(request)
        return response.id

    def stop_sandbox(self, sandbox_id: SandboxID, cleanup: bool = True) -> None:
        """
        Stop a running sandbox.

        Args:
            sandbox_id: The ID of the sandbox to stop, or a Sandbox object.
            cleanup: When True (default), backing resources (container, VM,
                temporary files) are removed.  When False the sandbox
                process is stopped but resources are left in place for
                inspection or later resumption.
        """
        sandbox_id = _resolve_id(sandbox_id)
        request = sandbox_pb2.StopRequest(id=sandbox_id, skip_cleanup=not cleanup)
        self.stub.StopSandbox(request)

    def exec_command(self, sandbox_id: SandboxID, command: Command) -> ExecResult:
        """
        Execute a command in a running sandbox.

        Args:
            sandbox_id: The ID of the sandbox, or a Sandbox object.
            command: The command to execute.  Accepts a plain string
                (``"ls -al"``), a single-element list containing a shell
                expression (``["ls -al | head"]``), or a traditional argv
                list (``["ls", "-al"]``).  Strings and single-element
                lists that contain whitespace or shell meta-characters are
                automatically wrapped in ``["sh", "-c", ...]``.

        Returns:
            The execution response containing stdout, stderr, and exit code.
        """
        sandbox_id = _resolve_id(sandbox_id)
        argv = _prepare_command(command)
        request = sandbox_pb2.ExecRequest(id=sandbox_id, command=argv)
        t0 = time.monotonic()
        resp = self.stub.ExecCommand(request)
        elapsed = (time.monotonic() - t0) * 1000.0
        return ExecResult(
            stdout=resp.stdout,
            stderr=resp.stderr,
            exit_code=resp.exit_code,
            elapsed_ms=elapsed,
        )

    def get_status(self, sandbox_id: SandboxID) -> str:
        """
        Get the status of a sandbox.

        Args:
            sandbox_id: The ID of the sandbox, or a Sandbox object.

        Returns:
            The status of the sandbox (e.g., "running", "stopped").
        """
        sandbox_id = _resolve_id(sandbox_id)
        request = sandbox_pb2.StatusRequest(id=sandbox_id)
        response = self.stub.GetStatus(request)
        return response.status

    def snapshot_sandbox(self, sandbox_id: SandboxID) -> dict:
        """
        Snapshot a running sandbox.

        Args:
            sandbox_id: The ID of the sandbox to snapshot, or a Sandbox object.

        Returns:
            A dictionary containing the paths to the snapshot file and memory file.
        """
        sandbox_id = _resolve_id(sandbox_id)
        request = sandbox_pb2.SnapshotRequest(id=sandbox_id)
        response = self.stub.SnapshotSandbox(request)
        return {
            "snapshot_path": response.snapshot_path,
            "mem_file_path": response.mem_file_path,
        }

    def resume_sandbox(self, sandbox_id: SandboxID) -> None:
        """
        Resume a paused or snapshotted sandbox.

        Args:
            sandbox_id: The ID of the sandbox to resume, or a Sandbox object.
        """
        sandbox_id = _resolve_id(sandbox_id)
        request = sandbox_pb2.ResumeRequest(id=sandbox_id)
        self.stub.ResumeSandbox(request)

    def run_python(self, sandbox_id: SandboxID, code: str) -> str:
        """
        Execute python code in a running sandbox.

        Args:
            sandbox_id: The ID of the sandbox, or a Sandbox object.
            code: The Python code snippet to execute.

        Returns:
            The standard output resulting from executing the Python code.
        """
        sandbox_id = _resolve_id(sandbox_id)
        request = sandbox_pb2.RunPythonRequest(id=sandbox_id, code=code)
        response = self.stub.RunPython(request)
        return response.result

    def get_platform_info(self, sandbox_id: SandboxID) -> Dict[str, Any]:
        """
        Get platform information from a running sandbox using Python.

        Args:
            sandbox_id: The ID of the sandbox, or a Sandbox object.

        Returns:
            A dictionary containing platform info (e.g. keys like
            ``"system"``, ``"node"``, ``"release"``, ``"version"``,
            ``"machine"``, ``"processor"``).
        """
        sandbox_id = _resolve_id(sandbox_id)
        request = sandbox_pb2.GetPlatformInfoRequest(id=sandbox_id)
        response = self.stub.GetPlatformInfo(request)
        info = response.info
        if isinstance(info, dict):
            return info
        return json.loads(info)

    def read_file(self, sandbox_id: SandboxID, path: str) -> bytes:
        """
        Read a file from a running sandbox.

        Args:
            sandbox_id: The ID of the sandbox, or a Sandbox object.
            path: The path of the file to read inside the sandbox.

        Returns:
            The contents of the file.
        """
        sandbox_id = _resolve_id(sandbox_id)
        request = sandbox_pb2.ReadFileRequest(id=sandbox_id, path=path)
        response = self.stub.ReadFile(request)
        return response.content

    def write_file(self, sandbox_id: SandboxID, path: str, content: bytes) -> None:
        """
        Write a file to a running sandbox.

        Args:
            sandbox_id: The ID of the sandbox, or a Sandbox object.
            path: The path of the file to write inside the sandbox.
            content: The file content as bytes.
        """
        sandbox_id = _resolve_id(sandbox_id)
        request = sandbox_pb2.WriteFileRequest(id=sandbox_id, path=path, content=content)
        self.stub.WriteFile(request)
