#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any

import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from scripts.qwen35_qwen3tts_voiceclone_smoke import generate_portuguese_answer, load_llm
from speech_bridge_gemma.mimi_overfit import resolve_device, set_seed


class AnswerRuntime:
    def __init__(self, args: argparse.Namespace) -> None:
        self.args = args
        self.loaded = False

    def load(self) -> None:
        if self.loaded:
            return
        set_seed(self.args.seed)
        self.device = resolve_device(self.args.device)
        self.tokenizer, self.llm, self.system, self.sft = load_llm(self.args, self.device)
        self.loaded = True

    def answer(self, payload: dict[str, Any]) -> dict[str, Any]:
        self.load()
        question = str(payload.get("question") or "").strip()
        context = str(payload.get("context") or "").strip()
        if not question:
            raise ValueError("question is required")
        answer = generate_portuguese_answer(
            llm=self.llm,
            tokenizer=self.tokenizer,
            question=question,
            system=self.system,
            device=self.device,
            max_new_tokens=self.args.max_answer_tokens,
            temperature=self.args.temperature,
            top_p=self.args.top_p,
            context=context,
        )
        return {
            "answer": answer,
            "system": self.system,
            "text_sft_loaded_layers": self.sft.get("text_sft_loaded_layers", 0),
        }


def build_app(runtime: AnswerRuntime) -> FastAPI:
    app = FastAPI()

    @app.get("/health")
    async def health() -> dict[str, Any]:
        return {"ok": True, "loaded": runtime.loaded}

    @app.post("/answer")
    async def answer(request: Request) -> JSONResponse:
        try:
            payload = json.loads((await request.body()).decode("utf-8") or "{}")
            return JSONResponse(runtime.answer(payload))
        except Exception as exc:
            return JSONResponse({"error": str(exc)}, status_code=500)

    return app


def parse_args() -> argparse.Namespace:
    root = Path(os.environ.get("S2S_ROOT", "/workspace/gemma1b-tts-integration"))
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", default=str(root))
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8767)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--llm-model", default="Qwen/Qwen3.5-0.8B")
    parser.add_argument("--llm-dtype", choices=["auto", "float32", "float16", "bfloat16"], default="float16")
    parser.add_argument("--text-sft-checkpoint", default="job_output/qwen3-text-sft-full-ptbr-20260606T022749Z/qwen_text_sft_layers.pt")
    parser.add_argument("--system", default="Voce e uma assistente brasileira. Responda sempre em portugues brasileiro, em uma frase curta, direta e natural.")
    parser.add_argument("--max-answer-tokens", type=int, default=48)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--top-p", type=float, default=0.9)
    parser.add_argument("--seed", type=int, default=1337)
    args = parser.parse_args()
    os.chdir(Path(args.root).resolve())
    return args


def main() -> int:
    args = parse_args()
    runtime = AnswerRuntime(args)
    uvicorn.run(build_app(runtime), host=args.host, port=args.port)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())