#!/usr/bin/env python3
"""Serve the bonsai + decider-4b MetaCog merge as ONE OpenAI-compatible model.

    thinker: bonsai 27B (ternary) ............ THINKER_URL  (default http://localhost:8010)
    judge:   decider-4b (/v1/systemone) ..... JUDGE_URL    (default http://localhost:8008)

    POST /v1/chat/completions {"messages": [...]} -> judge-picked answer.

Env: PORT (8200), THINKER_URL, THINKER_MODEL (bonsai), JUDGE_URL,
     N_PATHS (8), NO_CASCADE, MAX_TOKENS (2048), TEMPERATURE (0.8),
     STOP_CONFIDENCE (0.80), NO_RACE (unset).
"""
from __future__ import annotations

import json
import os
import time
import uuid

import re

def clean_answer(text: str) -> str:
    """Strip <think>...</think> reasoning blocks the thinker may leak."""
    return re.sub(r"<think>.*?</think>\s*", "", text, flags=re.S).strip()

from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from metacog import Config, MetaCog, OpenAICompatThinker, SystemOneJudge
from race import race

MODEL_ID = "heuristic-1"


def env(name: str, default: str) -> str:
    return os.environ.get(name, default)


thinker = OpenAICompatThinker(
    env("THINKER_URL", "http://localhost:8010"),
    model=env("THINKER_MODEL", "bonsai"),
)
judge = SystemOneJudge(base_url=env("JUDGE_URL", "http://localhost:8008"))
USE_RACE = env("NO_RACE", "") != "1"
STOP_CONFIDENCE = float(env("STOP_CONFIDENCE", "0.80"))
N_PATHS = int(env("N_PATHS", "8"))
MAX_TOKENS = int(env("MAX_TOKENS", "2048"))
TEMPERATURE = float(env("TEMPERATURE", "0.8"))
MC = MetaCog(
    thinker,
    judge,
    Config(
        mode="best_of_n",
        strategy="noul",
        n_paths=N_PATHS,
        greedy_anchor=True,
        cascade_confidence=None if env("NO_CASCADE", "") == "1" else 0.95,
        max_tokens=MAX_TOKENS,
        temperature=TEMPERATURE,
    ),
)


def problem_from(messages):
    for m in reversed(messages):
        if m.get("role") == "user":
            c = m.get("content", "")
            return c if isinstance(c, str) else json.dumps(c)
    return "\n".join(f"{m.get('role')}: {m.get('content')}" for m in messages)


class Handler(BaseHTTPRequestHandler):
    server_version = "MergeServer/1.0"

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

    def log_message(self, fmt, *args):
        print(f"[{time.strftime('%H:%M:%S')}] {self.command} {self.path}", flush=True)

    def do_GET(self):
        path = self.path.rstrip("/")
        if path in ("/health", "/v1/health", "/v1/models"):
            self._send(200, {"status": "ok", "model": MODEL_ID})
        else:
            self._send(404, {"error": "not found"})

    def do_POST(self):
        path = self.path.rstrip("/")
        if path not in ("/v1/chat/completions", "/chat/completions"):
            return self._send(404, {"error": "not found"})
        try:
            raw = self.rfile.read(int(self.headers.get("Content-Length", 0))) or b"{}"
            body = json.loads(raw)
        except Exception as e:
            return self._send(400, {"error": f"bad JSON: {e}"})
        messages = body.get("messages", [])
        if not messages:
            return self._send(400, {"error": "messages required"})
        t0 = time.time()
        try:
            if USE_RACE:
                r = race(problem_from(messages), thinker, judge,
                         n_paths=N_PATHS, stop_confidence=STOP_CONFIDENCE,
                         max_tokens=MAX_TOKENS, temperature=TEMPERATURE)
                answer, finished = clean_answer(r.answer), r.finished
                conf, toks, dt = r.confidence, r.thinker_tokens, r.seconds
            else:
                result = MC.run(problem_from(messages))
                dt = time.time() - t0
                rounds = result.trace.rounds
                conf = float(rounds[-1].verdict.confidence) if rounds else 0.0
                answer, finished, toks = (clean_answer(result.answer),
                                          result.finished,
                                          result.trace.thinker_tokens)
        except Exception as e:
            return self._send(502, {"error": f"pipeline failed: {e}"})
        print(f"[{time.strftime('%H:%M:%S')}] conf={conf:.2f} {dt:.1f}s", flush=True)
        self._send(
            200,
            {
                "id": f"merge-{uuid.uuid4().hex[:12]}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": body.get("model", MODEL_ID),
                "choices": [
                    {
                        "index": 0,
                        "message": {"role": "assistant", "content": answer},
                        "finish_reason": "stop" if finished else "length",
                    }
                ],
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": toks,
                    "total_tokens": toks,
                },
                "merge": {
                    "thinker": "bonsai",
                    "judge": "decider-4b",
                    "judge_confidence": round(conf, 4),
                    "seconds": round(dt, 2),
                },
            },
        )


def main() -> None:
    port = int(env("PORT", "8200"))
    print(f"{MODEL_ID} on :{port}", flush=True)
    ThreadingHTTPServer(("0.0.0.0", port), Handler).serve_forever()


if __name__ == "__main__":
    main()