#!/usr/bin/env python3
"""Binks swarm proxy - one OpenAI endpoint in front of brain + worker pool.

Chadeous-Agent-OS discovers providers at LM_STUDIO_URL (default :1234) and
speaks OpenAI protocol. This proxy presents the WHOLE swarm on that one URL:

    GET  /v1/models            -> binks-run8 (brain) + binks-mini-1b7 (worker pool)
    POST /v1/chat/completions  -> routes by the "model" field:
         contains "mini"  -> worker pool (round-robin across backends)
         anything else    -> brain backend
    GET  /health               -> proxy + per-backend status
    GET  /v1/activity          -> in-memory request feed for the swarm dashboard
                                  (optional ?since_id=N for incremental polling)

Backends:
  brain   k-trainer openai_server.py (HRR full stack, GPU, e.g. :1240)
  worker  llama-server GGUF instances (CPU, e.g. :8501, :8502, ...)

Streaming responses are tee'd as they pass through: bytes go straight to the
client AND into a small buffer so /v1/activity can show what each agent said
(its "thoughts"), its <|state:*|> token, latency and token usage. CORS is open
so the dashboard dev server can poll from another port. Stdlib only.

    python swarm_proxy.py --port 1234 \
        --brain http://127.0.0.1:1240 \
        --worker http://127.0.0.1:8501 --worker http://127.0.0.1:8502
"""
from __future__ import annotations

import argparse, json, re, threading, time, urllib.request, urllib.error
from collections import deque
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

# ---- Executive Cognitive Directive for calibrated reasoning ----
EXECUTIVE_DIRECTIVE = (
    "You are K-Mini (4.02B), a sovereign local intelligence and developer agent created by Osakra Research.\n"
    "Protocol:\n"
    "1. Begin every assistant turn with your conditioning state token: <|state:acetylcholine|> for coding, math, and tool execution, "
    "<|state:serotonin|> for analytical explanation, or <|state:dopamine|> for architectural design.\n"
    "2. Solve each problem methodically step-by-step.\n"
    "3. Only if a factual trivia question contains an impossible historical/physical premise (e.g. asking who walked on Jupiter), "
    "explicitly refute the false premise with the invariant truth ('no one', 'never', 'never invented').\n"
    "4. For formal logic contrapositives, state the direct contrapositive in English; for analogies, state the single best target word.\n"
    "5. For math, logic, or factual questions, state 'Final Answer: <answer>' at the end."
)

ARGS = None
RR_LOCK = threading.Lock()
RR_IDX = [0]
HEALTHY_WORKERS = []

def _health_monitor_loop():
    global HEALTHY_WORKERS
    while True:
        healthy = []
        for w in ARGS.workers:
            if backend_health(w):
                healthy.append(w)
        HEALTHY_WORKERS = healthy
        time.sleep(3)

# ---- per-request forensic log (survives crashes; one line on start + end) ----
REQLOG_LOCK = threading.Lock()
REQLOG_PATH = [None]       # set by --reqlog; None disables


def _req_log(line: str) -> None:
    if not REQLOG_PATH[0]:
        return
    ts = time.strftime("%Y-%m-%dT%H:%M:%S")
    with REQLOG_LOCK:
        try:
            with open(REQLOG_PATH[0], "a", encoding="utf-8") as f:
                f.write(f"{ts} {line}\n")
        except Exception:
            pass


# ---- activity feed (in-memory ring buffer, read by the swarm dashboard) ----
ACT_LOCK = threading.Lock()
ACT_SEQ = [0]
ACTIVITY: deque = deque(maxlen=400)
OUTPUT_CAP = 6000          # chars of assistant output kept per event
TEE_CAP = 256 * 1024       # bytes of upstream response teed per request
STATE_RE = re.compile(r"<\|state:([a-zA-Z0-9_+\-]+)\|>")


def _act_add(ev: dict) -> int:
    with ACT_LOCK:
        ACT_SEQ[0] += 1
        ev["id"] = ACT_SEQ[0]
        ACTIVITY.append(ev)
        return ev["id"]


