# ==============================================================================
# API DO BLAZERRHINO 3B: FASTAPI + LLAMA.CPP (DOCKER)
# ==============================================================================
# Autor: Davi Rediske / BlazerIA Team (2026)
# Este script baixa seu modelo GGUF automaticamente e cria uma rota de chat
# ultra rápida para integrar com o seu site ou Worker.
# ==============================================================================

import os
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

app = FastAPI(title="BlazerRhino 3B API - Docker")

# Ativa o CORS para que seu site BlazerIA faça chamadas para esta API sem bloqueios
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Configurações do seu repositório onde está o arquivo GGUF compilado do BlazerRhino
REPO_ID = "Davizig10jojo/BlazerRhino-3B-GGUF"
ARQUIVO_GGUF = "BlazerRhino-3B-Instruct.Q4_K_M.gguf"

print("📥 Baixando arquivo GGUF do Hugging Face Hub...")
try:
    caminho_modelo = hf_hub_download(repo_id=REPO_ID, filename=ARQUIVO_GGUF)
    print(f"✅ Modelo carregado com sucesso em: {caminho_modelo}")
except Exception as e:
    print(f"⚠️ Erro ao carregar arquivo personalizado ({e}). Tentando nome original...")
    # Fallback caso o nome padrão do Unsloth seja mantido no repositório
    caminho_modelo = hf_hub_download(repo_id=REPO_ID, filename="Qwen2.5-3B-Instruct.Q4_K_M.gguf")

print("🧠 Carregando o modelo na CPU gratuita do Space (16GB RAM)...")
# Otimizado para usar os 2 núcleos de processamento da CPU básica sem travar o servidor
llm = Llama(
    model_path=caminho_modelo,
    n_ctx=2048,          # Contexto confortável para conversas e estudos
    n_threads=2,         # Usa os 2 núcleos da CPU grátis do Space no limite máximo
    n_batch=512,
)
print("🚀 BlazerRhino 3B está online e pronto para receber requisições!")

class ChatRequest(BaseModel):
    prompt: str
    temperature: float = 0.7
    max_tokens: int = 1024

@app.get("/")
def home():
    return {
        "status": "online",
        "model": "BlazerRhino-3B-GGUF",
        "bypass": "Docker ativo via README.md Exploit!",
        "author": "Davi Rediske / BlazerIA"
    }

@app.post("/v1/chat")
async def chat(request: ChatRequest):
    try:
        # Executa a geração do texto na CPU básica do Space
        resposta = llm(
            request.prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            stop=["<|im_end|>"],
        )
        
        texto_resposta = resposta["choices"][0]["text"]
        return {"response": texto_resposta}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))