"""Python SDK of the sandbox gateway: sandboxes, exec, files, git, background sessions, LSP.

    from sandbox_gateway import Gateway
    gw = Gateway("https://sandboxes.example.com", "token")   # or SANDBOX_GATEWAY_URL / _TOKEN
    with gw.create("python:3.14.8-slim", cpus=1, memory_mb=1024) as sbx:
        print(sbx.exec("python -V").stdout)
        sbx.fs.upload(b"print(1)\\n", "/work/a.py")
        s = sbx.sessions.start("python -u /work/a.py")
        for chunk in s.logs(follow=True):
            print(chunk.decode(), end="")

The gateway and sandboxd only run argv and keep background sessions; files, git and LSP are built on
those two and need the tools in the image: coreutils/busybox (`stat`, `find`, `grep`, `chmod`), `git`
and the language server.
"""

from __future__ import annotations

import base64
import gzip
import io
import itertools
import json
import os
import threading
import time
import zlib
from collections.abc import Iterator
from concurrent.futures import Future
from dataclasses import dataclass
from typing import BinaryIO, Self

import httpx

POLL_SEC = 30
# A just-created sandbox may be unknown to another gateway replica until its informer sees the Pod.
CREATE_GRACE_SEC = 60
CHUNK = 1 << 20


# The gateway's HTTP API versions this SDK speaks (GET /version, "api"); bumped only by a breaking change.
API_VERSIONS = (1,)


class GatewayError(Exception):
    def __init__(self, message: str, code: str = ""):
        super().__init__(message)
        self.code = code  # the gateway's stable error code, e.g. "not_found", "invalid_token"


@dataclass
class ExecResult:
    exit_code: int
    stdout: str
    stderr: str
    timed_out: bool = False
    raw_stdout: bytes = b""


def _argv(cmd: str | list[str]) -> list[str]:
    return ["sh", "-c", cmd] if isinstance(cmd, str) else cmd


RETRIES = 6
RETRY_STATUS = (502, 503, 504)


def _idempotent(send, *, deadline: float | None = None) -> httpx.Response:
    """GET and DELETE: retried while a gateway replica restarts (502/503/504, connection errors).
    Raises TimeoutError once `deadline` has passed."""
    for attempt in range(RETRIES + 1):
        if deadline is not None and time.monotonic() >= deadline:
            raise TimeoutError("request deadline exceeded")
        try:
            r = send()
        except httpx.TransportError:
            if deadline is not None and time.monotonic() >= deadline:
                raise TimeoutError("request deadline exceeded")
            if attempt == RETRIES:
                raise
        else:
            if deadline is not None and time.monotonic() >= deadline:
                raise TimeoutError("request deadline exceeded")
            if r.status_code not in RETRY_STATUS or attempt == RETRIES:
                return r
        delay = min(2**attempt * 0.5, 8)
        time.sleep(delay if deadline is None else min(delay, max(0, deadline - time.monotonic())))
    raise AssertionError("unreachable")


def _check(r: httpx.Response) -> httpx.Response:
    if r.status_code >= 400:
        try:
            code = r.json().get("code", "")
        except ValueError:
            code = ""
        raise GatewayError(f"{r.request.method} {r.request.url.path}: {r.status_code} {r.text[:1000]}", code)
    return r


