#!/usr/bin/env python
"""Opal Router — the architecture-agnostic "merge" of un-mergeable models.

Different architectures (Llama / gpt-oss / Phi / Gemma / Liquid) cannot be
weight-merged. So Opal combines them at the SYSTEM level: one OpenAI-compatible
endpoint that classifies each request and routes it to the US specialist best
suited to it — Llama for tool-calling, gpt-oss for hard reasoning, the Opal-8B
merge for code/security, Phi for general, a VLM for images.

Zero dependencies (stdlib only). Runs locally over Ollama. No training.

  python opal_router.py            # serves http://localhost:11500/v1
  # point any OpenAI client at it with model="opal"

Env: OLLAMA=http://localhost:11434  PORT=11500
"""
import json, os, re, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

ROOT = os.path.dirname(os.path.abspath(__file__))
OLLAMA = os.environ.get("OLLAMA", "http://localhost:11434")
PORT = int(os.environ.get("PORT", "11500"))
ROUTES = json.load(open(os.path.join(ROOT, "configs", "routes.json")))["routes"]

CODE_RE = re.compile(r"```|\b(def |class |function|refactor|debug|implement|"
                     r"compile|stack ?trace|regex|SQL|python|javascript|rust|"
                     r"golang|typescript|unit test|api endpoint)\b", re.I)
REASON_RE = re.compile(r"\b(prove|proof|step by step|reason|derive|solve|"
                       r"why does|how many|calculate|logic|theorem|explain why|"
                       r"think through|chain of thought)\b", re.I)
SEC_RE = re.compile(r"\b(vulnerabilit|exploit|CVE|malware|detection|sigma rule|"
                    r"threat model|beacon|firmware|reverse engineer|yara)\b", re.I)


def available_models():
    try:
        with urllib.request.urlopen(f"{OLLAMA}/api/tags", timeout=10) as r:
            return {m["name"] for m in json.loads(r.read()).get("models", [])}
    except Exception:
        return set()


def classify(body):
    """Return a capability key for this request."""
    if body.get("tools") or body.get("functions"):
        return "tool"
    msgs = body.get("messages", [])
    # vision: any image content block
    for m in msgs:
        c = m.get("content")
        if isinstance(c, list) and any(
                isinstance(b, dict) and b.get("type") in ("image_url", "image")
                for b in c):
            return "vision"
    text = " ".join(m.get("content", "") for m in msgs
                    if isinstance(m.get("content"), str))[-4000:]
    if SEC_RE.search(text) or CODE_RE.search(text):
        return "code"
    if REASON_RE.search(text):
        return "reasoning"
    return "general"


def pick_model(cap, avail):
    want = ROUTES.get(cap, ROUTES["default"])
    if want in avail or f"{want}:latest" in avail:
        return want
    return ROUTES["default"]


def forward(body):
    data = json.dumps(body).encode()
    req = urllib.request.Request(f"{OLLAMA}/v1/chat/completions", data=data,
                                 headers={"Content-Type": "application/json",
                                          "Authorization": "Bearer ollama"})
    with urllib.request.urlopen(req, timeout=900) as r:
        return r.read()


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

    def _send(self, code, payload, ctype="application/json"):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.end_headers()
        self.wfile.write(payload if isinstance(payload, bytes) else payload.encode())

    def do_GET(self):
        if self.path.rstrip("/").endswith("/models"):
            self._send(200, json.dumps({"object": "list", "data": [
                {"id": "opal", "object": "model", "owned_by": "cognis-digital"}]}))
        else:
            self._send(200, json.dumps({"status": "ok", "routes": ROUTES}))

    def do_POST(self):
        n = int(self.headers.get("Content-Length", 0))
        try:
            body = json.loads(self.rfile.read(n) or b"{}")
        except Exception:
            return self._send(400, json.dumps({"error": "bad json"}))
        avail = available_models()
        cap = classify(body)
        model = pick_model(cap, avail)
        body["model"] = model
        body["stream"] = False  # router does non-streaming for simplicity
        print(f"[opal] route={cap:9} -> {model}", flush=True)
        try:
            out = forward(body)
        except Exception as e:
            return self._send(502, json.dumps({"error": str(e)}))
        # tag which specialist answered
        try:
            j = json.loads(out)
            j.setdefault("x_opal", {})["route"] = cap
            j["x_opal"]["model"] = model
            out = json.dumps(j).encode()
        except Exception:
            pass
        self._send(200, out)


if __name__ == "__main__":
    print(f"Opal Router on :{PORT} -> Ollama {OLLAMA}")
    print("routes:", json.dumps(ROUTES))
    print("available:", sorted(available_models()))
    ThreadingHTTPServer(("0.0.0.0", PORT), Handler).serve_forever()