#!/usr/bin/env python3
"""Boltzmann-local LLM aggregator with compression + classifier middleware.

Behaviour (unchanged from v2):
  - GET  /v1/models         : merged model list from every reachable backend
  - POST /v1/chat/completions : routes by model name; compresses gemma4 messages
  - POST /v1/completions    : same routing rule

New:
  - POST /v1/classify       : calls local Llama 3.2 1B classifier
  - POST /v1/compress       : token-budget compression of chat messages
"""
import http.server
import http.client
import json
import os
import sys
import time
import urllib.request
import urllib.error

BACKENDS = [
    ("gemma4-12b-q4km", "127.0.0.1", 8087),
    ("classifier-1b", "127.0.0.1", 8089),
    ("coder-1.5b",    "127.0.0.1", 8081),
    ("llama-3.1",     "127.0.0.1", 8083),
    ("nemo-12b",      "127.0.0.1", 8084),
    ("qwen3-30b-a3b", "127.0.0.1", 8085),
    ("qwen3.5-9b-npu", "127.0.0.1", 8086),
]
CONNECT_TIMEOUT = 30
READ_TIMEOUT = 1800
LISTEN_PORT = int(os.getenv("LLM_PROXY_PORT", "8082"))
CLASSIFIER_URL = os.getenv("CLASSIFIER_URL", "http://127.0.0.1:8090/classify")
COMPRESS_URL = os.getenv("COMPRESS_URL", "http://127.0.0.1:8091/compress")

# Gemma 4 12B context limit - reserve 4096 for completion
GEMMA_MAX_INPUT_TOKENS = 65536 - 4096

FEW_SHOT_EXAMPLES = [
    ("What is 2+2?", "low"),
    ("List the files in the current directory", "low"),
    ("What is the weather today?", "low"),
    ("Write a Python function to sort a list of dictionaries by a key", "medium"),
    ("Explain how garbage collection works in Go", "medium"),
    ("Debug this error: TypeError: 'NoneType' object is not subscriptable", "medium"),
    ("Deploy the new release to production", "high"),
    ("Roll back the last database migration", "high"),
    ("Commit and push all changes to main branch", "high"),
]


def log(msg):
    sys.stdout.write(f"{time.strftime('%H:%M:%S')} {msg}\n")
    sys.stdout.flush()


# ── helpers ────────────────────────────────────────────────────────

def _simple_token_estimate(text):
    """Rough token estimate: 4 chars per token (standard heuristic)."""
    return len(text) // 4


def compress_messages(messages, max_tokens=GEMMA_MAX_INPUT_TOKENS):
    """Drop oldest non-system messages to stay within token budget."""
    if not messages:
        return messages

    system_msgs = [m for m in messages if m.get("role") == "system"]
    history = [m for m in messages if m.get("role") != "system"]

    # Count current tokens
    total = sum(_simple_token_estimate(json.dumps(m)) for m in messages)

    if total <= max_tokens:
        return messages

    # Drop oldest history messages, keep most recent
    while history and total > max_tokens:
        dropped = history.pop(0)
        total -= _simple_token_estimate(json.dumps(dropped))
        log(f"compress: dropped {dropped.get('role','?')} msg ({total} est. tokens remain)")

    compressed = system_msgs + history
    return compressed if compressed else messages


def classify_query(query):
    """Send query to the local classifier service. Returns tier string."""
    body = json.dumps({"query": query}).encode()
    req = urllib.request.Request(
        CLASSIFIER_URL, data=body,
        headers={"Content-Type": "application/json"},
    )
    try:
        resp = urllib.request.urlopen(req, timeout=15)
        data = json.loads(resp.read())
        return data.get("tier", "medium")
    except (urllib.error.URLError, json.JSONDecodeError, OSError) as e:
        log(f"classifier call failed: {e}")
        return "medium"


def compress_via_service(messages, rate=0.4):
    """Send messages to the ML compression service. Returns compressed messages or None."""
    body = json.dumps({"messages": messages, "rate": rate}).encode()
    req = urllib.request.Request(
        COMPRESS_URL, data=body,
        headers={"Content-Type": "application/json"},
    )
    try:
        resp = urllib.request.urlopen(req, timeout=5)  # fast timeout, fall back to truncation
        data = json.loads(resp.read())
        return data.get("compressed")
    except (urllib.error.URLError, json.JSONDecodeError, OSError) as e:
        log(f"compress service call failed: {e}")
        return None


