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

# --- Авторизация ---
# Свой ключ для код-модели; если не задан — берём общий GINI_API_KEY
API_KEY = os.environ.get("GINI_CODE_API_KEY") or os.environ.get("GINI_API_KEY", "")

# --- Модель: Qwen2.5-Coder (специализирована на коде) ---
# Вариант по умолчанию — 7B (умнее). Для скорости см. блок ниже.
REPO_ID = "bartowski/Qwen2.5-Coder-7B-Instruct-GGUF"
FILENAME = "Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf"

# Быстрее, но проще — раскомментируй эти две строки и закомментируй две сверху:
# REPO_ID = "bartowski/Qwen2.5-Coder-3B-Instruct-GGUF"
# FILENAME = "Qwen2.5-Coder-3B-Instruct-Q4_K_M.gguf"

SYSTEM_PROMPT = (
    "Ты — Gini-Coder, помощник-программист. "
    "Отвечай кратко и по делу, выдавай рабочий код. "
    "Оборачивай код в markdown-блоки с указанием языка. "
    "Пояснения добавляй только там, где они реально нужны."
)

MODEL_PATH = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=4096,                # больше контекста под длинный код
    n_threads=int(os.environ.get("N_THREADS", "2")),
    verbose=False,
)

app = FastAPI()


class GenerateRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 512  # код обычно длиннее обычного ответа


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


@app.post("/generate")
def generate(req: GenerateRequest, authorization: str = Header(default="")):
    token = authorization.removeprefix("Bearer ").strip()
    if not API_KEY or token != API_KEY:
        raise HTTPException(status_code=401, detail="Unauthorized")

    out = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": req.prompt},
        ],
        max_tokens=req.max_new_tokens,
        temperature=0.2,        # низкая — код должен быть точным, не «фантазийным»
        top_p=0.9,
        repeat_penalty=1.05,
    )
    answer = out["choices"][0]["message"]["content"].strip()
    return {"response": answer}