#!/usr/bin/env python3
import json, os, time, uuid, urllib.request, urllib.error
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
from transformers import AutoTokenizer

MODEL_DIR=os.environ["DEVSTRAL_MODEL"]
BACKEND=os.environ.get("DEVSTRAL_BACKEND","http://192.0.2.11:8031").rstrip("/")
PUBLIC_MODEL=os.environ.get("DEVSTRAL_PUBLIC_MODEL","devstral-small2-24b")
HOST=os.environ.get("PROXY_HOST","192.0.2.11")
PORT=int(os.environ.get("PROXY_PORT","8025"))

TOKENIZER=AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)

def normalize_content(content):
    if isinstance(content,str):
        return content
    if isinstance(content,list):
        parts=[]
        for item in content:
            if isinstance(item,dict) and item.get("type") in ("text","input_text"):
                parts.append(str(item.get("text","")))
            elif isinstance(item,str):
                parts.append(item)
        return "\n".join(x for x in parts if x)
    if content is None:
        return ""
    return str(content)

def render_messages(messages):
    normalized=[]
    for m in messages:
        normalized.append({
            "role": str(m.get("role","user")),
            "content": normalize_content(m.get("content",""))
        })
    return TOKENIZER.apply_chat_template(
        normalized,
        tokenize=False,
        add_generation_prompt=True
    )

def backend_post(path,payload,timeout=600):
    req=urllib.request.Request(
        BACKEND+path,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type":"application/json"},
        method="POST"
    )
    with urllib.request.urlopen(req,timeout=timeout) as r:
        return r.status,json.loads(r.read().decode("utf-8"))

def backend_health():
    try:
        with urllib.request.urlopen(BACKEND+"/health",timeout=3) as r:
            return r.status==200
    except Exception:
        return False

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

    def log_message(self,fmt,*args):
        print("%s - %s" % (self.address_string(), fmt%args), flush=True)

    def send_json(self,status,obj):
        raw=json.dumps(obj,ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type","application/json; charset=utf-8")
        self.send_header("Content-Length",str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def do_GET(self):
        if self.path=="/health":
            if backend_health():
                self.send_json(200,{"status":"ok","backend":"ok"})
            else:
                self.send_json(503,{"status":"degraded","backend":"down"})
            return
        if self.path=="/v1/models":
            self.send_json(200,{
                "object":"list",
                "data":[{
                    "id":PUBLIC_MODEL,
                    "object":"model",
                    "created":int(time.time()),
                    "owned_by":"local-vllm-proxy",
                    "root":MODEL_DIR,
                    "parent":None
                }]
            })
            return
        self.send_json(404,{"error":{"message":"not found"}})

    def do_POST(self):
        try:
            length=int(self.headers.get("Content-Length","0"))
            body=json.loads(self.rfile.read(length) or b"{}")
        except Exception as e:
            self.send_json(400,{"error":{"message":"invalid json: "+repr(e)}})
            return

        if self.path=="/v1/chat/completions":
            try:
                messages=body.get("messages") or []
                prompt=render_messages(messages)
                payload={
                    "model":PUBLIC_MODEL,
                    "prompt":prompt,
                    "max_tokens":int(body.get("max_tokens",body.get("max_completion_tokens",256))),
                    "temperature":float(body.get("temperature",0.0)),
                    "top_p":float(body.get("top_p",1.0)),
                    "stream":False
                }
                if body.get("stop") is not None:
                    payload["stop"]=body["stop"]
                status,out=backend_post("/v1/completions",payload)
                text=((out.get("choices") or [{}])[0].get("text") or "")
                resp={
                    "id":"chatcmpl-"+uuid.uuid4().hex[:24],
                    "object":"chat.completion",
                    "created":int(time.time()),
                    "model":PUBLIC_MODEL,
                    "choices":[{
                        "index":0,
                        "message":{"role":"assistant","content":text},
                        "finish_reason":((out.get("choices") or [{}])[0].get("finish_reason"))
                    }],
                    "usage":out.get("usage",{})
                }
                self.send_json(200 if status==200 else status,resp)
            except urllib.error.HTTPError as e:
                raw=e.read().decode(errors="replace")
                self.send_json(e.code,{"error":{"message":"backend http error","body":raw}})
            except Exception as e:
                self.send_json(500,{"error":{"message":"proxy render/backend failure","detail":repr(e)}})
            return

        if self.path=="/v1/completions":
            try:
                status,out=backend_post("/v1/completions",body)
                self.send_json(status,out)
            except urllib.error.HTTPError as e:
                self.send_json(e.code,{"error":{"message":e.read().decode(errors="replace")}})
            except Exception as e:
                self.send_json(500,{"error":{"message":repr(e)}})
            return

        self.send_json(404,{"error":{"message":"not found"}})

print(f"DEVSTRAL_CHAT_PROXY_V2_START host={HOST} port={PORT} backend={BACKEND}",flush=True)
ThreadingHTTPServer((HOST,PORT),Handler).serve_forever()