class Gateway:
    def __init__(self, url: str | None = None, token: str | None = None):
        url = url or os.environ.get("SANDBOX_GATEWAY_URL")
        if not url:
            raise ValueError("sandbox gateway URL missing: pass url=... or set SANDBOX_GATEWAY_URL")
        token = token or os.environ.get("SANDBOX_GATEWAY_TOKEN", "")
        self.http = httpx.Client(base_url=url.rstrip("/"), headers={"Authorization": f"Bearer {token}"},
                                 timeout=httpx.Timeout(60, connect=10))
        self._api_checked = False

    def check_version(self) -> None:
        """Refuses a gateway whose API this SDK does not speak, before anything is created.
        A gateway without /version predates it and speaks API 1."""
        r = _idempotent(lambda: self.http.get("/version"))
        api = 1 if r.status_code == 404 else _check(r).json().get("api")
        if api not in API_VERSIONS:
            raise GatewayError(f"the gateway speaks API {api}, this SDK {API_VERSIONS}: "
                               f"update the SDK from {self.http.base_url}/clients/sandbox_gateway.py")
        self._api_checked = True

    def create(self, image: str, *, cpus: float, memory_mb: int, storage_mb: int | None = None,
               env: dict[str, str] | None = None, name: str | None = None, labels: dict[str, str] | None = None,
               wait: bool = True, timeout: float | None = None) -> Sandbox:
        """Creates a sandbox. With `wait`, blocks until it is ready (it may queue on the team's quota)
        for at most `timeout` seconds (None: indefinitely). A sandbox that does not come up is deleted."""
        if not self._api_checked:
            self.check_version()
        spec = {"image": image, "cpus": cpus, "memory_mb": memory_mb, "env": env or {}}
        if storage_mb:
            spec["storage_mb"] = storage_mb
        if name:
            spec["name"] = name
        if labels:
            spec["labels"] = labels
        sid = _check(self.http.post("/sandboxes", json=spec)).json()["id"]
        sbx = self.get(sid)
        if wait:
            try:
                sbx.wait_ready(timeout)
            except BaseException as error:
                try:
                    sbx.delete()
                except Exception as cleanup:
                    error.add_note(f"sandbox cleanup failed: {cleanup}")
                raise
        return sbx

    def get(self, sandbox_id: str) -> Sandbox:
        return Sandbox(self.http, f"/sandboxes/{sandbox_id}", sandbox_id)

    def list(self, **labels: str | None) -> list[dict]:
        """The team's sandboxes. list(tool="rl", run=None) selects label tool=rl with any value of run."""
        q = [("label", k if v is None else f"{k}={v}") for k, v in labels.items()]
        return _check(self.http.get("/sandboxes", params=q)).json()


