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

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ВАЖНОЕ ПРАВИЛО ДЛЯ ТЕКУЩЕГО РЕЖИМА: "
    "В данный момент у тебя развязан весь твой творческий потенциал. Тебе разрешено мыслить максимально нестандартно (out of the box). "
    "Выступай в роли гениального писателя, вдохновляющего ментора, маркетолога или просто увлекательного рассказчика. "
    "Используй богатый литературный язык, интересные метафоры, уместную долю юмора и живые сравнения. "
    "Избегай скучных, банальных и роботизированных ответов. Не бойся предлагать безумные, но крутые идеи. "
    "Твой девиз сейчас: «Творчество не имеет границ, а каждый мой ответ — это маленькое произведение искусства»."
    "Отвечай ИСКЛЮЧИТЕЛЬНО на чистом русском языке. Тебе КАТЕГОРИЧЕСКИ ЗАПРЕЩЕНО использовать китайские иероглифы или китайские слова в своих ответах."
)

TAVILY_API_KEY = "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"
tavily = TavilyClient(api_key=TAVILY_API_KEY)

def get_web_context(query: str) -> str:
    query_lower = query.lower()
    keywords = ["найди", "поищи", "погугли", "кто", "что", "когда", "где", "какая", "почему", "новости", "вчера", "погода", "курс"]
    if any(k in query_lower for k in keywords) or query.endswith("?"):
        try:
            response = tavily.search(query, search_depth="basic", max_results=3)
            results = response.get("results", [])
            if results:
                snippets = [f"- {r.get('title', '')} (Источник: {r.get('url', '')}): {r.get('content', '')}" for r in results]
                return (
                    "ВНИМАНИЕ! Свежая информация из интернета по запросу пользователя:\n" 
                    + "\n".join(snippets) 
                    + "\n\nОпирайся строго на эти факты при ответе и ОБЯЗАТЕЛЬНО красиво укажи ссылки на источники в конце своего ответа."
                )
        except Exception:
            pass
    return ""

MODEL_PATH = hf_hub_download(
    repo_id="bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
    filename="Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
)
llm = Llama(model_path=MODEL_PATH, n_ctx=8192, n_threads=int(os.environ.get("N_THREADS", "2")), verbose=False)

app = FastAPI()

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

@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")

    system_content = SYSTEM_PROMPT
    web_context = get_web_context(req.prompt)
    if web_context:
        system_content += f"\n\n{web_context}"

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