"""HTTP-based client for the Sandbox Orchestrator.

This module provides :class:`HTTPClient`, a drop-in alternative to the
gRPC-based :class:`SandboxClient` that communicates with the orchestrator
over its REST/JSON HTTP API (default port 8080).

Every public method on :class:`HTTPClient` mirrors the corresponding method
on :class:`SandboxClient` so that callers can switch transports without
changing application logic.
"""

from __future__ import annotations

import json as _json
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union
from urllib.parse import urljoin, urlparse

import requests


# ---------------------------------------------------------------------------
# Response data-classes
# ---------------------------------------------------------------------------
# These give attribute-style access identical to the protobuf message objects
# returned by the gRPC client, so callers can treat the two transports
# interchangeably.


@dataclass(frozen=True)
class ExecResult:
    """Result of executing a command inside a sandbox.

    Attributes mirror ``sandbox_pb2.ExecResponse``.
    """

    stdout: str = ""
    stderr: str = ""
    exit_code: int = 0


@dataclass(frozen=True)
class MetricsResult:
    """Point-in-time resource-usage metrics for a sandbox.

    Attributes mirror ``sandbox_pb2.MetricsResponse``.
    """

    memory_usage_bytes: int = 0
    memory_limit_bytes: int = 0
    memory_percent: float = 0.0
    cpu_percent: float = 0.0
    pids_current: int = 0
    net_rx_bytes: int = 0
    net_tx_bytes: int = 0
    block_read_bytes: int = 0
    block_write_bytes: int = 0


@dataclass(frozen=True)
class SnapshotResult:
    """Paths produced by a snapshot operation.

    Attributes mirror ``sandbox_pb2.SnapshotResponse``.
    """

    snapshot_path: str = ""
    mem_file_path: str = ""


# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------


class SandboxHTTPError(Exception):
    """Raised when the orchestrator returns a non-success HTTP status."""

    def __init__(self, status_code: int, detail: str) -> None:
        self.status_code = status_code
        self.detail = detail
        super().__init__(f"HTTP {status_code}: {detail}")


class SandboxNotFoundError(SandboxHTTPError):
    """Raised when the requested sandbox does not exist (HTTP 404)."""


# ---------------------------------------------------------------------------
# Client
# ---------------------------------------------------------------------------


