#!/usr/bin/env python3
"""sandbox-agent — stdlib-only in-sandbox telemetry capture agent (REQ-3-003).

D-031: this file is copied into the sandbox namespace and runs against the
system Python — no third-party packages are importable there, so this module
depends on the standard library ONLY (the test suite enforces this with an
AST scan of the file).

What it does:
  * wraps a non-interactive `/bin/sh` REPL: each stdin line is executed via
    `sh -c` inside the workspace and reported as `stdin` -> `command` ->
    `stdout` -> `run_result`/`test_result` events;
  * polls the workspace tree (~250 ms) and emits `file_diff` events
    (created/modified/deleted with unified diffs) plus periodic `activity`
    heartbeats;
  * streams events to ai-service as TelemetryEvent-shaped JSON frames over a
    raw-socket RFC 6455 WebSocket client (no `websockets` package exists in
    the namespace — the client handshake + frame codec is implemented here);
  * at-least-once delivery (D-026): every event is appended to an fsync'd
    JSONL spool file inside the workdir BEFORE any send attempt; on
    disconnect the spool grows; after reconnect (exponential backoff) the
    spool is flushed oldest-first. The server dedups on (learner, task, seq)
    so replayed duplicates are harmless — loss is not tolerated.

Configured entirely through env baked at spawn time:
  NC_LEARNER_ID / NC_TASK_ID / NC_INGEST_URL / NC_SANDBOX_ID   (required)
  NC_WORKSPACE            workspace root to watch/run in (default: cwd)
  NC_SPOOL                spool path (default: <workspace>/.nc-agent/spool.jsonl)
  NC_POLL_INTERVAL_S / NC_ACTIVITY_INTERVAL_S / NC_COMMAND_TIMEOUT_S
  NC_BACKOFF_BASE_S / NC_BACKOFF_MAX_S                          (optional knobs)

Sequencing survives process restarts (incl. SIGKILL): on boot the spool is
replayed into the pending queue and `seq` resumes at max(spooled seq) + 1.
"""

from __future__ import annotations

import base64
import difflib
import hashlib
import json
import os
import secrets
import socket
import ssl
import struct
import subprocess
import sys
import threading
import time
import urllib.parse
from collections import deque
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

_WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
_EVENT_KINDS = frozenset(
    {"command", "file_diff", "run_result", "test_result", "activity", "stdin", "stdout"}
)
_AGENT_DIR_PREFIX = ".nc-"  # agent-private paths (spool) are excluded from watching
_MAX_DIFF_BYTES = 64 * 1024  # files larger than this are reported truncated, no diff
_MAX_OUTPUT_CHARS = 64 * 1024  # captured stdout/stderr tail cap per command
_HANDSHAKE_MAX_BYTES = 64 * 1024


# --------------------------------------------------------------------------- config


@dataclass(frozen=True)
class AgentConfig:
    """Runtime configuration, normally built from `NC_*` env baked at spawn."""

    learner_id: str
    task_id: str
    ingest_url: str
    sandbox_id: str
    workspace: Path
    spool_path: Path
    poll_interval_s: float = 0.25
    activity_interval_s: float = 5.0
    command_timeout_s: float = 30.0
    backoff_base_s: float = 0.25
    backoff_max_s: float = 8.0

    def __post_init__(self) -> None:
        for name in ("learner_id", "task_id", "ingest_url", "sandbox_id"):
            if not getattr(self, name):
                raise ValueError(f"missing required config: NC_{name.upper()}")

    @classmethod
    def from_env(cls, env: Mapping[str, str] | None = None) -> AgentConfig:
        src = os.environ if env is None else env
        workspace = Path(src.get("NC_WORKSPACE") or os.getcwd()).resolve()
        return cls(
            learner_id=src.get("NC_LEARNER_ID", ""),
            task_id=src.get("NC_TASK_ID", ""),
            ingest_url=src.get("NC_INGEST_URL", ""),
            sandbox_id=src.get("NC_SANDBOX_ID", ""),
            workspace=workspace,
            spool_path=Path(
                src.get("NC_SPOOL") or (workspace / ".nc-agent" / "spool.jsonl")
            ),
            poll_interval_s=float(src.get("NC_POLL_INTERVAL_S", "0.25")),
            activity_interval_s=float(src.get("NC_ACTIVITY_INTERVAL_S", "5.0")),
            command_timeout_s=float(src.get("NC_COMMAND_TIMEOUT_S", "30.0")),
            backoff_base_s=float(src.get("NC_BACKOFF_BASE_S", "0.25")),
            backoff_max_s=float(src.get("NC_BACKOFF_MAX_S", "8.0")),
        )