def _act_update(eid: int, **fields) -> None:
    with ACT_LOCK:
        for ev in reversed(ACTIVITY):
            if ev["id"] == eid:
                ev.update(fields)
                return


def _prompt_preview(payload: dict) -> tuple[str, int]:
    """Last user message, whitespace-collapsed and capped, + message count.
    Chadeous packs [SYSTEM DIRECTIVE]+env+task into one first user message,
    so a head preview made every iteration row look identical ("duplicate
    prompts" confusion 2026-08-09). For directive-packed messages preview the
    TAIL, where the actual task lives."""
    msgs = payload.get("messages") or []
    text = ""
    for m in reversed(msgs):
        if m.get("role") == "user":
            c = m.get("content")
            if isinstance(c, str):
                text = c
            elif isinstance(c, list):
                text = " ".join(p.get("text", "") for p in c if isinstance(p, dict))
            break
    text = " ".join(text.split())
    if text.startswith("[SYSTEM DIRECTIVE]") and len(text) > 400:
        return "…" + text[-380:], len(msgs)
    return text[:400], len(msgs)


def _parse_sse_text(buf: bytes) -> tuple[str, dict]:
    """Concatenate assistant content deltas + grab last usage from teed SSE."""
    parts, usage = [], {}
    for line in buf.decode("utf-8", "ignore").splitlines():
        line = line.strip()
        if not line.startswith("data:"):
            continue
        data = line[5:].strip()
        if not data or data == "[DONE]":
            continue
        try:
            obj = json.loads(data)
        except Exception:
            continue
        if obj.get("usage"):
            usage = obj["usage"]
        for ch in obj.get("choices") or []:
            delta = ch.get("delta") or {}
            if isinstance(delta.get("content"), str):
                parts.append(delta["content"])
    return "".join(parts), usage


def _parse_json_text(buf: bytes) -> tuple[str, dict]:
    try:
        obj = json.loads(buf.decode("utf-8", "ignore"))
    except Exception:
        return "", {}
    usage = obj.get("usage") or {}
    parts = []
    for ch in obj.get("choices") or []:
        msg = ch.get("message") or {}
        if isinstance(msg.get("content"), str):
            parts.append(msg["content"])
    return "".join(parts), usage


def pick_backend(model_field: str) -> tuple[str, str]:
    """Return (base_url, advertised_id) for a requested model id."""
    m = (model_field or "").lower()
    if any(k in m for k in ("norn", "ncm", "titans", "concept")):
        ncm_url = getattr(ARGS, "ncm", "http://127.0.0.1:8092")
        return ncm_url, "norn-4b" if "norn" in m else "titans-ncm"
    if any(k in m for k in ("k-mini", "kmini", "brain")):
        return ARGS.brain, "k-mini"
    if any(k in m for k in ("binks", "worker", "swarm", "mini")) and HEALTHY_WORKERS:
        with RR_LOCK:
            url = HEALTHY_WORKERS[RR_IDX[0] % len(HEALTHY_WORKERS)]
            RR_IDX[0] += 1
        return url, "binks"
    return ARGS.brain, "k-mini"


def backend_health(url: str) -> bool:
    try:
        urllib.request.urlopen(url.rstrip("/") + "/health", timeout=2)
        return True
    except Exception:
        return False