class HTTPClient:
    """Client for the Sandbox Orchestrator REST/JSON API.

    The public interface intentionally matches :class:`SandboxClient` so that
    users can swap transports with minimal code changes.

    Parameters
    ----------
    host:
        Hostname or IP of the orchestrator.  Ignored when *address* is given.
    port:
        HTTP port of the orchestrator (default ``8080``).
    address:
        Full base URL **or** ``host:port/path`` string.  When a path component
        is present it is used as a prefix for all requests (useful for reverse-
        proxy setups such as ``localhost:8092/v1/service/sandbox``).  Overrides
        *host* and *port*.
    timeout:
        Default request timeout in seconds.  ``None`` means no timeout.
    session:
        Optional :class:`requests.Session` to use.  When not provided a new
        session is created internally.  Passing your own session is handy for
        custom TLS settings, retries, or authentication.

    Usage::

        from flash_sandbox import HTTPClient

        client = HTTPClient(host="localhost", port=8080)

        sid = client.start_sandbox(type="docker", image="alpine:latest",
                                   command=["tail", "-f", "/dev/null"])
        status = client.get_status(sid)
        result = client.exec_command(sid, ["echo", "hello"])
        print(result.stdout)
        client.stop_sandbox(sid)
        client.close()
    """

    # --------------------------------------------------------------------- #
    # Construction / teardown
    # --------------------------------------------------------------------- #

    def __init__(
        self,
        host: Optional[str] = None,
        port: int = 8080,
        address: Optional[str] = None,
        timeout: Optional[float] = 30.0,
        session: Optional[requests.Session] = None,
    ) -> None:
        if address is not None:
            self._base_url = self._normalize_address(address)
        elif host is not None:
            self._base_url = f"http://{host}:{port}"
        else:
            raise ValueError("Must provide either 'address' or 'host'")

        # Guarantee a trailing slash so that urljoin works correctly.
        if not self._base_url.endswith("/"):
            self._base_url += "/"

        self.timeout = timeout
        self._session = session or requests.Session()
        self._owns_session = session is None  # only close if we created it

    # -- helpers for address normalisation -------------------------------- #

    @staticmethod
    def _normalize_address(address: str) -> str:
        """Turn a user-supplied *address* into a proper base URL."""
        addr = address.strip()
        # Add a scheme when missing so urlparse works correctly.
        if not addr.startswith("http://") and not addr.startswith("https://"):
            addr = f"http://{addr}"
        parsed = urlparse(addr)
        scheme = parsed.scheme or "http"
        host = parsed.hostname or "localhost"
        port = parsed.port or (443 if scheme == "https" else 80)
        path = parsed.path.rstrip("/")
        return f"{scheme}://{host}:{port}{path}"

    # -- context-manager -------------------------------------------------- #

    def close(self) -> None:
        """Close the underlying HTTP session (if we own it)."""
        if self._owns_session:
            self._session.close()

    def __enter__(self) -> "HTTPClient":
        return self

    def __exit__(
        self,
        exc_type: Any,
        exc_val: Any,
        exc_tb: Any,
    ) -> None:
        self.close()

    # --------------------------------------------------------------------- #
    # Internal request helpers
    # --------------------------------------------------------------------- #

    def _url(self, path: str) -> str:
        """Build an absolute URL for the given *path* (no leading slash)."""
        return urljoin(self._base_url, path)

    def _raise_for_status(self, resp: requests.Response) -> None:
        """Raise a :class:`SandboxHTTPError` when the response is not 2xx."""
        if resp.ok:
            return
        detail = resp.text.strip()
        if resp.status_code == 404:
            raise SandboxNotFoundError(resp.status_code, detail)
        raise SandboxHTTPError(resp.status_code, detail)

    def _get(self, path: str, **kwargs: Any) -> requests.Response:
        resp = self._session.get(self._url(path), timeout=self.timeout, **kwargs)
        self._raise_for_status(resp)
        return resp

    def _post(self, path: str, json: Any = None, **kwargs: Any) -> requests.Response:
        resp = self._session.post(self._url(path), json=json, timeout=self.timeout, **kwargs)
        self._raise_for_status(resp)
        return resp

    def _delete(self, path: str, **kwargs: Any) -> requests.Response:
        resp = self._session.delete(self._url(path), timeout=self.timeout, **kwargs)
        self._raise_for_status(resp)
        return resp

    # --------------------------------------------------------------------- #
    # Public API – mirrors SandboxClient (gRPC)
    # --------------------------------------------------------------------- #

    def start_sandbox(
        self,
        type: str,
        image: str,
        command: Optional[List[str]] = None,
        memory_mb: int = 512,
        cpu_cores: float = 1.0,
        kernel_image: Optional[str] = None,
        initrd_path: Optional[str] = None,
        snapshot_path: Optional[str] = None,
        mem_file_path: Optional[str] = None,
    ) -> str:
        """Start a new sandbox and return its ID.

        Parameters
        ----------
        type:
            Sandbox backend – ``"docker"``, ``"gvisor"``, or
            ``"firecracker"``.
        image:
            Container / rootfs image name.
        command:
            Entrypoint command to run inside the sandbox.
        memory_mb:
            Memory limit in megabytes.
        cpu_cores:
            CPU cores limit.
        kernel_image:
            *(Firecracker only)* Path to the kernel image.
        initrd_path:
            *(Firecracker only)* Path to the initrd.
        snapshot_path:
            When restoring from a snapshot, path to the snapshot file.
        mem_file_path:
            When restoring from a snapshot, path to the memory file.

        Returns
        -------
        str
            The unique ID of the newly created sandbox.
        """
        payload: Dict[str, Any] = {
            "type": type,
            "image": image,
            "command": command or [],
            "memory_mb": memory_mb,
            "cpu_cores": cpu_cores,
        }
        if kernel_image is not None:
            payload["kernel_image"] = kernel_image
        if initrd_path is not None:
            payload["initrd_path"] = initrd_path
        if snapshot_path is not None:
            payload["snapshot_path"] = snapshot_path
        if mem_file_path is not None:
            payload["mem_file_path"] = mem_file_path

        data = self._post("sandboxes", json=payload).json()
        return data["id"]

    def stop_sandbox(self, sandbox_id: str) -> None:
        """Stop and remove a running sandbox.

        Parameters
        ----------
        sandbox_id:
            ID of the sandbox to stop.
        """
        self._delete(f"sandboxes/{sandbox_id}")

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

        Parameters
        ----------
        sandbox_id:
            ID of the target sandbox.
        command:
            The command (and arguments) to execute.

        Returns
        -------
        ExecResult
            An object with ``stdout``, ``stderr``, and ``exit_code``
            attributes – the same shape as ``sandbox_pb2.ExecResponse``.
        """
        data = self._post(
            f"sandboxes/{sandbox_id}/exec",
            json={"command": command},
        ).json()
        return ExecResult(
            stdout=data.get("stdout", ""),
            stderr=data.get("stderr", ""),
            exit_code=int(data.get("exit_code", 0)),
        )

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

        Parameters
        ----------
        sandbox_id:
            ID of the sandbox.

        Returns
        -------
        str
            Status such as ``"running"`` or ``"stopped"``.
        """
        data = self._get(f"sandboxes/{sandbox_id}").json()
        return data.get("status", "")

    def get_metrics(self, sandbox_id: str) -> MetricsResult:
        """Return point-in-time resource-usage metrics for a sandbox.

        Parameters
        ----------
        sandbox_id:
            ID of the sandbox.

        Returns
        -------
        MetricsResult
            Dataclass with memory, CPU, network, and block-I/O counters.
        """
        data = self._get(f"sandboxes/{sandbox_id}/metrics").json()
        return MetricsResult(
            memory_usage_bytes=int(data.get("memory_usage_bytes", 0)),
            memory_limit_bytes=int(data.get("memory_limit_bytes", 0)),
            memory_percent=float(data.get("memory_percent", 0.0)),
            cpu_percent=float(data.get("cpu_percent", 0.0)),
            pids_current=int(data.get("pids_current", 0)),
            net_rx_bytes=int(data.get("net_rx_bytes", 0)),
            net_tx_bytes=int(data.get("net_tx_bytes", 0)),
            block_read_bytes=int(data.get("block_read_bytes", 0)),
            block_write_bytes=int(data.get("block_write_bytes", 0)),
        )

    def snapshot_sandbox(self, sandbox_id: str) -> Dict[str, str]:
        """Create a snapshot of a running sandbox.

        Parameters
        ----------
        sandbox_id:
            ID of the sandbox to snapshot.

        Returns
        -------
        dict
            ``{"snapshot_path": "...", "mem_file_path": "..."}`` – same
            shape returned by :meth:`SandboxClient.snapshot_sandbox`.
        """
        data = self._post(f"sandboxes/{sandbox_id}/snapshot").json()
        return {
            "snapshot_path": data.get("snapshot_path", ""),
            "mem_file_path": data.get("mem_file_path", ""),
        }

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

        Parameters
        ----------
        sandbox_id:
            ID of the sandbox to resume.

        .. note::

            This calls ``POST /sandboxes/{id}/resume``.  The endpoint must
            be registered on the orchestrator (it is available in the gRPC
            API but may not yet be wired into the HTTP router).
        """
        self._post(f"sandboxes/{sandbox_id}/resume")