# --------------------------------------------------------------------------- spool


class Spool:
    """Append-only JSONL spool with per-append fsync (survives SIGKILL).

    `rewrite` swaps in a compacted file atomically (tmp file + os.replace).
    Lines are stored without trailing newlines in memory, one per line on disk.
    """

    def __init__(self, path: Path) -> None:
        self._path = path
        path.parent.mkdir(parents=True, exist_ok=True)

    @property
    def path(self) -> Path:
        return self._path

    def append(self, line: str) -> None:
        with self._path.open("a", encoding="utf-8") as fh:
            fh.write(line + "\n")
            fh.flush()
            os.fsync(fh.fileno())

    def read_all(self) -> list[str]:
        if not self._path.exists():
            return []
        with self._path.open("r", encoding="utf-8") as fh:
            return [line.rstrip("\n") for line in fh if line.strip()]

    def rewrite(self, lines: list[str]) -> None:
        tmp = self._path.with_name(self._path.name + ".tmp")
        with tmp.open("w", encoding="utf-8") as fh:
            for line in lines:
                fh.write(line + "\n")
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(tmp, self._path)


# --------------------------------------------------------------- websocket codec


def _encode_frame(opcode: int, payload: bytes) -> bytes:
    """RFC 6455 client frame: FIN set, always masked (servers require it)."""
    header = bytearray([0x80 | opcode])
    n = len(payload)
    if n < 126:
        header.append(0x80 | n)
    elif n < 65536:
        header.append(0x80 | 126)
        header += struct.pack("!H", n)
    else:
        header.append(0x80 | 127)
        header += struct.pack("!Q", n)
    mask = secrets.token_bytes(4)
    header += mask
    masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
    return bytes(header) + masked