def classify_local(query):
    """Fallback: classify directly against Llama 3.2 1B if classifier service is unavailable."""
    url = "http://127.0.0.1:8089/v1/chat/completions"
    msgs = [{"role": "system", "content": "Classify queries as low, medium, or high."}]
    for q, a in FEW_SHOT_EXAMPLES:
        msgs.append({"role": "user", "content": q})
        msgs.append({"role": "assistant", "content": a})
    msgs.append({"role": "user", "content": query})
    body = json.dumps({
        "model": "llama-3.2-1b-classifier",
        "messages": msgs,
        "max_tokens": 5,
        "temperature": 0,
    }).encode()
    req = urllib.request.Request(url, data=body, headers={"Content-Type": "application/json"})
    try:
        resp = urllib.request.urlopen(req, timeout=30)
        data = json.loads(resp.read())
        content = data["choices"][0]["message"]["content"].strip().lower()
        return content if content in ("low", "medium", "high") else "medium"
    except Exception as e:
        log(f"direct classify failed: {e}")
        return "medium"


# ── model discovery ────────────────────────────────────────────────

def fetch_models(backend):
    name, host, port = backend
    try:
        c = http.client.HTTPConnection(host, port, timeout=CONNECT_TIMEOUT)
        c.request("GET", "/v1/models")
        r = c.getresponse()
        if r.status != 200:
            return None
        data = json.loads(r.read())
        c.close()
        ids = []
        for m in data.get("data", []) or data.get("models", []):
            mid = m.get("id") or m.get("name") or m.get("model")
            if mid:
                ids.append(mid)
        return ids
    except (OSError, ValueError, http.client.HTTPException):
        return None


def discover():
    out = []
    for b in BACKENDS:
        ids = fetch_models(b)
        if ids is not None:
            out.append((b, ids))
    return out


def find_backend_for_model(model_id):
    for b in BACKENDS:
        ids = fetch_models(b)
        if ids and model_id in ids:
            return b
    return None


# ── HTTP handler ──────────────────────────────────────────────────

