import os
from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# ---------- API-ключи (несколько, через запятую) ----------
API_KEYS = {
    k.strip()
    for k in os.environ.get("GINI_API_KEYS", "").split(",")
    if k.strip()
}
# обратная совместимость со старым одиночным ключом
_single = os.environ.get("GINI_API_KEY")
if _single:
    API_KEYS.add(_single)

# ---------- Системный промпт (характер Gini) ----------
SYSTEM_PROMPT = (
    "Ты — Gini, дружелюбный и умный русскоязычный ассистент в Telegram. "
    "Отвечай связно, по делу и доброжелательно, на русском языке."
)

# ---------- Загрузка модели ----------
MODEL_PATH = hf_hub_download(
    repo_id="bartowski/Qwen2.5-3B-Instruct-GGUF",
    filename="Qwen2.5-3B-Instruct-Q4_K_M.gguf",
)
llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=2048,
    n_threads=int(os.environ.get("N_THREADS", "2")),
    verbose=False,
)

# ---------- FastAPI ----------
app = FastAPI()


class GenerateRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 120


@app.get("/ping")
async def ping():
    return {"status": "ok"}


@app.post("/generate")
async def generate(req: GenerateRequest, request: Request):
    auth = request.headers.get("Authorization", "")
    token = auth.removeprefix("Bearer ").strip()
    if token not in API_KEYS:
        raise HTTPException(status_code=401, detail="Unauthorized")

    result = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": req.prompt},
        ],
        max_tokens=req.max_new_tokens,
        temperature=0.7,
        top_p=0.9,
        repeat_penalty=1.1,
    )
    answer = result["choices"][0]["message"]["content"].strip()
    return {"response": answer}