"""
Beryl Chat API — FastAPI raw inference proxy (Docker Space)
Calls api-inference.huggingface.co directly (resolves inside HF network).
POST /run/predict  →  { data: [messages_json, model_key] }
"""
import os, json, requests
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

HF_TOKEN = os.environ.get("HF_TOKEN", "")
# Featherless AI provider via HF Router — resolves everywhere, supports Qwen/Mistral
HF_API   = "https://router.huggingface.co/featherless-ai/v1/chat/completions"
HEADERS  = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}

MODELS = {
    "qwen":   "Qwen/Qwen2.5-72B-Instruct",
    "qwen7b": "Qwen/Qwen2.5-7B-Instruct",
    "glm":    "Qwen/Qwen2.5-7B-Instruct",
    "hermes": "NousResearch/Hermes-3-Llama-3.1-8B",
    "auto":   "Qwen/Qwen2.5-7B-Instruct",
}

EMOTION_KW = ["feel","lonely","sad","love","companion","miss",
              "emotional","relationship","heart","care","hurt"]

app = FastAPI(title="Beryl Chat API")


def route_model(messages: list, model_key: str) -> str:
    if model_key and model_key not in ("auto", ""):
        return MODELS.get(model_key, MODELS["auto"])
    last = (messages[-1].get("content","") if messages else "").lower()
    return MODELS["glm"] if any(k in last for k in EMOTION_KW) else MODELS["auto"]


def do_chat(messages: list, model_key: str) -> dict:
    model = route_model(messages, model_key)
    payload = {
        "model":       model,
        "messages":    messages,
        "max_tokens":  600,
        "temperature": 0.78,
        "stream":      False,
    }
    r = requests.post(HF_API, headers=HEADERS, json=payload, timeout=30)
    if r.status_code == 200:
        data = r.json()
        return {
            "response": data["choices"][0]["message"]["content"],
            "model":    model.split("/")[-1],
            "ok":       True,
        }
    # Fallback: try Qwen2.5-7B if primary model unavailable/quota
    if r.status_code in (429, 503, 400) and model != MODELS["qwen7b"]:
        payload["model"] = MODELS["qwen7b"]
        r2 = requests.post(HF_API, headers=HEADERS, json=payload, timeout=30)
        if r2.status_code == 200:
            data2 = r2.json()
            return {
                "response": data2["choices"][0]["message"]["content"],
                "model":    "Qwen2.5-7B-Instruct",
                "ok":       True,
            }
    return {"ok": False, "error": f"HTTP {r.status_code}: {r.text[:200]}", "response": ""}


@app.get("/health")
def health():
    return {"ok": True, "version": "1.0.0", "service": "beryl-chat-api"}


@app.post("/predict")
@app.post("/run/predict")
async def predict(request: Request):
    try:
        body = await request.json()
        data = body.get("data", [])
        messages_json = data[0] if len(data) > 0 else "[]"
        model_key     = data[1] if len(data) > 1 else "auto"
        messages      = json.loads(messages_json)
        result        = do_chat(messages, model_key)
        return JSONResponse({"data": [json.dumps(result)]})
    except Exception as e:
        err = json.dumps({"ok": False, "error": str(e), "response": ""})
        return JSONResponse({"data": [err]}, status_code=200)