class Sandbox:
    def __init__(self, http: httpx.Client, base: str, sandbox_id: str = ""):
        self.http, self.base, self.id = http, base, sandbox_id
        self.fs = Files(self)
        self.git = Git(self)
        self.sessions = Sessions(self)

    def __enter__(self) -> Self:
        return self

    def __exit__(self, *exc) -> None:
        try:
            self.delete()
        except Exception as cleanup:
            if exc[1] is None:
                raise
            exc[1].add_note(f"sandbox cleanup failed: {cleanup}")

    def state(self, wait: float = 0) -> dict:
        return _check(_idempotent(lambda: self.http.get(self.base, params={"wait": wait}, timeout=wait + 30))).json()

    def wait_ready(self, timeout: float | None = None) -> None:
        """Wait at most timeout seconds; None waits indefinitely, zero fails immediately."""
        grace = time.monotonic() + CREATE_GRACE_SEC
        deadline = time.monotonic() + timeout if timeout is not None else None

        def poll():
            remaining = POLL_SEC + 30 if deadline is None else max(0, deadline - time.monotonic())
            return self.http.get(self.base, params={"wait": min(POLL_SEC, remaining)}, timeout=remaining)

        while True:
            try:
                r = _idempotent(poll, deadline=deadline)
            except TimeoutError as error:
                raise GatewayError(f"sandbox {self.id} not ready after {timeout}s") from error
            if r.status_code == 404 and time.monotonic() < grace:
                time.sleep(1 if deadline is None else min(1, max(0, deadline - time.monotonic())))
                continue
            if (s := _check(r).json())["state"] in ("ready", "failed"):
                break
        if s["state"] == "failed":
            raise GatewayError(f"sandbox {self.id} failed: {s.get('reason')}")

    def diagnostics(self) -> dict:
        """State, Pod, containers (reasons, exit codes, resources) and events; does not renew the lease."""
        return _check(_idempotent(lambda: self.http.get(f"{self.base}/diagnostics"))).json()

    def collect_diagnostics(self, dest: str | os.PathLike, *, log_lines: int = 200) -> dict[str, str]:
        """Saves diagnostics.json and main.log into the directory dest (e.g. before delete). Best
        effort: returns {file: error} for what could not be read, and saves the rest."""
        os.makedirs(dest, exist_ok=True)
        errors = {}
        for name, get in (("diagnostics.json", lambda: json.dumps(self.diagnostics(), indent=2).encode()),
                          ("main.log", lambda: _check(_idempotent(lambda: self.http.get(
                              f"{self.base}/logs", params={"container": "main", "tail": log_lines}))).content)):
            try:
                data = get()
            except (GatewayError, httpx.HTTPError) as e:
                errors[name] = str(e)[:500]
                continue
            with open(os.path.join(dest, name), "wb") as f:
                f.write(data)
        return errors

    def delete(self) -> None:
        r = _idempotent(lambda: self.http.delete(self.base))
        if r.status_code != 404:  # a retried delete may find it gone
            _check(r)

    def exec(self, cmd: str | list[str], *, stdin: bytes = b"", cwd: str | None = None,
             env: dict[str, str] | None = None, timeout: float | None = None,
             container: str | None = None) -> ExecResult:
        """Run to completion. A string goes through `sh -c`, a list is argv as is."""
        argv = _argv(cmd)
        out = io.BytesIO()
        code, err, timed_out = self.exec_stream(argv, stdin=stdin, stdout=out, timeout=timeout,
                                                container=container, env=env, cwd=cwd)
        raw = out.getvalue()
        return ExecResult(code, raw.decode(errors="replace"), err.decode(errors="replace"), timed_out, raw)

    def exec_stream(self, argv: list[str], *, stdin: bytes | BinaryIO = b"", stdout: BinaryIO,
                    timeout: float | None = None, container: str | None = None,
                    env: dict[str, str] | None = None, cwd: str | None = None) -> tuple[int, bytes, bool]:
        """argv with stdin (bytes or a file) streamed in and stdout streamed to a file:
        memory stays flat whatever the size. Returns (exit code, stderr, timed out)."""
        line = json.dumps({"argv": argv, "timeout_sec": timeout or 0, "env": env, "cwd": cwd}).encode() + b"\n"
        headers = {}
        if isinstance(stdin, bytes):
            body = line + stdin
            if len(body) >= 16 << 10:  # file uploads: gzipped on the wire, the gateway unpacks them
                body, headers = gzip.compress(body, compresslevel=1), {"Content-Encoding": "gzip"}
        else:
            def chunks(f=stdin):
                z = zlib.compressobj(1, wbits=31)
                yield z.compress(line)
                while chunk := f.read(CHUNK):
                    yield z.compress(chunk)
                yield z.flush()
            body, headers = chunks(), {"Content-Encoding": "gzip"}
        with self.http.stream("POST", f"{self.base}/exec", content=body, headers=headers,
                              params={"container": container} if container else None,
                              timeout=httpx.Timeout(timeout + 60 if timeout else None, connect=10)) as r:
            if r.status_code >= 400:
                r.read()
                _check(r)
            left, err = int(r.headers["X-Stdout-Length"]), bytearray()
            for chunk in r.iter_bytes(CHUNK):
                if left:
                    stdout.write(chunk[:left])
                    chunk, left = chunk[left:], max(0, left - len(chunk))
                err += chunk
            if left:
                raise GatewayError(f"exec: stdout cut short by {left} bytes")
            return int(r.headers["X-Return-Code"]), bytes(err), r.headers.get("X-Timed-Out") == "1"

    def run(self, cmd: str | list[str], **kw) -> ExecResult:
        """exec that raises on a non-zero exit."""
        r = self.exec(cmd, **kw)
        if r.exit_code != 0:
            raise GatewayError(f"{cmd!r} exited {r.exit_code}: {r.stderr[-2000:]}")
        return r


@dataclass
class FileInfo:
    name: str
    size: int
    mode: str  # octal, like "644"
    is_dir: bool
    mtime: int


