#!/usr/bin/env python3
import argparse
import json
import os
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import mlx.core as mx
from mlx_lm import generate, load


def build_prompt(tokenizer, messages):
    if hasattr(tokenizer, "apply_chat_template"):
        try:
            return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        except Exception:
            pass

    parts = []
    for message in messages:
        role = str(message.get("role", "user")).upper()
        content = str(message.get("content", "")).strip()
        if content:
            parts.append(f"{role}: {content}")
    parts.append("ASSISTANT:")
    return "\n\n".join(parts)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=11436)
    parser.add_argument("--model-path", required=True)
    args = parser.parse_args()

    model, tokenizer = load(args.model_path)

    def embed_text(text):
        encoded = tokenizer.encode(str(text or "").strip(), add_special_tokens=True)
        if not encoded:
            eos_token_id = getattr(tokenizer, "eos_token_id", None)
            encoded = [eos_token_id if eos_token_id is not None else 0]

        tokens = mx.array([encoded])
        encoder = getattr(model, "model", model)
        hidden_states = encoder(tokens)
        if isinstance(hidden_states, tuple):
            hidden_states = hidden_states[0]
        if len(hidden_states.shape) == 3:
            hidden_states = hidden_states[0]
        pooled = mx.mean(hidden_states, axis=0) if len(hidden_states.shape) > 1 else hidden_states
        pooled = pooled.astype(mx.float32)
        norm = mx.sqrt(mx.sum(pooled * pooled))
        mx.eval(pooled, norm)
        norm_value = float(norm.item()) if hasattr(norm, "item") else float(norm)
        if norm_value > 0:
            pooled = pooled / norm
            mx.eval(pooled)
        return [float(value) for value in pooled.tolist()]

    class Handler(BaseHTTPRequestHandler):
        def respond_json(self, payload, status=200):
            body = json.dumps(payload).encode("utf-8")
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            if self.path == "/health":
                return self.respond_json({
                    "ok": True,
                    "pid": os.getpid(),
                    "model_path": args.model_path,
                    "embedding_ready": True
                })
            return self.respond_json({"error": "Not found"}, 404)

        def do_POST(self):
            if self.path not in {"/chat", "/embed"}:
                return self.respond_json({"error": "Not found"}, 404)

            try:
                content_length = int(self.headers.get("Content-Length", "0"))
                payload = json.loads(self.rfile.read(content_length) or b"{}")
                if self.path == "/embed":
                    text = payload.get("text", "")
                    return self.respond_json({"embedding": embed_text(text)})
                messages = payload.get("messages") or []
                max_tokens = int(payload.get("max_tokens") or 700)
                temperature = float(payload.get("temperature") or 0.2)
                prompt = build_prompt(tokenizer, messages)
                output = generate(
                    model,
                    tokenizer,
                    prompt=prompt,
                    max_tokens=max_tokens,
                    temp=temperature,
                    verbose=False
                )
                return self.respond_json({"content": output})
            except Exception as error:
                return self.respond_json({"error": str(error)}, 500)

        def log_message(self, _format, *_args):
            return

    server = ThreadingHTTPServer((args.host, args.port), Handler)
    server.serve_forever()


if __name__ == "__main__":
    main()