class WsConnection:
    """Minimal blocking RFC 6455 client over a raw socket (stdlib only)."""

    def __init__(self, sock: socket.socket) -> None:
        self._sock = sock
        self._write_lock = threading.Lock()

    @classmethod
    def connect(cls, url: str, timeout_s: float = 5.0) -> WsConnection:
        parts = urllib.parse.urlsplit(url)
        if parts.scheme not in ("ws", "wss"):
            raise ValueError(f"unsupported scheme in NC_INGEST_URL: {parts.scheme!r}")
        host = parts.hostname or "localhost"
        port = parts.port or (443 if parts.scheme == "wss" else 80)
        path = parts.path or "/"
        if parts.query:
            path += "?" + parts.query

        sock = socket.create_connection((host, port), timeout=timeout_s)
        if parts.scheme == "wss":
            sock = ssl.create_default_context().wrap_socket(sock, server_hostname=host)

        key = base64.b64encode(secrets.token_bytes(16)).decode("ascii")
        request = (
            f"GET {path} HTTP/1.1\r\n"
            f"Host: {host}:{port}\r\n"
            "Upgrade: websocket\r\n"
            "Connection: Upgrade\r\n"
            f"Sec-WebSocket-Key: {key}\r\n"
            "Sec-WebSocket-Version: 13\r\n\r\n"
        )
        sock.sendall(request.encode("ascii"))
        response = cls._read_http_response(sock)
        cls._validate_handshake(response, key)
        return cls(sock)

    @staticmethod
    def _read_http_response(sock: socket.socket) -> bytes:
        buf = b""
        while b"\r\n\r\n" not in buf:
            chunk = sock.recv(4096)
            if not chunk:
                raise ConnectionError("server closed during WebSocket handshake")
            buf += chunk
            if len(buf) > _HANDSHAKE_MAX_BYTES:
                raise ConnectionError("handshake response exceeded size cap")
        return buf.split(b"\r\n\r\n", 1)[0]

    @staticmethod
    def _validate_handshake(response: bytes, key: str) -> None:
        head = response.decode("latin-1")
        lines = head.split("\r\n")
        if not lines or " 101" not in lines[0]:
            raise ConnectionError(f"handshake rejected: {lines[0] if lines else '<empty>'}")
        headers = {}
        for line in lines[1:]:
            if ":" in line:
                name, _, value = line.partition(":")
                headers[name.strip().lower()] = value.strip()
        expect = base64.b64encode(
            hashlib.sha1((key + _WS_GUID).encode("ascii")).digest()
        ).decode("ascii")
        if headers.get("sec-websocket-accept") != expect:
            raise ConnectionError("bad Sec-WebSocket-Accept in handshake response")

    # -- send ------------------------------------------------------------
    def send_text(self, text: str) -> None:
        with self._write_lock:
            self._sock.sendall(_encode_frame(0x1, text.encode("utf-8")))

    def _send_frame(self, opcode: int, payload: bytes) -> None:
        with self._write_lock:
            self._sock.sendall(_encode_frame(opcode, payload))

    # -- receive ---------------------------------------------------------
    def recv_message(self, timeout_s: float) -> tuple[int, bytes] | None:
        """Return (opcode, payload) for a data/close frame, or None on timeout.

        Ping frames are answered with pong internally and never surfaced;
        pongs are swallowed. Fragmented messages are reassembled. Raises
        ConnectionError/OSError when the socket breaks.
        """
        deadline = time.monotonic() + timeout_s
        fragments = bytearray()
        frag_opcode = 0
        while True:
            frame = self._recv_one_frame(deadline)
            if frame is None:
                return None
            fin, opcode, payload = frame
            if opcode == 0x9:  # ping
                self._send_frame(0xA, payload)
                continue
            if opcode == 0xA:  # pong
                continue
            if opcode == 0x0:  # continuation
                fragments += payload
            else:
                fragments = bytearray(payload)
                frag_opcode = opcode
            if fin:
                return frag_opcode, bytes(fragments)

    def _recv_one_frame(self, deadline: float) -> tuple[bool, int, bytes] | None:
        header = self._read_exact(2, deadline)
        if header is None:
            return None
        b0, b1 = header[0], header[1]
        fin = bool(b0 & 0x80)
        opcode = b0 & 0x0F
        length = b1 & 0x7F
        if length == 126:
            ext = self._read_exact(2, deadline)
            if ext is None:
                return None
            length = struct.unpack("!H", ext)[0]
        elif length == 127:
            ext = self._read_exact(8, deadline)
            if ext is None:
                return None
            length = struct.unpack("!Q", ext)[0]
        mask = self._read_exact(4, deadline) if (b1 & 0x80) else b""
        if mask is None:
            return None
        payload = self._read_exact(length, deadline) if length else b""
        if payload is None:
            return None
        if mask:
            payload = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
        return fin, opcode, payload

    def _read_exact(self, n: int, deadline: float) -> bytes | None:
        buf = bytearray()
        while len(buf) < n:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return None
            self._sock.settimeout(remaining)
            try:
                chunk = self._sock.recv(n - len(buf))
            except TimeoutError:
                return None
            if not chunk:
                raise ConnectionError("peer closed the WebSocket connection")
            buf += chunk
        return bytes(buf)

    def close(self) -> None:
        try:
            self._sock.close()
        except OSError:
            pass


# --------------------------------------------------------------------------- agent