class Files:
    def __init__(self, sbx: Sandbox):
        self.sbx = sbx

    def list(self, path: str) -> list[FileInfo]:
        script = ('cd "$1" || exit 1; for f in * .[!.]* ..?*; do [ -e "$f" ] || [ -L "$f" ] || continue; '
                  "stat -c '%s\t%a\t%F\t%Y\t%n' \"$f\"; done")
        out = self.sbx.run(["sh", "-c", script, "sh", path]).stdout
        entries = []
        for line in out.splitlines():
            size, mode, kind, mtime, name = line.split("\t", 4)
            entries.append(FileInfo(name, int(size), mode, kind == "directory", int(mtime)))
        return entries

    def info(self, path: str) -> FileInfo:
        size, mode, kind, mtime, name = self.sbx.run(["stat", "-c", "%s\t%a\t%F\t%Y\t%n", path]).stdout.rstrip("\n").split("\t", 4)
        return FileInfo(name, int(size), mode, kind == "directory", int(mtime))

    def search(self, path: str, pattern: str) -> list[str]:
        """Paths under `path` whose name matches the glob `pattern`."""
        return self.sbx.run(["find", path, "-name", pattern]).stdout.splitlines()

    def grep(self, path: str, text: str) -> list[tuple[str, int, str]]:
        """(file, line, content) of every line containing `text` (fixed string) under `path`."""
        r = self.sbx.exec(["grep", "-rnIF", "--", text, path])
        if r.exit_code > 1:
            raise GatewayError(r.stderr)
        hits = []
        for line in r.stdout.splitlines():
            f, n, content = line.split(":", 2)
            hits.append((f, int(n), content))
        return hits

    def chmod(self, path: str, mode: str | int, *, recursive: bool = False) -> None:
        mode = format(mode, "o") if isinstance(mode, int) else mode
        self.sbx.run(["chmod", *(["-R"] if recursive else []), mode, path])

    def chown(self, path: str, owner: str, *, recursive: bool = False) -> None:
        self.sbx.run(["chown", *(["-R"] if recursive else []), owner, path])

    def mkdir(self, path: str) -> None:
        self.sbx.run(["mkdir", "-p", path])

    def remove(self, path: str) -> None:
        self.sbx.run(["rm", "-rf", "--", path])

    def move(self, src: str, dst: str) -> None:
        self.sbx.run(["mv", "--", src, dst])

    def upload(self, data: bytes, path: str) -> None:
        self.sbx.run(["sh", "-c", 'mkdir -p "$(dirname "$1")" && cat >"$1"', "sh", path], stdin=data)

    def download(self, path: str) -> bytes:
        out = io.BytesIO()
        self.download_to(path, out)
        return out.getvalue()

    def download_to(self, path: str, dest: str | os.PathLike | BinaryIO) -> None:
        """Streams the sandbox file at `path` to `dest`: a local path or a binary file object."""
        if not hasattr(dest, "write"):
            with open(dest, "wb") as f:
                return self.download_to(path, f)
        code, err, _ = self.sbx.exec_stream(["cat", "--", path], stdout=dest)
        if code != 0:
            raise GatewayError(f"download {path}: exit {code}: {err.decode(errors='replace')[-2000:]}")

    def upload_from(self, src: str | os.PathLike | BinaryIO, path: str) -> None:
        """Streams a local file (path or binary file object) to `path` in the sandbox."""
        if not hasattr(src, "read"):
            with open(src, "rb") as f:
                return self.upload_from(f, path)
        code, err, _ = self.sbx.exec_stream(["sh", "-c", 'mkdir -p "$(dirname "$1")" && cat >"$1"', "sh", path],
                                            stdin=src, stdout=io.BytesIO())
        if code != 0:
            raise GatewayError(f"upload {path}: exit {code}: {err.decode(errors='replace')[-2000:]}")


