"""Agent gateway for the opencode-inside example capsule.

Adapted from an earlier agent gateway. There the gateway was
also the sandbox's only door out (LLM forward + web proxy). Here m64 owns that: the
capsule sits on an internal network, its only egress is the m64 gate (HTTP(S)_PROXY), which
lets out what the user granted at serve time (this capsule needs openrouter.ai:443, so
--allow-external-network), adds the OpenRouter key itself (the capsule holds only a
placeholder) and logs every connection. So this file keeps just the harness-neutral control API and
drives the local `opencode serve` on 127.0.0.1:4096.

Routes (the agent.session kind; the m64 front proxies only these and /health):

  GET  /health                          200 when opencode answers
  GET  /v1/agent                        what this is
  POST /v1/sessions                     new session, fresh empty directory
  POST /v1/sessions/{id}/messages       {"text": ...}; wait=true (default) blocks for the reply
  GET  /v1/sessions/{id}/events         SSE stream of this session's events
  GET  /v1/sessions/{id}/files[?path=]  list a dir / download a file
  PUT  /v1/sessions/{id}/files?path=    upload a file (raw body)
  POST /v1/sessions/{id}/abort          stop the running turn
"""
import asyncio
import json
import logging
import mimetypes
import os
import stat
import time
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from urllib.parse import quote

import httpx
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse

PORT = int(os.environ.get("M64_PORT", "9500"))
OPENCODE = os.environ.get("OPENCODE_URL", "http://127.0.0.1:4096")
PROVIDER = os.environ.get("AGENT_PROVIDER", "openrouter-compat")
MODEL = os.environ.get("AGENT_MODEL", "deepseek/deepseek-v4.1-flash")

# What the user granted this run when serving it (m64 serve --allow-external-network /
# --allow-internal-network); the m64 gate enforces it.
_REACH = {
    "": "nothing: no egress at all",
    "external": "the public internet only (never the LAN or tailnet)",
    "internal": "the LAN and tailnet only",
    "external,internal": "the public internet, the LAN and the tailnet",
}
NETWORK = _REACH.get(os.environ.get("M64_NETWORK", ""), os.environ.get("M64_NETWORK", ""))
WORK = Path(os.environ.get("WORK_DIR", "/scratch/work"))
SESSIONS_DIR = WORK / "sessions"
REGISTRY = WORK / ".gateway" / "sessions.json"

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("agent_gateway")


# ── Session registry (lives in /scratch, so it goes with the run) ────────────────

def _load_registry() -> dict:
    try:
        return json.loads(REGISTRY.read_text())
    except Exception:
        return {}


def _save_registry(reg: dict) -> None:
    REGISTRY.parent.mkdir(parents=True, exist_ok=True)
    tmp = REGISTRY.with_suffix(".tmp")
    tmp.write_text(json.dumps(reg, indent=2))
    tmp.replace(REGISTRY)


def _session(reg: dict, sid: str) -> dict:
    if sid not in reg:
        raise HTTPException(404, f"unknown session '{sid}'")
    return reg[sid]


def _safe_path(root: Path, rel: str) -> Path:
    if "\x00" in rel:
        raise HTTPException(400, "invalid path")
    try:
        p = (root / rel.lstrip("/")).resolve()
    except (OSError, ValueError, RuntimeError):  # e.g. a symlink loop
        raise HTTPException(400, "invalid path")
    if p != root.resolve() and root.resolve() not in p.parents:
        raise HTTPException(400, "path escapes the session directory")
    return p


def _kind(p: Path) -> str | None:
    """"file", "dir", "other" (FIFO, socket, device...), or None if it does not exist."""
    try:
        mode = p.stat().st_mode
    except (FileNotFoundError, NotADirectoryError):
        return None
    except OSError:
        return "other"
    return "file" if stat.S_ISREG(mode) else "dir" if stat.S_ISDIR(mode) else "other"


# Files are written by the agent, so their bytes are untrusted. They are only ever served
# as downloads, sandboxed, never sniffed, and typed as one of these inert types or bytes.
_INERT_TYPES = {"text/plain", "image/png", "image/jpeg", "image/gif", "image/webp"}
_DOWNLOAD_HEADERS = {"X-Content-Type-Options": "nosniff", "Content-Security-Policy": "sandbox"}


def _open_regular(p: Path, flags: int):
    """Open p as a regular file, never blocking on a FIFO or following a final symlink."""
    try:
        fd = os.open(p, flags | os.O_NONBLOCK | os.O_NOFOLLOW, 0o644)
    except OSError:
        raise HTTPException(400, "not a regular file")
    if not stat.S_ISREG(os.fstat(fd).st_mode):
        os.close(fd)
        raise HTTPException(400, "not a regular file")
    os.set_blocking(fd, True)
    return os.fdopen(fd, "rb" if flags == os.O_RDONLY else "wb")