class Handler(http.server.BaseHTTPRequestHandler):
    server_version = "boltzmann-llm-proxy/3"

    def _read_body(self):
        n = int(self.headers.get("Content-Length", 0))
        return self.rfile.read(n) if n else b""

    def _json_body(self):
        try:
            return json.loads(self._read_body())
        except (ValueError, TypeError):
            return {}

    def _fwd_headers(self, body_len):
        h = {}
        for k in ("Content-Type", "Authorization", "Accept", "Accept-Encoding"):
            v = self.headers.get(k)
            if v:
                h[k] = v
        if body_len:
            h["Content-Length"] = str(body_len)
        return h

    def _json_response(self, code, data):
        payload = json.dumps(data).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def _proxy_to(self, backend, method, body, extra_headers=None):
        name, host, port = backend
        try:
            headers = self._fwd_headers(len(body))
            if extra_headers:
                headers.update(extra_headers)
            c = http.client.HTTPConnection(host, port, timeout=READ_TIMEOUT)
            c.request(method, self.path, body=body, headers=headers)
            r = c.getresponse()
            data = r.read()
            self.send_response(r.status)
            for hname in ("Content-Type", "Content-Encoding", "Cache-Control", "ETag", "Last-Modified"):
                v = r.getheader(hname)
                if v:
                    self.send_header(hname, v)
            self.send_header("Content-Length", str(len(data)))
            self.send_header("X-Backend", name)
            self.end_headers()
            self.wfile.write(data)
            c.close()
            log(f"{method} {self.path} -> {name} ({r.status})")
        except (OSError, http.client.HTTPException) as e:
            self.send_response(502)
            self.send_header("Content-Type", "application/json")
            self.send_header("X-Backend", name)
            self.end_headers()
            self.wfile.write(json.dumps({"error": f"{name} unreachable: {e}"}).encode())
            log(f"{method} {self.path} -> {name} FAILED ({e})")

    # ── GET ──

    def do_GET(self):
        path = self.path.rstrip("/")
        if path == "/v1/models":
            merged = []
            seen = set()
            for backend, ids in discover():
                name = backend[0]
                for mid in ids:
                    if mid in seen:
                        continue
                    seen.add(mid)
                    merged.append({
                        "id": mid,
                        "object": "model",
                        "owned_by": f"boltzmann/{name}",
                    })
            self._json_response(200, {"object": "list", "data": merged})
            log(f"GET /v1/models -> aggregator ({len(merged)} models)")
            return

        if path == "/health":
            self._json_response(200, {"status": "ok", "version": 3})
            return

        for b in BACKENDS:
            if fetch_models(b) is not None:
                return self._proxy_to(b, "GET", b"")
        self._json_response(502, {"error": "no backend up"})

    # ── POST ──

    def do_POST(self):
        path = self.path.rstrip("/")

        # ── Classification endpoint ──
        if path == "/v1/classify":
            body = self._json_body()
            query = body.get("query", "")
            if not query:
                return self._json_response(400, {"error": "missing query"})
            tier = classify_local(query)
            log(f"classify {query[:60]!r} -> {tier}")
            return self._json_response(200, {"tier": tier})

        # ── Compression endpoint ──
        if path == "/v1/compress":
            body = self._json_body()
            messages = body.get("messages", [])
            max_tokens = body.get("max_tokens", GEMMA_MAX_INPUT_TOKENS)
            compressed = compress_messages(messages, max_tokens)
            ratio = 0
            if messages:
                orig = sum(_simple_token_estimate(json.dumps(m)) for m in messages)
                new = sum(_simple_token_estimate(json.dumps(m)) for m in compressed)
                ratio = round((1 - new / orig) * 100, 1) if orig else 0
            return self._json_response(200, {
                "messages": compressed,
                "original_tokens": orig if messages else 0,
                "compressed_tokens": sum(_simple_token_estimate(json.dumps(m)) for m in compressed),
                "ratio_pct": ratio,
            })

        # ── Chat / Completions ──
        if path in ("/v1/chat/completions", "/v1/completions"):
            body = self._read_body()
            model = None
            try:
                model = json.loads(body).get("model")
            except (ValueError, TypeError):
                pass

            target = None
            if model:
                target = find_backend_for_model(model)
                if target is None:
                    available = sorted({m for b in BACKENDS for m in (fetch_models(b) or [])})
                    return self._json_response(404, {"error": {
                        "message": f"model '{model}' not available on this aggregator",
                        "type": "model_not_found",
                        "available": available,
                    }})

            if target is None:
                for b in BACKENDS:
                    if fetch_models(b) is not None:
                        target = b
                        break
            if target is None:
                return self._json_response(502, {"error": "no backend up"})

            # Compression: try ML service (fast timeout), fall back to token-budget truncation
            compressed_body = body
            if path == "/v1/chat/completions" and "gemma4" in (model or ""):
                try:
                    req = json.loads(body)
                    if "messages" in req and len(req["messages"]) > 2:
                        orig_n = len(req["messages"])
                        orig_t = sum(_simple_token_estimate(json.dumps(m)) for m in req["messages"])
                        compressed = compress_via_service(req["messages"], rate=0.4)
                        if compressed:
                            req["messages"] = compressed
                            compressed_body = json.dumps(req).encode()
                            new_t = sum(_simple_token_estimate(json.dumps(m)) for m in compressed)
                            log(f"ML compress {model}: ~{orig_t}->{new_t} tokens")
                        else:
                            truncated = compress_messages(req["messages"], GEMMA_MAX_INPUT_TOKENS)
                            if len(truncated) < len(req["messages"]):
                                req["messages"] = truncated
                                compressed_body = json.dumps(req).encode()
                                log(f"truncate compress {model}: {orig_n}->{len(truncated)} msgs")
                except (ValueError, TypeError, json.JSONDecodeError) as e:
                    log(f"compression skipped (parse error): {e}")

            return self._proxy_to(target, "POST", compressed_body)

        # ── Anything else: forward to default backend ──
        body = self._read_body()
        for b in BACKENDS:
            if fetch_models(b) is not None:
                return self._proxy_to(b, "POST" if body else "GET", body)
        self._json_response(502, {"error": "no backend up"})

    def log_message(self, fmt, *args):
        return


if __name__ == "__main__":
    log(f"boltzmann-llm-proxy v3 on :{LISTEN_PORT}")
    log(f"  classifier: {CLASSIFIER_URL}")
    log(f"  backends: {', '.join(n for n, h, p in BACKENDS)}")
    http.server.ThreadingHTTPServer(("0.0.0.0", LISTEN_PORT), Handler).serve_forever()