class Git:
    """git CLI in the sandbox. Credentials go as an HTTP header for one command (not into
    .git/config) and in a session's environment rather than argv: the gateway logs argv."""

    def __init__(self, sbx: Sandbox):
        self.sbx = sbx

    def _git(self, path: str | None, *args: str, username: str | None = None,
             password: str | None = None) -> ExecResult:
        argv = ["git", *(["-C", path] if path else []), *args]
        if not password:
            return self.sbx.run(argv)
        basic = base64.b64encode(f"{username or 'git'}:{password}".encode()).decode()
        s = self.sbx.sessions.start(argv, env={"GIT_CONFIG_COUNT": "1", "GIT_CONFIG_KEY_0": "http.extraHeader",
                                               "GIT_CONFIG_VALUE_0": f"Authorization: Basic {basic}",
                                               "GIT_TERMINAL_PROMPT": "0"})
        try:
            code = s.wait()
            r = ExecResult(code, s.output().decode(errors="replace"), s.output("stderr").decode(errors="replace"))
        finally:
            s.kill()
        if code != 0:
            raise GatewayError(f"{argv!r} exited {code}: {r.stderr[-2000:]}")
        return r

    def clone(self, url: str, path: str, *, branch: str | None = None, commit: str | None = None,
              username: str | None = None, password: str | None = None) -> None:
        self._git(None, "clone", *(["--branch", branch] if branch else []), url, path,
                  username=username, password=password)
        if commit:
            self._git(path, "checkout", commit)

    def status(self, path: str) -> dict:
        """{"branch", "ahead", "behind", "files": [(xy, path)]} from `git status --porcelain --branch`."""
        out = self._git(path, "status", "--porcelain=v1", "--branch").stdout.splitlines()
        head, files = out[0][3:], [(line[:2], line[3:]) for line in out[1:]]
        head = head.removeprefix("No commits yet on ")
        branch, _, track = head.partition("...")
        counts = {"ahead": 0, "behind": 0}
        if "[" in track:  # "origin/main [ahead 1, behind 2]" or "[gone]"
            for part in track[track.index("[") + 1:-1].split(", "):
                word, _, n = part.partition(" ")
                if word in counts:
                    counts[word] = int(n)
        return {"branch": branch, **counts, "files": files}

    def add(self, path: str, files: list[str] | None = None) -> None:
        self._git(path, "add", "--", *(files or ["."]))

    def commit(self, path: str, message: str, *, author: str, email: str) -> str:
        self._git(path, "-c", f"user.name={author}", "-c", f"user.email={email}", "commit", "-m", message)
        return self._git(path, "rev-parse", "HEAD").stdout.strip()

    def branches(self, path: str) -> list[str]:
        return self._git(path, "branch", "--format=%(refname:short)").stdout.split()

    def checkout(self, path: str, branch: str, *, create: bool = False) -> None:
        self._git(path, "checkout", *(["-b"] if create else []), branch)

    def push(self, path: str, *, username: str | None = None, password: str | None = None) -> None:
        self._git(path, "push", username=username, password=password)

    def pull(self, path: str, *, username: str | None = None, password: str | None = None) -> None:
        self._git(path, "pull", username=username, password=password)


class Session:
    """A background process in the sandbox (sandboxd/sessions.go)."""

    def __init__(self, sbx: Sandbox, info: dict):
        self.sbx, self.id, self.info = sbx, info["id"], info
        self.url = f"{sbx.base}/sessions/{self.id}"

    def status(self) -> dict:
        self.info = _check(self.sbx.http.get(self.url)).json()
        return self.info

    def logs(self, stream: str = "stdout", *, offset: int = 0, follow: bool = False) -> Iterator[bytes]:
        """Output from `offset`; with `follow` it streams until the process exits."""
        params = {"stream": stream, "offset": offset, "follow": int(follow)}
        with self.sbx.http.stream("GET", f"{self.url}/logs", params=params,
                                  timeout=httpx.Timeout(None, connect=10)) as r:
            if r.status_code >= 400:
                r.read()
                _check(r)
            yield from r.iter_bytes()

    def output(self, stream: str = "stdout") -> bytes:
        return b"".join(self.logs(stream))

    def write(self, data: bytes | str, *, close: bool = False) -> None:
        _check(self.sbx.http.post(f"{self.url}/stdin", content=data.encode() if isinstance(data, str) else data,
                                  params={"close": 1} if close else None))

    def wait(self, poll: float = 0.2) -> int:
        while (s := self.status())["running"]:
            time.sleep(poll)
        if s.get("output_error"):
            raise GatewayError("session output: " + s["output_error"])
        return s["exit_code"]

    def kill(self) -> None:
        _check(self.sbx.http.delete(self.url))


class Sessions:
    def __init__(self, sbx: Sandbox):
        self.sbx = sbx

    def start(self, cmd: str | list[str], *, cwd: str | None = None, env: dict[str, str] | None = None) -> Session:
        body = {"argv": _argv(cmd), "cwd": cwd or "", "env": env or {}}
        return Session(self.sbx, _check(self.sbx.http.post(f"{self.sbx.base}/sessions", json=body)).json())

    def list(self) -> list[dict]:
        return _check(self.sbx.http.get(f"{self.sbx.base}/sessions")).json()

    def get(self, session_id: str) -> Session:
        s = Session(self.sbx, {"id": session_id})
        s.status()
        return s


