#!/usr/bin/env python3
"""engine.py -- THE ENGINE PROTOCOL: one running engine, its stream record, the SSE and NDJSON readers, the trap.

What this file is, for a reader who opened it cold (SPEC Δ6, the bench-client lane §3.6). The ollama-vs-vLLM bench
measures three engines with one client, one clock and one rule set: vLLM (engine_vllm.py), Ollama (engine_ollama.py)
and llama-server (engine_llamaserver.py, llama.cpp's server as Ollama bundles it). Each adapter is an `Engine`:

    start()            launch from servelines.argv_env, as a plain CHILD of the arm (no new session): the stop
                       guard's signal to the arm's process group reaches it
    ready(timeout_s)   seconds to ready; R-3 on an exit, a fatal log line or the timeout
    load()             Ollama: one unscored request that loads the model (its load_duration kept); others: no-op
    readback()         argv, env (whitelisted keys), the boot lines parsed, the internals (N-5), log offsets
    raw_body(S, ...)   the parity path's body (item 1): never an `ignore_eos` key
    chat_row_body(t)   the labelled chat row's body, with the engine's own thinking switch
    raw_path() / chat_path()
    decode_stream(resp, t0, rec)   the frames -> the stream record (below), on the client's clock
    tokenize_raw(text, add_special)   the IDs this engine evaluates for a raw string (R-P2)
    render_chat_ids(user_text)        the chat render as IDs (R-P1)
    visible_count(text)               the tokenizer count of visible text (the hidden-token gate's fallback)
    pieces(ids)                       each ID's piece, specials kept (the gate's exact method)
    server_clock(rec) / cached_tokens(rec) / prompt_count(rec)
    log_offset() / log_since(offset)  the log bytes a level spans (all-layers, slot lines)
    stop()             SIGINT -> SIGTERM -> SIGKILL, every descendant: seat.SeatProcess's own stop, inherited

THE STREAM RECORD is the seat kit's `stream_once` record generalised; a field keeps its name where its meaning is
unchanged (ttft_ms, t_first_token_s, t_finish_s, decode_tok_s, prompt_tokens, completion_tokens, finish_reason,
completion_sha256, completion_head, token_frames, status, fail_kind, fail_reasons, unscored_*), and it adds
generated_ids (when the engine sends them), prompt_ids_sha256 (parity.ids_sha of the prompt IDs the engine says it
evaluated: R-P2's own sha, so the wire's IDs match the parity record of the same string), cached_tokens,
server_clock{..., rule}, hidden_surplus / hidden_method, and path / engine / posture / nonce.

ONE TIMING RULE for all three engines (the kit's Δ5 rule, seatlib.RULES["engine_timing"]): t_first_token_s is the first
frame that carries a generated token, t_finish_s the frame carrying the last generated token, and decode_tok_s =
(completion_tokens - 1) / (t_finish_s - t_first_token_s). Each adapter's `decode_stream` says which frame that is.
"""
import hashlib
import http.client
import json
import os
import re
import signal
import subprocess
import time

import pen
import seat as S
import seatlib as L
import servelines as SL

