#!/usr/bin/env python3
"""
Laya HTTP server — bọc model Laya (decision model) thành REST API cho backend LearnAI.

Hỗ trợ 2 backend:
  --backend mlx    : `pip install laya-mlx`  (Apple Silicon, không cần PyTorch)
  --backend torch  : `pip install laya`      (PyTorch, Linux/Windows/WSL; hỗ trợ laya-multilingual)

Chạy:
  python server.py --backend torch --model convaiinnovations/laya-multilingual \
      --host 0.0.0.0 --port 8791 --token <TOKEN>

Hợp đồng API (xem README.md trong thư mục này và backend/src/services/layaService.ts):
  POST /predict  { state, questions }  ->  { model, answers, timing_ms }
  GET  /health                         ->  { status, model, backend }

LƯU Ý: API Python của Laya có thể thay đổi giữa các phiên bản. Script này gọi qua lớp
mỏng `_predict()` và tự dò vài tên hàm phổ biến (predict / system_one). Nếu phiên bản
bạn cài khác đi, chỉ cần sửa `_predict()`.
"""

from __future__ import annotations

import argparse
import os
import threading
import time
from typing import Any, Dict

# Model card Laya khuyến nghị: transformers dò TensorFlow/JAX lúc import có thể gây deadlock
# khi nạp model (abseil runtime của TF). Đặt trước khi import transformers/laya.
os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("USE_JAX", "0")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

from fastapi import FastAPI, Header, HTTPException, Request

parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=["mlx", "torch"], default="torch")
parser.add_argument("--model", default=os.environ.get("LAYA_MODEL", "convaiinnovations/laya-multilingual"))
parser.add_argument("--host", default=os.environ.get("LAYA_HOST", "0.0.0.0"))
parser.add_argument(
    "--port",
    type=int,
    # HF Spaces: app phải nghe cổng 7860 (hoặc khai báo app_port trong README của Space).
    default=int(os.environ.get("LAYA_PORT") or os.environ.get("PORT") or "8791"),
)
parser.add_argument("--token", default=os.environ.get("LAYA_SERVER_TOKEN", ""))
parser.add_argument("--max-len", type=int, default=0, help="0 = dùng mặc định của checkpoint")
args = parser.parse_args()

app = FastAPI(title="Laya server cho LearnAI")
_agent = None
_load_lock = threading.Lock()


def _load_agent() -> Any:
    """Nạp checkpoint (chỉ nạp một lần, an toàn khi gọi từ nhiều luồng)."""
    global _agent
    if _agent is None:
        with _load_lock:
            if _agent is None:
                started = time.perf_counter()
                if args.backend == "mlx":
                    import laya_mlx as laya  # type: ignore
                else:
                    import laya  # type: ignore

                _agent = laya.load(args.model)
                if args.max_len and hasattr(_agent, "cfg"):
                    _agent.cfg["max_len"] = args.max_len
                print(
                    f"[laya-server] đã nạp {args.model} (backend={args.backend}) "
                    f"trong {time.perf_counter() - started:.1f}s"
                )
    return _agent


@app.on_event("startup")
def _warmup_on_startup() -> None:
    """
    Mở cổng TRƯỚC, nạp model ở luồng nền. Bắt buộc với HF Spaces/Docker:
    health check phải trả 200 sớm, nếu không Space bị coi là "khởi động lỗi".
    Đặt LAYA_WARMUP=0 nếu muốn nạp thật sự lười (chỉ khi có request đầu tiên).
    """
    if os.environ.get("LAYA_WARMUP", "1") == "0":
        return
    threading.Thread(target=_load_agent, name="laya-warmup", daemon=True).start()


def _predict(state: Dict[str, Any], questions: Dict[str, Any]) -> Dict[str, Any]:
    """
    Lớp mỏng bọc API Python của Laya. Sửa ở ĐÂY nếu phiên bản Laya của bạn đổi tên hàm.

    Theo model card (https://huggingface.co/convaiinnovations/laya-multilingual):
        agent = laya.load("convaiinnovations/laya-multilingual")
        result = agent.predict(state, questions)
    ⇒ ưu tiên gọi PHƯƠNG THỨC `agent.predict`, sau đó mới tới các biến thể cũ.
    """
    agent = _load_agent()

    if hasattr(agent, "predict"):
        return agent.predict(state, questions)

    if args.backend == "mlx":
        import laya_mlx as laya  # type: ignore
    else:
        import laya  # type: ignore

    if hasattr(laya, "predict"):
        return laya.predict(agent, state, questions)
    if hasattr(laya, "system_one"):
        return laya.system_one(agent, state, questions)
    raise RuntimeError("Không tìm thấy hàm predict/system_one trong gói laya đã cài.")


def _extract_score(answer: Any) -> float:
    if isinstance(answer, dict):
        for key in ("score", "value", "level"):
            if key in answer and answer[key] is not None:
                return float(answer[key])
        if "probability" in answer and answer["probability"] is not None:
            return float(answer["probability"])
    return float(answer)


def _extract_probability(answer: Any) -> float:
    if isinstance(answer, dict):
        for key in ("probability", "prob", "p"):
            if key in answer and answer[key] is not None:
                return float(answer[key])
        if "score" in answer and answer["score"] is not None:
            score = float(answer["score"])
            return score if 0.0 <= score <= 1.0 else 0.5
    return float(answer)


@app.get("/health")
def health() -> Dict[str, Any]:
    return {"status": "ok", "backend": args.backend, "model": args.model, "loaded": _agent is not None}


@app.get("/warmup")
def warmup() -> Dict[str, Any]:
    """Nạp model ngay (đồng bộ) — tiện gọi sau khi Space vừa "thức" để tránh timeout."""
    started = time.perf_counter()
    _load_agent()
    return {
        "status": "ok",
        "model": args.model,
        "loaded": True,
        "loaded_in_ms": round((time.perf_counter() - started) * 1000, 1),
    }


@app.post("/predict")
async def predict(request: Request, authorization: str | None = Header(default=None)) -> Dict[str, Any]:
    if args.token:
        expected = f"Bearer {args.token}"
        if authorization != expected:
            raise HTTPException(status_code=401, detail="Token không hợp lệ.")

    payload = await request.json()
    state = payload.get("state")
    questions = payload.get("questions")
    if not state or not questions:
        raise HTTPException(status_code=422, detail="Thiếu 'state' hoặc 'questions'.")

    started = time.perf_counter()
    try:
        raw = _predict(state, questions)
    except Exception as exc:  # noqa: BLE001 — trả lỗi rõ ràng cho backend
        raise HTTPException(status_code=500, detail=f"Laya lỗi khi suy luận: {exc}") from exc

    raw_answers = raw.get("answers", raw) if isinstance(raw, dict) else {}
    answers: Dict[str, Any] = {}
    for name, answer in raw_answers.items():
        if isinstance(answer, dict) and any(k in answer for k in ("probability", "prob", "p")):
            answers[name] = {
                "probability": _extract_probability(answer),
                "answer": bool(answer.get("answer", answer.get("label", False))),
            }
        else:
            answers[name] = {"score": _extract_score(answer)}

    return {
        "model": payload.get("model") or args.model,
        "answers": answers,
        "timing_ms": round((time.perf_counter() - started) * 1000, 1),
    }


if __name__ == "__main__":
    import uvicorn

    # Không nạp model ở đây để cổng được mở ngay (xem _warmup_on_startup).
    print(f"[laya-server] {args.model} (backend={args.backend}) — nghe {args.host}:{args.port}")
    uvicorn.run(app, host=args.host, port=args.port)