class Agent:
    """Wires capture (shell + workspace watcher) to the framed event stream."""

    def __init__(self, config: AgentConfig) -> None:
        self.config = config
        self._spool = Spool(config.spool_path)
        self._pending: deque[str] = deque()
        self._seq = 0
        self._emit_lock = threading.Lock()  # serializes seq + spool + flush
        self._conn_lock = threading.Lock()  # guards _conn swaps
        self._conn: WsConnection | None = None
        self._last_sent: str | None = None  # one-line replay margin, see below
        self._stop = threading.Event()
        self._threads: list[threading.Thread] = []
        self._baseline: dict[str, tuple[int, int, str | None]] = {}
        self._resume_from_spool()

    # -- durability ------------------------------------------------------
    def _resume_from_spool(self) -> None:
        highest = -1
        for line in self._spool.read_all():
            self._pending.append(line)
            try:
                seq = int(json.loads(line).get("seq", -1))
            except (ValueError, AttributeError):
                continue
            highest = max(highest, seq)
        self._seq = highest + 1

    # -- event construction ---------------------------------------------
    def _wire_frame(self, spooled_line: str) -> str:
        """Spool format -> wire format: strip URL-owned identity fields.

        The spool keeps full events (local durability + restart recovery).
        The ingest endpoint binds identity at the WS handshake (query params)
        and rejects frames carrying learner_id/task_id (`extra="forbid"`
        anti-spoofing), so the wire frame carries only seq/kind/payload/ts.
        """
        import json as _json

        full = _json.loads(spooled_line)
        wire = {
            k: full[k]
            for k in ("seq", "kind", "payload", "ts")
        }
        if full.get("sandbox_id"):
            wire["sandbox_id"] = full["sandbox_id"]
        return _json.dumps(wire)

    def _next_event(self, kind: str, payload: dict[str, Any]) -> dict[str, Any]:
        if kind not in _EVENT_KINDS:
            raise ValueError(f"unknown event kind: {kind!r}")
        event = {
            "learner_id": self.config.learner_id,
            "task_id": self.config.task_id,
            "seq": self._seq,
            "kind": kind,
            "payload": payload,
            "ts": datetime.now(UTC).isoformat(),
            "sandbox_id": self.config.sandbox_id,
        }
        self._seq += 1
        return event

    # -- emission / flush (D-026) ----------------------------------------
    def emit(self, kind: str, payload: dict[str, Any]) -> dict[str, Any]:
        """Spool-then-send. Never blocks on reconnect; loss is impossible."""
        with self._emit_lock:
            line = json.dumps(self._next_event(kind, payload))
            self._spool.append(line)  # durable BEFORE any send attempt
            self._pending.append(line)
            self._flush_locked()
            return json.loads(line)

    def _flush_locked(self) -> None:
        conn = self._current_conn()
        while self._pending and conn is not None:
            line = self._pending[0]
            try:
                conn.send_text(self._wire_frame(line))
            except (ConnectionError, OSError):
                self._drop_conn()
                return
            self._pending.popleft()
            self._last_sent = line  # kept until a later send proves delivery
        if not self._pending and self._last_sent is not None:
            # Compact, but retain the most recently sent line: a send into a
            # silently-dead socket "succeeds" once at TCP level, so the last
            # line is only confirmed-sent once a later write works. Retention
            # is cheap; the server dedups on (learner, task, seq).
            self._spool.rewrite([self._last_sent])

    def replay_margin(self) -> None:
        """Requeue the last-sent line after a detected disconnect."""
        with self._emit_lock:
            if self._last_sent is not None and (
                not self._pending or self._pending[0] != self._last_sent
            ):
                self._pending.appendleft(self._last_sent)
                self._spool.rewrite(list(self._pending))
            self._last_sent = None

    # -- connection supervision ------------------------------------------
    def _current_conn(self) -> WsConnection | None:
        with self._conn_lock:
            return self._conn

    def _set_conn(self, conn: WsConnection | None) -> None:
        with self._conn_lock:
            self._conn = conn

    def _drop_conn(self) -> None:
        conn = self._current_conn()
        self._set_conn(None)
        if conn is not None:
            conn.close()
        self.replay_margin()

    def is_connected(self) -> bool:
        return self._current_conn() is not None

    def wait_connected(self, timeout_s: float) -> bool:
        return self._wait_for(lambda: self.is_connected(), timeout_s)

    def wait_disconnected(self, timeout_s: float) -> bool:
        return self._wait_for(lambda: not self.is_connected(), timeout_s)

    def _wait_for(self, pred: Any, timeout_s: float) -> bool:
        deadline = time.monotonic() + timeout_s
        while time.monotonic() < deadline:
            if pred():
                return True
            time.sleep(0.02)
        return pred()

    def _supervisor_loop(self) -> None:
        """Maintain the WS connection: connect, flush backlog, read, backoff."""
        backoff = self.config.backoff_base_s
        while not self._stop.is_set():
            if self._current_conn() is None:
                try:
                    conn = WsConnection.connect(self.config.ingest_url)
                except (ConnectionError, OSError, ValueError, TimeoutError):
                    self._stop.wait(backoff)
                    backoff = min(self.config.backoff_max_s, backoff * 2)
                    continue
                self._set_conn(conn)
                self._last_sent = None
                backoff = self.config.backoff_base_s
                with self._emit_lock:  # ordered against concurrent emit()s
                    self._flush_locked()
            else:
                conn = self._current_conn()
            if conn is None:
                continue
            try:
                frame = conn.recv_message(timeout_s=1.0)
            except (ConnectionError, OSError):
                self._drop_conn()
                continue
            if frame is None:
                continue
            opcode, _payload = frame
            if opcode == 0x8:  # server close frame
                self._drop_conn()

    # -- workspace watcher ------------------------------------------------
    def _snapshot_workspace(self) -> dict[str, tuple[int, int, str | None]]:
        """Map rel path -> (mtime_ns, size, text-or-None-if-too-large)."""
        snap: dict[str, tuple[int, int, str | None]] = {}
        root = self.config.workspace
        if not root.is_dir():
            return snap
        for dirpath, dirnames, filenames in os.walk(root):
            dirnames[:] = sorted(
                d for d in dirnames if not d.startswith(_AGENT_DIR_PREFIX)
            )
            for name in sorted(filenames):
                if name.startswith(_AGENT_DIR_PREFIX):
                    continue
                path = Path(dirpath) / name
                try:
                    st = path.stat()
                except OSError:
                    continue
                rel = path.relative_to(root).as_posix()
                text: str | None = None
                if st.st_size <= _MAX_DIFF_BYTES:
                    try:
                        text = path.read_text(encoding="utf-8", errors="replace")
                    except OSError:
                        pass
                snap[rel] = (st.st_mtime_ns, st.st_size, text)
        return snap

    def _file_diff_payload(self, rel: str, change: str, old: str | None, new: str | None) -> dict:
        payload: dict[str, Any] = {"path": rel, "change": change}
        if old is None and new is None:
            payload["truncated"] = True
            return payload
        diff = "".join(
            difflib.unified_diff(
                (old or "").splitlines(keepends=True),
                (new or "").splitlines(keepends=True),
                fromfile=f"a/{rel}",
                tofile=f"b/{rel}",
            )
        )
        payload["diff"] = diff
        payload["size"] = len(new or "")
        return payload

    def _watcher_loop(self) -> None:
        # Baseline is taken in start() before it returns, so any change made
        # after start() completes is guaranteed to be observed.
        baseline = self._baseline
        last_heartbeat = time.monotonic()
        while not self._stop.wait(self.config.poll_interval_s):
            current = self._snapshot_workspace()
            for rel in sorted(current.keys() | baseline.keys()):
                if rel not in baseline and rel in current:
                    self.emit(
                        "file_diff",
                        self._file_diff_payload(rel, "created", None, current[rel][2]),
                    )
                elif rel in baseline and rel not in current:
                    self.emit(
                        "file_diff",
                        self._file_diff_payload(rel, "deleted", baseline[rel][2], None),
                    )
                else:
                    old_stat, new_stat = baseline[rel], current[rel]
                    if old_stat[:2] != new_stat[:2] and old_stat[2] != new_stat[2]:
                        self.emit(
                            "file_diff",
                            self._file_diff_payload(
                                rel, "modified", old_stat[2], new_stat[2]
                            ),
                        )
            baseline = current
            if time.monotonic() - last_heartbeat >= self.config.activity_interval_s:
                self.emit("activity", {"state": "idle", "spooled": len(self._pending)})
                last_heartbeat = time.monotonic()

    # -- shell wrapper ------------------------------------------------------
    @staticmethod
    def _is_test_command(cmd: str) -> bool:
        return "test" in cmd.lower()

    def run_command(self, line: str) -> dict[str, Any] | None:
        """Run one REPL line; emits stdin/command/stdout/run|test_result."""
        line = line.strip()
        if not line:
            return None
        self.emit("stdin", {"line": line})
        self.emit("activity", {"state": "command", "spooled": len(self._pending)})
        self.emit("command", {"cmd": line})
        started = time.monotonic()
        timed_out = False
        exit_code: int | None = None
        out: str | bytes = ""
        err: str | bytes = ""
        try:
            proc = subprocess.run(
                ["sh", "-c", line],
                cwd=self.config.workspace,
                capture_output=True,
                timeout=self.config.command_timeout_s,
                text=True,
                errors="replace",
            )
            exit_code, out, err = proc.returncode, proc.stdout, proc.stderr
        except subprocess.TimeoutExpired as exc:
            timed_out = True
            # TimeoutExpired output attrs are always bytes (even in text mode).
            out = exc.stdout or b""
            err = exc.stderr or b""
        duration = time.monotonic() - started
        for stream, data in (("stdout", out), ("stderr", err)):
            if isinstance(data, bytes):
                data = data.decode(errors="replace")
            if data:
                self.emit("stdout", {"stream": stream, "data": data[-_MAX_OUTPUT_CHARS:]})
        kind = "test_result" if self._is_test_command(line) else "run_result"
        result = self.emit(
            kind,
            {
                "cmd": line,
                "exit_code": exit_code,
                "duration_s": round(duration, 6),
                "timed_out": timed_out,
            },
        )
        self.emit("activity", {"state": "idle", "spooled": len(self._pending)})
        return result

    # -- lifecycle ----------------------------------------------------------
    def start(self) -> None:
        self.config.workspace.mkdir(parents=True, exist_ok=True)
        self._baseline = self._snapshot_workspace()
        self.emit("activity", {"state": "starting", "spooled": len(self._pending)})
        self._threads = [
            threading.Thread(target=self._supervisor_loop, daemon=True, name="nc-ws"),
            threading.Thread(target=self._watcher_loop, daemon=True, name="nc-watch"),
        ]
        for thread in self._threads:
            thread.start()

    def stop(self) -> None:
        if self._stop.is_set():
            return
        try:
            self.emit("activity", {"state": "stopped", "spooled": len(self._pending)})
        finally:
            self._stop.set()
            self._drop_conn()
            for thread in self._threads:
                thread.join(timeout=3)
            for thread in self._threads:
                thread.join(timeout=3)
        with self._emit_lock:
            self._spool.rewrite(list(self._pending))


