import grpc
from typing import List, Optional
from urllib.parse import urlparse

from .proto import sandbox_pb2
from .proto import sandbox_pb2_grpc


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: str) -> None:
        """
        Stop a running sandbox.

        Args:
            sandbox_id: The ID of the sandbox to stop.
        """
        request = sandbox_pb2.StopRequest(id=sandbox_id)
        self.stub.StopSandbox(request)

    def exec_command(self, sandbox_id: str, command: List[str]) -> sandbox_pb2.ExecResponse:
        """
        Execute a command in a running sandbox.

        Args:
            sandbox_id: The ID of the sandbox.
            command: The command to execute.

        Returns:
            The execution response containing stdout, stderr, and exit code.
        """
        request = sandbox_pb2.ExecRequest(id=sandbox_id, command=command)
        return self.stub.ExecCommand(request)

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

        Args:
            sandbox_id: The ID of the sandbox.

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

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

        Args:
            sandbox_id: The ID of the sandbox to snapshot.

        Returns:
            A dictionary containing the paths to the snapshot file and memory file.
        """
        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: str) -> None:
        """
        Resume a paused or snapshotted sandbox.

        Args:
            sandbox_id: The ID of the sandbox to resume.
        """
        request = sandbox_pb2.ResumeRequest(id=sandbox_id)
        self.stub.ResumeSandbox(request)

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

        Args:
            sandbox_id: The ID of the sandbox.
            code: The Python code snippet to execute.

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

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

        Args:
            sandbox_id: The ID of the sandbox.

        Returns:
            A JSON-formatted string containing the platform info dict.
        """
        request = sandbox_pb2.GetPlatformInfoRequest(id=sandbox_id)
        response = self.stub.GetPlatformInfo(request)
        return response.info