def _download_type(name: str) -> str:
    guessed = mimetypes.guess_type(name)[0]
    return guessed if guessed in _INERT_TYPES else "application/octet-stream"


# ── opencode translator ─────────────────────────────────────────────────────────

class Opencode:
    """`opencode serve` HTTP API. Every call carries ?directory= so each session works in
    its own directory under $WORK/sessions."""

    def __init__(self, base: str):
        self.base = base.rstrip("/")
        # trust_env=False: this is loopback; never route it through HTTPS_PROXY.
        self.http = httpx.AsyncClient(base_url=self.base, timeout=httpx.Timeout(30, read=None),
                                      trust_env=False)

    async def healthy(self) -> bool:
        try:
            return (await self.http.get("/global/health", timeout=3)).status_code == 200
        except Exception:
            return False

    async def create(self, directory: str, title: str | None) -> str:
        r = await self.http.post("/session", params={"directory": directory},
                                 json={"title": title} if title else {})
        r.raise_for_status()
        return r.json()["id"]

    @staticmethod
    def _prompt_body(text: str) -> dict:
        return {"parts": [{"type": "text", "text": text}],
                "model": {"providerID": PROVIDER, "modelID": MODEL}}

    async def send(self, sid: str, directory: str, text: str) -> dict:
        r = await self.http.post(f"/session/{sid}/message", params={"directory": directory},
                                 json=self._prompt_body(text))
        r.raise_for_status()
        reply = self._normalize(r.json())
        # One turn = several assistant messages (one per tool step); the response is only
        # the last. Gather the tool calls of the whole turn from the history.
        turn = []
        for m in reversed(await self.history(sid, directory)):
            if m["role"] == "user":
                break
            turn.append(m)
        reply["tools"] = [t for m in reversed(turn) for t in m["tools"]]
        reply["steps"] = len(turn)
        return reply

    async def send_async(self, sid: str, directory: str, text: str) -> None:
        r = await self.http.post(f"/session/{sid}/prompt_async", params={"directory": directory},
                                 json=self._prompt_body(text))
        r.raise_for_status()

    async def history(self, sid: str, directory: str) -> list:
        r = await self.http.get(f"/session/{sid}/message", params={"directory": directory})
        r.raise_for_status()
        return [self._normalize(m) for m in r.json()]

    async def abort(self, sid: str, directory: str) -> None:
        await self.http.post(f"/session/{sid}/abort", params={"directory": directory})

    async def events(self, sid: str, directory: str):
        """Yield SSE frames for events that belong to `sid`."""
        async with self.http.stream("GET", "/event", params={"directory": directory}) as r:
            async for line in r.aiter_lines():
                if not line.startswith("data:"):
                    continue
                try:
                    ev = json.loads(line[5:])
                except ValueError:
                    continue
                props = ev.get("properties") or {}
                owner = (props.get("sessionID")
                         or (props.get("info") or {}).get("sessionID")
                         or (props.get("part") or {}).get("sessionID")
                         or (props.get("info") or {}).get("id"))
                if owner == sid or ev.get("type") == "server.connected":
                    yield f"data: {json.dumps(ev)}\n\n"

    @staticmethod
    def _normalize(msg: dict) -> dict:
        info = msg.get("info") or {}
        parts = msg.get("parts") or []
        text = "".join(p.get("text", "") for p in parts if p.get("type") == "text")
        tools = [{"tool": p.get("tool"),
                  "status": (p.get("state") or {}).get("status"),
                  "input": (p.get("state") or {}).get("input"),
                  "output": ((p.get("state") or {}).get("output") or "")[:4000]}
                 for p in parts if p.get("type") == "tool"]
        out = {"id": info.get("id"), "role": info.get("role"), "text": text, "tools": tools}
        if info.get("role") == "assistant":
            out["tokens"] = info.get("tokens")
            out["finish"] = info.get("finish")
            if info.get("error"):
                out["error"] = info["error"]
        return out


harness = Opencode(OPENCODE)


def _report(phase: str, detail: str) -> None:
    """Append a phase line to m64's status file (if running under m64)."""
    if path := os.environ.get("M64_STATUS"):
        with open(path, "a") as f:
            f.write(json.dumps({"phase": phase, "detail": detail}) + "\n")


async def _report_ready_when_up() -> None:
    while not await harness.healthy():
        await asyncio.sleep(1)
    _report("ready", f"opencode up, model {MODEL} via OpenRouter")


@asynccontextmanager
async def _lifespan(app: FastAPI):
    SESSIONS_DIR.mkdir(parents=True, exist_ok=True)
    task = asyncio.create_task(_report_ready_when_up())
    yield
    task.cancel()


app = FastAPI(title="opencode-inside agent gateway", docs_url=None, redoc_url=None,
              openapi_url=None, lifespan=_lifespan)