def main() -> int:
    try:
        config = AgentConfig.from_env()
    except ValueError as exc:
        print(f"sandbox-agent: {exc}", file=sys.stderr)
        return 2
    agent = Agent(config)
    agent.start()
    try:
        # REPL mode: each stdin line is executed and reported (interactive use).
        # Daemon mode: when stdin is closed/absent (the sandbox backend spawns
        # the agent with stdin=DEVNULL), keep streaming workspace diffs +
        # activity until SIGTERM/SIGINT so the agent's lifecycle is tied to
        # the sandbox (destroy() reaps it) rather than to stdin EOF.
        if sys.stdin is None or sys.stdin.closed:  # pragma: no cover - defensive
            agent._stop.wait()  # noqa: SLF001 - daemon block
        else:
            line = sys.stdin.readline()
            while line:
                agent.run_command(line)
                line = sys.stdin.readline()
            if not agent._stop.is_set() and not sys.stdin.isatty():  # noqa: SLF001
                # EOF on a pipe (DEVNULL): daemonize — watch + stream until killed.
                import signal

                signal.signal(signal.SIGTERM, lambda *_: agent._stop.set())  # noqa: SLF001
                agent._stop.wait()  # noqa: SLF001
    except KeyboardInterrupt:
        pass
    finally:
        agent.stop()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())