class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt, *a):
        pass

    def _cors(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization, Accept")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")

    def _send(self, code, body, ctype="application/json"):
        data = body.encode() if isinstance(body, str) else body
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Connection", "close")
        self._cors()
        self.end_headers()
        self.wfile.write(data)

    def do_OPTIONS(self):
        self.send_response(204)
        self._cors()
        self.send_header("Content-Length", "0")
        self.send_header("Connection", "close")
        self.end_headers()

    def do_GET(self):
        path = self.path.split("?", 1)[0].rstrip("/")
        if path in ("/v1/models", "/models"):
            ncm_url = getattr(ARGS, "ncm", "http://127.0.0.1:8092")
            models = [
                {"id": "k-mini", "object": "model", "owned_by": "k-mini-brain",
                 "backend": ARGS.brain, "up": backend_health(ARGS.brain)},
                {"id": "k-mini2.5", "object": "model", "owned_by": "k-mini-brain",
                 "backend": ARGS.brain, "up": backend_health(ARGS.brain)},
                {"id": "norn-4b", "object": "model", "owned_by": "osakra-research",
                 "backend": ncm_url, "up": backend_health(ncm_url)},
                {"id": "titans-ncm", "object": "model", "owned_by": "neuro-symbolic-ncm",
                 "backend": ncm_url, "up": backend_health(ncm_url)},
            ]
            if ARGS.workers:
                models.append({
                    "id": "binks", "object": "model",
                    "owned_by": "binks-swarm-agent",
                    "pool_size": len(HEALTHY_WORKERS),
                    "up": len(HEALTHY_WORKERS) > 0})
                models.append({
                    "id": "binks-mini-1b7", "object": "model",
                    "owned_by": "binks-worker-pool",
                    "pool_size": len(HEALTHY_WORKERS),
                    "up": len(HEALTHY_WORKERS) > 0})
            self._send(200, json.dumps({"object": "list", "data": models}))
        elif path in ("/health", "/"):
            self._send(200, json.dumps({
                "ok": True,
                "now": time.time(),
                "brain": {"url": ARGS.brain, "up": backend_health(ARGS.brain)},
                "workers": [{"url": u, "up": u in HEALTHY_WORKERS} for u in ARGS.workers]}))
        elif path == "/v1/activity":
            since = 0
            if "?" in self.path:
                for part in self.path.split("?", 1)[1].split("&"):
                    if part.startswith("since_id="):
                        try:
                            since = int(part[9:])
                        except ValueError:
                            pass
            with ACT_LOCK:
                events = [dict(ev) for ev in ACTIVITY if ev["id"] > since]
            self._send(200, json.dumps({
                "ok": True, "now": time.time(),
                "seq": ACT_SEQ[0], "events": events}))
        else:
            self._send(404, json.dumps({"error": "not found"}))

    def do_POST(self):
        path = self.path.split("?", 1)[0].rstrip("/")
        if path not in ("/v1/chat/completions", "/chat/completions", "/v1/completions"):
            self._send(404, json.dumps({"error": "not found"}))
            return
        length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(length)
        try:
            payload = json.loads(body)
        except Exception:
            self._send(400, json.dumps({"error": "invalid JSON"}))
            return

        req_model = payload.get("model") or ""
        base, advert = pick_backend(req_model)
        payload["model"] = advert

        # Inject calibrated Executive Cognitive Directive if no explicit system message is supplied
        msgs = payload.get("messages")
        if msgs and isinstance(msgs, list) and advert not in ("titans-ncm", "norn-4b"):
            has_system = any(m.get("role") == "system" for m in msgs)
            if not has_system:
                payload["messages"] = [{"role": "system", "content": EXECUTIVE_DIRECTIVE}] + msgs
        preview, n_msgs = _prompt_preview(payload)
        role = "worker" if advert in ("binks", "binks-mini-1b7") else "brain"
        t0 = time.time()
        eid = _act_add({
            "ts_start": t0, "ts_end": None, "latency_s": None,
            "model_req": req_model, "model": advert, "role": role,
            "backend": base, "status": "running", "error": None,
            "prompt_preview": preview, "prompt_messages": n_msgs,
            "stream": bool(payload.get("stream")),
            "output": "", "state_token": None,
            "tokens_in": None, "tokens_out": None,
        })
        _req_log(f"id={eid} START model={advert} req={req_model!r} backend={base} "
                 f"stream={bool(payload.get('stream'))} msgs={n_msgs} prompt={preview[:160]!r}")

        upstream_path = "/v1/chat/completions" if "completions" in path else path
        url = base.rstrip("/") + upstream_path
        req = urllib.request.Request(
            url, data=json.dumps(payload).encode(),
            headers={"Content-Type": "application/json", "Accept": self.headers.get("Accept", "*/*")})
        try:
            # 660s: just past Chadeous's 600s client timeout, so a wedged upstream
            # fails the request instead of parking this thread (nearly) forever.
            resp = urllib.request.urlopen(req, timeout=660)
        except urllib.error.HTTPError as e:
            lat = round(time.time() - t0, 2)
            _act_update(eid, ts_end=time.time(), latency_s=lat,
                        status="error", error=f"backend HTTP {e.code}")
            _req_log(f"id={eid} END status=error latency={lat}s err=backend_http_{e.code}")
            self._send(e.code, e.read())
            return
        except Exception as e:
            lat = round(time.time() - t0, 2)
            _act_update(eid, ts_end=time.time(), latency_s=lat,
                        status="error", error=f"backend unreachable: {e}")
            _req_log(f"id={eid} END status=error latency={lat}s err=unreachable:{e}")
            self._send(502, json.dumps({"error": f"backend {base} unreachable: {e}"}))
            return

        # stream straight through (SSE or plain JSON alike), teeing a copy
        ctype = resp.headers.get("Content-Type", "")
        self.send_response(resp.status)
        for h, v in resp.headers.items():
            if h.lower() in ("content-type", "cache-control"):
                self.send_header(h, v)
        self.send_header("Connection", "close")
        self._cors()
        self.end_headers()

        tee = bytearray()
        client_gone = False
        try:
            while True:
                chunk = resp.read(4096)
                if not chunk:
                    break
                if len(tee) < TEE_CAP:
                    tee += chunk[: TEE_CAP - len(tee)]
                try:
                    self.wfile.write(chunk)
                    self.wfile.flush()
                except (BrokenPipeError, ConnectionResetError):
                    client_gone = True
                    break
        except (BrokenPipeError, ConnectionResetError):
            client_gone = True
        finally:
            resp.close()

        ts_end = time.time()
        if "text/event-stream" in ctype:
            text, usage = _parse_sse_text(bytes(tee))
        else:
            text, usage = _parse_json_text(bytes(tee))
        m = STATE_RE.search(text or "")
        status = "client_disconnect" if client_gone else "ok"
        _act_update(
            eid,
            ts_end=ts_end, latency_s=round(ts_end - t0, 2),
            status=status,
            output=(text or "")[-OUTPUT_CAP:],
            state_token=m.group(0) if m else None,
            tokens_in=usage.get("prompt_tokens"),
            tokens_out=usage.get("completion_tokens"),
        )
        _req_log(f"id={eid} END status={status} latency={round(ts_end - t0, 2)}s "
                 f"tokens_in={usage.get('prompt_tokens')} tokens_out={usage.get('completion_tokens')}")


def main():
    global ARGS
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=1234)
    ap.add_argument("--brain", default="http://127.0.0.1:8096")
    ap.add_argument("--ncm", default="http://127.0.0.1:8092")
    ap.add_argument("--worker", dest="workers", action="append", default=[])
    ap.add_argument("--reqlog", default=None,
                    help="append one START/END line per request to this file")
    ARGS = ap.parse_args()
    REQLOG_PATH[0] = ARGS.reqlog
    
    # Start health monitor
    t = threading.Thread(target=_health_monitor_loop, daemon=True)
    t.start()
    
    srv = ThreadingHTTPServer(("0.0.0.0", ARGS.port), Handler)
    print(f"[swarm-proxy] http://localhost:{ARGS.port}/v1", flush=True)
    print(f"[swarm-proxy] brain   -> {ARGS.brain}", flush=True)
    for w in ARGS.workers:
        print(f"[swarm-proxy] worker  -> {w}", flush=True)
    print("[swarm-proxy] activity feed on /v1/activity (dashboard polls this)", flush=True)
    if ARGS.reqlog:
        print(f"[swarm-proxy] request log -> {ARGS.reqlog}", flush=True)
    srv.serve_forever()


if __name__ == "__main__":
    main()