import os
from huggingface_hub import hf_hub_download, login
from llama_cpp import Llama
from fastapi import FastAPI
import uvicorn
import time
from concurrent.futures import ThreadPoolExecutor

# === Авторизация ===
token = os.getenv("TOKEN_FACE_GRAINED")
if token:
    login(token=token)

# === Загрузка модели ===
print("Скачивание модели...")
model_path = hf_hub_download(
    repo_id="Sanya2025/Proxima_B-instruct-Q5_0-GGUF",
    filename="proxima_b-instruct-q5_0.gguf",
    token=token,
    local_dir="/tmp/model",
    local_dir_use_symlinks=False
)
print(f"Модель: {model_path}")

print("Инициализация llama.cpp...")
llm = Llama(
    model_path=model_path,
    n_ctx=4096,
    n_threads=4,
    verbose=False
)
print("Модель готова!")

# === ThreadPool для неблокирующей генерации ===
executor = ThreadPoolExecutor(max_workers=1)

app = FastAPI()

@app.get("/")
def health():
    return {"status": "ok", "model": "proxima-b-instruct-q5_0"}

# !!! СИНХРОННЫЙ endpoint (не async) !!!
@app.post("/v1/chat/completions")
def chat(request: dict):
    print(f"[{time.strftime('%H:%M:%S')}] Запрос: {str(request)[:200]}")
    
    messages = request.get("messages", [])
    max_tokens = request.get("max_tokens", 128)
    
    print(f"Генерация начата (max_tokens={max_tokens})...")
    start = time.time()
    
    try:
        response = llm.create_chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            temperature=request.get("temperature", 0.7),
            stop=request.get("stop", []),
            stream=False
        )
        print(f"Генерация завершена за {time.time()-start:.1f} сек")
        return response
        
    except Exception as e:
        print(f"ОШИБКА генерации: {e}")
        return {"error": str(e)}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)