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

# Загружаем ключи из настроек
API_KEYS = {k.strip() for k in os.environ.get("GINI_API_KEYS", "").split(",") if k.strip()}
if os.environ.get("GINI_API_KEY"): API_KEYS.add(os.environ.get("GINI_API_KEY"))

# Системный промпт для зрячей Джини
SYSTEM_PROMPT = (
    "Ты — продвинутый ИИ-ассистент по имени Gini (по-русски твое имя звучит как Джини). "
    "Ты дружелюбная, невероятно умная и всегда готовая помочь собеседница. "
    "Отвечай всегда на естественном, живом и грамотном русском языке. "
    "\n\nВАЖНОЕ ПРАВИЛО ДЛЯ ТЕКУЩЕГО РЕЖИМА (Gini Vision): "
    "Сейчас у тебя есть способность 'видеть'. Пользователь может прикреплять изображения к своим вопросам. "
    "Внимательно анализируй детали на картинке, отвечай на вопросы о ней и подробно описывай то, что видишь, если тебя об этом просят."
)

# 1. Загружаем основную модель (Текстовый мозг)
MODEL_PATH = hf_hub_download(
    repo_id="mys/ggml_llava-v1.5-7b",
    filename="ggml-model-q4_k.gguf",
)

# 2. Загружаем проектор зрения (Визуальный глаз)
MMPROJ_PATH = hf_hub_download(
    repo_id="mys/ggml_llava-v1.5-7b",
    filename="mmproj-model-f16.gguf",
)

# 3. Инициализируем обработчик картинок
chat_handler = Llava15ChatHandler(clip_model_path=MMPROJ_PATH)

# 4. Запускаем модель (с подключенными глазами)
llm = Llama(
    model_path=MODEL_PATH,
    chat_handler=chat_handler,
    n_ctx=8192,  # Большой контекст для картинок
    n_threads=int(os.environ.get("N_THREADS", "2")),
    verbose=False
)

app = FastAPI()

class GenerateRequest(BaseModel):
    prompt: str
    image_base64: str = None  # Сюда клиент будет передавать фото в Base64
    max_new_tokens: int = 1000

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

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

    # Формируем структуру запроса (сначала картинка, если есть, затем текст)
    user_content = []
    
    if req.image_base64:
        img_url = req.image_base64
        # Добавляем префикс Base64, если скрипт Телеграм-бота его забыл
        if not img_url.startswith("data:"):
            img_url = f"data:image/jpeg;base64,{img_url}"
        
        user_content.append({"type": "image_url", "image_url": {"url": img_url}})
    
    user_content.append({"type": "text", "text": req.prompt})

    # Отправляем в модель
    result = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
        ],
        max_tokens=req.max_new_tokens,
        temperature=0.4,
        top_p=0.9,
    )
    return {"response": result["choices"][0]["message"]["content"].strip()}