#!/usr/bin/env python3
"""Ollama shim for SMG local testing.

SMG's worker registration probes three endpoints Ollama doesn't have:
  1. GET /health       — reachability gate (fatal if missing)
  2. GET /server_info  — backend detection (200 = "sglang") AND the source
                         of the worker's model ID via metadata labels
Without them the router ends up with zero workers, or a worker registered
as UNKNOWN_MODEL_ID that no model name can route to.

This shim listens on :11435, answers /health and /server_info itself
(model name fetched live from Ollama's /v1/models), and transparently
forwards everything else to Ollama on :11434 (including streaming).

Run:  python3 ollama_health_shim.py [listen_port] [upstream_port]
      (defaults: listen 11435, upstream 11434)
Then: --worker-urls http://host.docker.internal:11435

Two engines for load-balancing demos = two pairs:
  OLLAMA_HOST=127.0.0.1:11436 ollama serve
  python3 ollama_health_shim.py 11437 11436
"""

import http.client
import json
import sys
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

UPSTREAM_HOST = "localhost"
LISTEN_PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 11435
UPSTREAM_PORT = int(sys.argv[2]) if len(sys.argv) > 2 else 11434

HOP_BY_HOP = {"host", "connection", "content-length", "transfer-encoding", "keep-alive"}


def _ollama_first_model():
    try:
        url = f"http://{UPSTREAM_HOST}:{UPSTREAM_PORT}/v1/models"
        with urllib.request.urlopen(url, timeout=5) as r:
            return json.load(r)["data"][0]["id"]
    except Exception:
        return "unknown"


def _respond_json(handler, payload):
    body = json.dumps(payload).encode()
    handler.send_response(200)
    handler.send_header("Content-Type", "application/json")
    handler.send_header("Content-Length", str(len(body)))
    handler.end_headers()
    handler.wfile.write(body)


class Shim(BaseHTTPRequestHandler):
    def _handle(self, method):
        if self.path == "/health":
            _respond_json(self, {"status": "ok"})
            return

        if self.path == "/server_info":
            model = _ollama_first_model()
            _respond_json(self, {
                "model_id": model,
                "served_model_name": model,
                "version": "0.0.0-ollama-shim",
                "max_total_tokens": 131072,
                "is_embedding": False,
            })
            return

        length = int(self.headers.get("Content-Length") or 0)
        body = self.rfile.read(length) if length else None
        headers = {k: v for k, v in self.headers.items() if k.lower() not in HOP_BY_HOP}

        conn = http.client.HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT, timeout=600)
        try:
            conn.request(method, self.path, body=body, headers=headers)
            resp = conn.getresponse()
            self.send_response(resp.status)
            for k, v in resp.getheaders():
                if k.lower() not in HOP_BY_HOP:
                    self.send_header(k, v)
            self.end_headers()
            while chunk := resp.read(65536):
                self.wfile.write(chunk)
                self.wfile.flush()
        except (BrokenPipeError, ConnectionResetError):
            pass  # client went away mid-stream
        finally:
            conn.close()

    def do_GET(self):
        self._handle("GET")

    def do_POST(self):
        self._handle("POST")

    def do_DELETE(self):
        self._handle("DELETE")

    def log_message(self, *args):
        pass


if __name__ == "__main__":
    print(f"Shim on :{LISTEN_PORT} -> Ollama on :{UPSTREAM_PORT} (adds /health + /server_info)")
    ThreadingHTTPServer(("0.0.0.0", LISTEN_PORT), Shim).serve_forever()