STREAM_TIMEOUT_S = 900.0
BARRIER_TIMEOUT_S = 120.0
CALL_TIMEOUT_S = 60.0
STREAM_TIMINGS = ("ttft_ms", "decode_tok_s", "t_first_token_s", "t_finish_s")
UNSCORED_PREFIX = "unscored_"
ENGINES = ("vllm", "ollama", "llama-server")
_PY = re.compile(r"python[0-9.]*$")
#: Every CPU this process could use when the kit was imported -- before the client pins itself (client.pin_client,
#: N-12). An engine is launched on ALL of them: a child inherits its parent's affinity, and an engine squeezed onto the
#: client's two SMT siblings would measure the pin, not the engine.
ALL_CPUS = frozenset(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else None


def _release_pin():
    """preexec_fn of every engine launch: the child runs on ALL_CPUS whatever the client pinned itself to."""
    if ALL_CPUS:
        os.sched_setaffinity(0, ALL_CPUS)


# ------------------------------------------------------------------------------------------------ #
# processes
# ------------------------------------------------------------------------------------------------ #

def program_argv(argv):
    """An argv from the program token on: `python <script> args` (an entry-point script) reads as `<script> args`."""
    argv = list(argv or [])
    if len(argv) > 1 and _PY.search(os.path.basename(argv[0])):
        return argv[1:]
    return argv


def engine_of_argv(argv):
    """'vllm' | 'ollama' | 'llama-server' | None: the three engines as this kit launches them (the trap's check)."""
    if S.is_vllm_serve(argv):
        return "vllm"
    prog = program_argv(argv)
    if not prog:
        return None
    base = os.path.basename(prog[0])
    if base == "ollama" and prog[1:2] == ["serve"]:
        return "ollama"
    if base == "llama-server":
        return "llama-server"
    return None


#: A log stamp with no zone and no year (vLLM's `INFO MM-DD hh:mm:ss`): never written outward (the pen's UTC rule).
LOG_STAMP = re.compile(r"(?:\((?:EngineCore|APIServer)[^)]*\)\s*)?(?:(?:DEBUG|INFO|WARNING|ERROR|CRITICAL)\s+)?"
                       r"\b\d{2}-\d{2} \d{2}:\d{2}:\d{2}(?:\.\d+)?\s*")


def strip_log_stamp(line):
    """A log line kept as evidence without its zone-less stamp (correctness N-1): the message and its [file:line]."""
    return LOG_STAMP.sub("", str(line)) if line is not None else None


def proc_argv(pid):
    return S._argv(pid)


def proc_environ(pid):
    """/proc/<pid>/environ as a dict, or None when unreadable."""
    try:
        with open(f"/proc/{pid}/environ", "rb") as fh:
            raw = fh.read()
    except OSError:
        return None
    out = {}
    for item in raw.split(b"\0"):
        if b"=" in item:
            k, v = item.split(b"=", 1)
            out[k.decode("utf-8", "replace")] = v.decode("utf-8", "replace")
    return out


class EngineProcess(S.SeatProcess):
    """One engine's process: seat.SeatProcess's discipline (pids, alive, the stop), a different launch."""

    def __init__(self, name, argv, env, log_path, pid_path):  # noqa: D401 -- no super(): no seat, no serve line
        self.seat = {"id": name, "port": None}
        self.argv = list(argv)
        self.env = dict(env)
        self.log_path = log_path
        self.pid_path = pid_path
        self.proc = None
        self.t_start = None
        self.known_pids = set()
        self.last_stop = None

    def start(self):
        fh = pen.private_log(self.log_path)
        try:
            self.proc = subprocess.Popen(self.argv, env=self.env, stdout=fh, stderr=subprocess.STDOUT,
                                         stdin=subprocess.DEVNULL, preexec_fn=_release_pin)
        finally:
            fh.close()
        self.t_start = time.monotonic()
        self.known_pids.add(self.proc.pid)
        if self.pid_path:
            pen.write_private_text(self.pid_path, f"{self.proc.pid}\n")
        return self.proc.pid


def reap_engines(run_dir, wait_s=10.0, log=print):
    """run_arms.sh's trap: every engine an arm left behind (its pidfile under <run>/private/) is stopped -- but only a
    live process whose argv is one of the three engines as this kit launches them (engine_of_argv), so a recycled PID
    is never signalled -- with its descendants: SIGTERM, then SIGKILL after `wait_s`. Returns what was reaped."""
    import glob
    reaped = []
    for path in sorted(glob.glob(os.path.join(run_dir, "private", "*.pid")) + glob.glob(os.path.join(run_dir, "*.pid"))):
        try:
            pid = int(open(path, encoding="utf-8").read().strip())
        except (OSError, ValueError):
            pid = None
        if pid and S._alive(pid):
            kind = engine_of_argv(proc_argv(pid))
            if kind is None:
                log(f"trap: {os.path.basename(path)} names pid {pid}, alive but not an engine "
                    f"({' '.join(proc_argv(pid))[:120]!r}): left alone")
            else:
                family = [pid] + S.descendants(pid)
                log(f"trap: {kind} was left running (pid {pid}, {os.path.basename(path)}): SIGTERM, then SIGKILL")
                for sig in (signal.SIGTERM, signal.SIGKILL):
                    for p in family:
                        try:
                            os.kill(p, sig)
                        except (ProcessLookupError, PermissionError):
                            pass
                    end = time.monotonic() + wait_s
                    while time.monotonic() < end and any(S._alive(p) for p in family):
                        time.sleep(0.2)
                    if not any(S._alive(p) for p in family):
                        break
                reaped.append({"pid": pid, "engine": kind, "pidfile": os.path.basename(path),
                               "left_alive": [p for p in family if S._alive(p)]})
        try:
            os.unlink(path)
        except OSError:
            pass
    return reaped


# ------------------------------------------------------------------------------------------------ #
# HTTP
# ------------------------------------------------------------------------------------------------ #

def http_call(port, method, path, body=None, timeout=CALL_TIMEOUT_S, headers=None):
    """(status, parsed JSON or text) on loopback, or (None, error text). `headers` adds to the request's (arm A's
    Authorization)."""
    conn = http.client.HTTPConnection("127.0.0.1", int(port), timeout=timeout)
    try:
        data = None if body is None else json.dumps(body).encode("utf-8")
        h = {"Content-Type": "application/json"} if data else {}
        h.update(headers or {})
        conn.request(method, path, body=data, headers=h)
        r = conn.getresponse()
        text = r.read().decode("utf-8", "replace")
        try:
            return r.status, json.loads(text)
        except ValueError:
            return r.status, text
    except (OSError, http.client.HTTPException) as e:
        return None, f"{type(e).__name__}: {e}"
    finally:
        conn.close()


def sse_events(resp):
    """(t_monotonic, payload bytes) per `data:` line; ends at `[DONE]` (drained) or at the end of the body."""
    while True:
        line = resp.readline()
        if not line:
            return
        line = line.strip()
        if not line.startswith(b"data:"):
            continue
        payload = line[5:].strip()
        if payload == b"[DONE]":
            resp.read()
            return
        yield time.monotonic(), payload


def ndjson_events(resp):
    """(t_monotonic, line bytes) per non-empty NDJSON line."""
    while True:
        line = resp.readline()
        if not line:
            return
        line = line.strip()
        if line:
            yield time.monotonic(), line


# ------------------------------------------------------------------------------------------------ #
# the stream record
# ------------------------------------------------------------------------------------------------ #

def new_record(engine, posture, path, nonce_key=None, nonce_text=None):
    return {"engine": engine, "posture": posture, "path": path, "nonce": nonce_key, "nonce_text": nonce_text,
            "status": None, "fail_kind": None, "fail_reasons": [], "http_status": None, "error": None,
            "ttft_ms": None, "t_first_token_s": None, "t_finish_s": None, "t_end_s": None, "decode_tok_s": None,
            "prompt_tokens": None, "completion_tokens": None, "finish_reason": None, "completion_sha256": None,
            "completion_head": None, "completion_text": "", "token_frames": 0, "frames": 0, "bad_frames": 0,
            "generated_ids": None, "prompt_ids_sha256": None, "cached_tokens": None, "cached_why": None,
            "server_clock": None, "hidden_surplus": None, "hidden_method": None, "flags": []}


def finish_record(rec):
    """The figures every engine derives the same way once its frames are read (seatlib.RULES["engine_timing"])."""
    text = rec.get("completion_text") or ""
    rec["completion_sha256"] = hashlib.sha256(text.encode("utf-8")).hexdigest()
    rec["completion_head"] = text[:120]
    tf, tl, n = rec.get("t_first_token_s"), rec.get("t_finish_s"), rec.get("completion_tokens")
    if tf is not None:
        rec["ttft_ms"] = round(tf * 1000.0, 3)
    if n and n > 1 and tf is not None and tl is not None and tl > tf:
        rec["decode_tok_s"] = round((n - 1) / (tl - tf), 3)
    return rec


def unscore(rec):
    """A failed stream carries no speed and no first token anywhere: each is kept only as unscored_<name>."""
    for k in STREAM_TIMINGS:
        kept = UNSCORED_PREFIX + k
        if rec.get(k) is not None or kept not in rec:
            rec[kept] = rec.get(k)
        rec[k] = None
    return rec


# ------------------------------------------------------------------------------------------------ #
# the protocol
# ------------------------------------------------------------------------------------------------ #

class Engine:
    """One running engine instance. Subclasses fill in the engine's own wire (see the module docstring)."""

    key = None
    #: The frame this engine's decode_stream reads each timing from, in words (printed beside every row).
    frame_rule = None

    def __init__(self, posture_key, ctx, run_dir, log=print):
        self.posture = SL.posture(posture_key)
        assert self.posture["engine"] == self.key, (posture_key, self.key)
        self.ctx = dict(ctx)
        self.port = int(ctx["port"])
        self.argv, self.env = SL.argv_env(self.posture, self.ctx)
        #: the serving variables the arm's environment carried, which argv_env dropped (names only; the fold's I-3)
        self.env_dropped = (SL.inherited_serving_env(self.ctx.get("base_env"))
                            if self.posture["engine"] in ("vllm", "ollama") else None)
        self.run_dir = run_dir
        _rows, priv = pen.dirs(run_dir)
        stem = f"{self.key}-{self.posture['key']}"
        self.stem = stem
        self.log_path = os.path.join(priv, f"{stem}.serve.log")
        self.pid_path = os.path.join(priv, f"{stem}.pid")
        self.proc = EngineProcess(stem, self.argv, self.env, self.log_path, self.pid_path)
        self.log = log
        self._pieces = {}
        self.boot_s = None
        self.started_utc = None

    # ---- process ---------------------------------------------------------------------------- #

    def start(self):
        self.started_utc = L.utc_now()
        return self.proc.start()

    def pids(self):
        return self.proc.pids()

    def stop(self, wait_s=None):
        return self.proc.stop(wait_s)

    def alive(self):
        return self.proc.alive()

    def log_offset(self):
        try:
            return os.path.getsize(self.log_path)
        except OSError:
            return 0

    def log_since(self, offset=0, end=None):
        try:
            with open(self.log_path, "rb") as fh:
                fh.seek(int(offset or 0))
                data = fh.read() if end is None else fh.read(max(0, int(end) - int(offset or 0)))
            return data.decode("utf-8", "replace")
        except OSError:
            return ""

    def runner_log_since(self, offset=0, end=None):
        return self.log_since(offset, end)

    def wait_ready(self, timeout_s, probe, fatal=None):
        """Seconds from start until `probe()` is true; R-3 on an exit, a fatal log line or the timeout."""
        t0 = self.proc.t_start
        deadline = t0 + timeout_s
        while True:
            self.proc.pids()
            if not self.proc.alive():
                raise L.Refused("R-3", f"{self.posture['key']} exited rc {self.proc.proc.returncode} before it was "
                                       f"ready on port {self.port}", {"log_tail": self.log_since()[-2000:]})
            if fatal is not None:
                m = fatal.search(self.log_since())
                if m:
                    raise L.Refused("R-3", f"{self.posture['key']}'s log printed {m.group(0)!r}: it did not load whole",
                                    {"log_tail": self.log_since()[-2000:]})
            if probe():
                self.boot_s = round(time.monotonic() - t0, 2)
                return self.boot_s
            if time.monotonic() > deadline:
                raise L.Refused("R-3", f"{self.posture['key']} was not ready within {timeout_s:g} s on port "
                                       f"{self.port}", {"log_tail": self.log_since()[-2000:]})
            time.sleep(0.2)

    def load(self):
        return {"load": "n/a", "why": f"{self.key} loads its model at boot"}

    #: True when a retokenize-method surplus can be replayed with logprobs (client.replay_hidden): Ollama only.
    replays_hidden = False

    def replay_exact_surplus(self, rec):
        return None

    def env_serving(self):
        """The serving variables of the engine's OWN /proc environ (what it really started with), or None unread."""
        env = proc_environ(self.proc.proc.pid) if self.proc.proc is not None else None
        return None if env is None else {k: v for k, v in env.items() if SL.is_serving_var(k)}

    # ---- calls ------------------------------------------------------------------------------ #

    def call(self, method, path, body=None, port=None, timeout=CALL_TIMEOUT_S):
        return http_call(port or self.port, method, path, body, timeout, headers=self.auth_headers(port))

    def auth_headers(self, port=None):
        """The Authorization a call to this engine's own port carries: arm A's canary key when its posture passes one
        (ctx["api_key"]); none otherwise, and never to another port (Ollama's runner)."""
        key = self.ctx.get("api_key")
        return {"Authorization": f"Bearer {key}"} if key and (port is None or port == self.port) else {}

    def pieces(self, ids):
        """Each ID's piece, specials kept, one call per ID not yet seen (the gate's exact method)."""
        out = []
        for i in ids:
            if i not in self._pieces:
                self._pieces[i] = self.detokenize([i])
            out.append(self._pieces[i])
        return out

    # ---- what every subclass supplies ------------------------------------------------------ #

    def raw_path(self):
        raise NotImplementedError

    def chat_path(self):
        raise NotImplementedError

    def raw_body(self, rendered, *, max_tokens, **road):
        raise NotImplementedError

    def chat_row_body(self, user_text, *, max_tokens):
        raise NotImplementedError

    def decode_stream(self, resp, t0, rec):
        raise NotImplementedError

    def tokenize_raw(self, text, *, add_special):
        raise NotImplementedError

    def render_chat_ids(self, user_text):
        raise NotImplementedError

    def visible_count(self, text):
        raise NotImplementedError

    def detokenize(self, ids):
        raise NotImplementedError

    def server_clock(self, rec):
        raise NotImplementedError

    def cached_tokens(self, rec):
        return rec.get("cached_tokens"), rec.get("cached_why")

    def prompt_count(self, rec):
        return rec.get("prompt_tokens")

    def readback(self):
        raise NotImplementedError


def decode_chat_sse(resp, t0, rec):
    """An OpenAI /v1/chat/completions stream (every chat row, every engine): the kit's role-frame rule (sweep.stream_once,
    SPEC Δ5) -- the role frame is not a token; each later choice frame is; the finish frame is the last token's."""
    text = []
    rec.setdefault("thinking_seen", False)
    for t, payload in sse_events(resp):
        rec["frames"] += 1
        try:
            obj = json.loads(payload)
        except ValueError:
            rec["bad_frames"] += 1
            continue
        if obj.get("error"):
            rec["error"] = str(obj["error"])[:300]
            continue
        for ch in (obj.get("choices") or [])[:1]:
            # vLLM under return_token_ids: each chunk's own IDs (the chat row's exact hidden-token method: fold N-2)
            if ch.get("token_ids"):
                rec["generated_ids"] = (rec.get("generated_ids") or []) + [int(x) for x in ch["token_ids"]]
        if obj.get("usage"):
            u = obj["usage"]
            rec["usage"] = u
            rec["prompt_tokens"], rec["completion_tokens"] = u.get("prompt_tokens"), u.get("completion_tokens")
            det = u.get("prompt_tokens_details") or {}
            if det.get("cached_tokens") is not None:
                rec["cached_tokens"] = det["cached_tokens"]
        for ch in (obj.get("choices") or [])[:1]:
            delta = ch.get("delta") or {}
            c = delta.get("content")
            rc = delta.get("reasoning_content") or delta.get("reasoning")
            if "role" in delta and not (c or rc) and not ch.get("finish_reason"):
                continue
            rec["token_frames"] += 1
            if rec["t_first_token_s"] is None:
                rec["t_first_token_s"] = t - t0
            if rc:
                rec["thinking_seen"] = True
            if c:
                text.append(c)
            if ch.get("finish_reason"):
                rec["finish_reason"] = ch["finish_reason"]
                rec["t_finish_s"] = t - t0
    rec["completion_text"] = "".join(text)
    rec["t_end_s"] = time.monotonic() - t0
    if rec.get("finish_reason") is None and not rec.get("error"):
        rec["error"] = "the stream ended without a finish_reason"
    return rec