class LspServer:
    """A language server in the sandbox, spoken to over JSON-RPC through a session's stdin and stdout.

        lsp = LspServer(sbx, ["pylsp"], "/work")        # typescript-language-server --stdio, gopls, ...
        lsp.did_open("/work/a.py", "python", text)
        lsp.completions("/work/a.py", line=0, character=3)
    """

    def __init__(self, sbx: Sandbox, argv: list[str], root: str):
        self.root = root
        self.session = sbx.sessions.start(argv, cwd=root)
        self._ids = itertools.count(1)
        # The reader thread and callers share these: _lock keeps "dead" and "pending" consistent.
        self._lock = threading.Lock()
        self._pending: dict[int, Future] = {}
        self._dead = False
        self._reader = threading.Thread(target=self._read, daemon=True)
        self._reader.start()
        self.request("initialize", {"processId": None, "rootUri": f"file://{root}", "capabilities": {}})
        self.notify("initialized", {})

    def _send(self, msg: dict) -> None:
        body = json.dumps(msg).encode()
        self.session.write(b"Content-Length: %d\r\n\r\n" % len(body) + body)

    def _read(self) -> None:
        buf = b""
        try:
            for chunk in self.session.logs(follow=True):
                buf += chunk
                while (end := buf.find(b"\r\n\r\n")) >= 0:
                    headers = dict(h.split(": ", 1) for h in buf[:end].decode().split("\r\n") if ": " in h)
                    n = int(headers["Content-Length"])
                    if len(buf) < end + 4 + n:
                        break
                    msg, buf = json.loads(buf[end + 4:end + 4 + n]), buf[end + 4 + n:]
                    if "method" in msg:
                        if "id" in msg:  # a request from the server: answer it, never touch _pending
                            self._answer(msg)
                        continue
                    with self._lock:
                        f = self._pending.pop(msg.get("id"), None)
                    if f:
                        if "error" in msg:
                            f.set_exception(GatewayError(f"lsp: {msg['error']}"))
                        else:
                            f.set_result(msg.get("result"))
        finally:
            with self._lock:
                self._dead = True
                pending, self._pending = self._pending, {}
            for f in pending.values():
                f.set_exception(GatewayError("language server exited"))

    def _answer(self, msg: dict) -> None:
        """workspace/configuration gets nulls (defaults), anything else "method not found"."""
        if msg["method"] == "workspace/configuration":
            reply = {"result": [None] * len(msg.get("params", {}).get("items", []))}
        else:
            reply = {"error": {"code": -32601, "message": f"unsupported: {msg['method']}"}}
        try:
            self._send({"jsonrpc": "2.0", "id": msg["id"], **reply})
        except GatewayError:
            pass

    def request(self, method: str, params: dict, timeout: float = 60):
        rid = next(self._ids)
        with self._lock:
            if self._dead:
                raise GatewayError("language server exited")
            self._pending[rid] = f = Future()
        try:
            self._send({"jsonrpc": "2.0", "id": rid, "method": method, "params": params})
            return f.result(timeout)
        finally:
            with self._lock:
                self._pending.pop(rid, None)

    def notify(self, method: str, params: dict) -> None:
        self._send({"jsonrpc": "2.0", "method": method, "params": params})

    def did_open(self, path: str, language_id: str, text: str) -> None:
        self.notify("textDocument/didOpen", {"textDocument": {
            "uri": f"file://{path}", "languageId": language_id, "version": 1, "text": text}})

    def completions(self, path: str, line: int, character: int):
        return self.request("textDocument/completion", {
            "textDocument": {"uri": f"file://{path}"}, "position": {"line": line, "character": character}})

    def document_symbols(self, path: str):
        return self.request("textDocument/documentSymbol", {"textDocument": {"uri": f"file://{path}"}})

    def workspace_symbols(self, query: str):
        return self.request("workspace/symbol", {"query": query})

    def stop(self) -> None:
        try:
            self.request("shutdown", {}, timeout=10)
            self.notify("exit", {})
        finally:
            self.session.kill()