@app.get("/health")
async def health():
    ok = await harness.healthy()
    return JSONResponse({"status": "ok" if ok else "starting", "harness": "opencode"},
                        status_code=200 if ok else 503)


@app.get("/v1/agent")
def describe():
    return {
        "harness": "opencode", "llm": MODEL, "llm_provider": "OpenRouter",
        "sandbox": {
            "network": f"internal docker network; its only way out is the m64 gate, which lets out "
                       f"{NETWORK}, adds the OpenRouter key itself (the agent never holds it) "
                       "and logs every connection",
            "filesystem": "each session gets a fresh empty directory under the run's scratch",
            "privileges": "non-root (1000:1000), all capabilities dropped, no-new-privileges",
        },
    }


@app.post("/v1/sessions", status_code=201)
async def create_session(request: Request):
    body = await _json(request)
    reg = _load_registry()
    d = SESSIONS_DIR / uuid.uuid4().hex[:12]
    d.mkdir(parents=True)
    sid = await harness.create(str(d), body.get("title"))
    reg[sid] = {"id": sid, "title": body.get("title"), "created": time.time(),
                "directory": str(d)}
    _save_registry(reg)
    log.info("SESSION_CREATE %s dir=%s", sid, d)
    return {k: v for k, v in reg[sid].items() if k != "directory"}


@app.post("/v1/sessions/{sid}/messages")
async def send_message(sid: str, request: Request):
    body = await _json(request)
    text = body.get("text") or body.get("prompt")
    if not text or not isinstance(text, str):
        raise HTTPException(400, "body needs 'text'")
    s = _session(_load_registry(), sid)
    if body.get("wait", True):
        t0 = time.time()
        reply = await harness.send(sid, s["directory"], text)
        return {**reply, "elapsed_s": round(time.time() - t0, 1)}
    await harness.send_async(sid, s["directory"], text)
    return JSONResponse({"session_id": sid, "status": "busy",
                         "events": f"GET /v1/sessions/{sid}/events"}, status_code=202)


@app.post("/v1/sessions/{sid}/abort")
async def abort(sid: str):
    s = _session(_load_registry(), sid)
    await harness.abort(sid, s["directory"])
    return {"aborted": sid}


@app.get("/v1/sessions/{sid}/events")
async def events(sid: str):
    s = _session(_load_registry(), sid)
    return StreamingResponse(harness.events(sid, s["directory"]), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache"})


@app.get("/v1/sessions/{sid}/files")
async def get_files(sid: str, path: str = "."):
    s = _session(_load_registry(), sid)
    root = Path(s["directory"])
    p = _safe_path(root, path)
    kind = _kind(p)
    if kind is None:
        raise HTTPException(404, f"{path} not found")
    if kind == "other":
        raise HTTPException(400, f"{path} is not a regular file or directory")
    if kind == "dir":
        entries = []
        for c in sorted(p.iterdir()):
            k = _kind(c) or "other"
            entries.append({"name": c.name, "type": k,
                            "size": c.stat().st_size if k == "file" else None})
        return {"path": str(p.relative_to(root)) if p != root else ".", "entries": entries}
    # The agent may swap the file for a FIFO after the check above: open without
    # blocking and check what was actually opened.
    f = _open_regular(p, os.O_RDONLY)
    size = os.fstat(f.fileno()).st_size

    def chunks():
        with f:
            while data := f.read(65536):
                yield data

    headers = {**_DOWNLOAD_HEADERS, "Content-Length": str(size),
               "Content-Disposition": "attachment; filename*=UTF-8''" + quote(p.name)}
    return StreamingResponse(chunks(), media_type=_download_type(p.name), headers=headers)


@app.put("/v1/sessions/{sid}/files", status_code=201)
async def put_file(sid: str, request: Request, path: str = ""):
    if not path.strip("/."):
        raise HTTPException(400, "?path= names the file to write")
    s = _session(_load_registry(), sid)
    p = _safe_path(Path(s["directory"]), path)
    if _kind(p) not in (None, "file"):
        raise HTTPException(400, f"{path} exists and is not a regular file")
    try:
        p.parent.mkdir(parents=True, exist_ok=True)
    except (FileExistsError, NotADirectoryError):
        raise HTTPException(400, f"a parent of {path} is not a directory")
    data = await request.body()
    with _open_regular(p, os.O_WRONLY | os.O_CREAT | os.O_TRUNC) as f:
        f.write(data)
    return {"written": path, "bytes": len(data)}


async def _json(request: Request) -> dict:
    raw = await request.body()
    if not raw:
        return {}
    try:
        data = json.loads(raw)
    except ValueError:
        raise HTTPException(400, "body must be JSON")
    if not isinstance(data, dict):
        raise HTTPException(400, "body must be a JSON object")
    return data


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")