#!/usr/bin/env python3
"""Query classifier service wrapping Llama 3.2 1B on :8089.

Endpoints:
  POST /classify  {"query": str} → {"tier": "low"|"medium"|"high"}
  GET  /health                   → {"status": "ok"}
"""
import http.server
import json
import os
import sys
import time
import urllib.request
import urllib.error

CLASSIFIER_URL = os.getenv("CLASSIFIER_URL", "http://127.0.0.1:8089/v1/chat/completions")
LISTEN_PORT = int(os.getenv("CLASSIFIER_LISTEN_PORT", "8090"))

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.stderr.write(f"{time.strftime('%H:%M:%S')} {msg}\n")
    sys.stderr.flush()


def build_messages(query):
    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})
    return msgs


def classify(query):
    body = json.dumps({
        "model": "llama-3.2-1b-classifier",
        "messages": build_messages(query),
        "max_tokens": 5,
        "temperature": 0,
    }).encode()
    req = urllib.request.Request(CLASSIFIER_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()
        if content in ("low", "medium", "high"):
            return content
        log(f"unexpected classifier output: {content!r} for query {query!r}")
        return "medium"
    except (urllib.error.URLError, json.JSONDecodeError, KeyError, IndexError) as e:
        log(f"classifier error: {e}")
        return "medium"


class Handler(http.server.BaseHTTPRequestHandler):
    server_version = "classifier/1"

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

    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 do_GET(self):
        if self.path == "/health":
            return self._json_response(200, {"status": "ok"})
        self._json_response(404, {"error": "not found"})

    def do_POST(self):
        if self.path == "/classify":
            body = self._read_body()
            query = body.get("query", "")
            if not query:
                return self._json_response(400, {"error": "missing query"})
            tier = classify(query)
            log(f"classify {query[:60]!r} -> {tier}")
            return self._json_response(200, {"tier": tier})
        self._json_response(404, {"error": "not found"})

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


if __name__ == "__main__":
    log(f"classifier-server on :{LISTEN_PORT}, backend {CLASSIFIER_URL}")
    http.server.ThreadingHTTPServer(("0.0.0.0", LISTEN_PORT), Handler).serve_forever()