import os
import sys
import json
import logging
import random
import hashlib
import re 
import chromadb
import pysqlite3
from datetime import datetime

# =============================================================================
# 1. HACK DO SQLITE
# =============================================================================
sys.modules["sqlite3"] = sys.modules.pop("pysqlite3") 
from huggingface_hub import snapshot_download
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
from openai import OpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded

# =============================================================================
# LOGGING & CONFIG
# =============================================================================
def log_force(msg):
    print(f"👉 {msg}", flush=True)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="FoddaciTron RAG (Router + Classic Prompt)")

# --- MODELOS ---
MODEL_MAIN = "deepseek-ai/DeepSeek-V3.2" 
MODEL_ROUTER = "deepseek-ai/DeepSeek-V3.2"
MODEL_EMBEDDING = "intfloat/multilingual-e5-small"

DATASET_REPO_ID = "ricsrdocasro/FoddaciBrainWiki"

# --- RATE LIMIT ---
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
MAX_INPUT_CHARS = 500  # <--- DEFINA AQUI O LIMITE (500 é um bom tamanho)

@app.exception_handler(RateLimitExceeded)
async def custom_rate_limit_handler(request: Request, exc: RateLimitExceeded):
    return JSONResponse(
        status_code=429,
        content={"detail": "CALMA AÍ, VICIADO! 🛑 Muita mensagem. Vai tocar grama um pouco."}
    )

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

SILICON_KEY = os.environ.get("SILICONFLOW_API_KEY")
client = OpenAI(api_key=SILICON_KEY, base_url="https://api.siliconflow.com/v1")

# =============================================================================
# 💰 CONTROLE FINANCEIRO & PRIVACIDADE
# =============================================================================
DAILY_QUOTA = 10
USAGE_FILE = "user_tracker.json"

def get_client_ip_hash(request: Request):
    forwarded = request.headers.get("X-Forwarded-For")
    if forwarded:
        real_ip = forwarded.split(",")[0]
    else:
        real_ip = request.client.host
    return hashlib.sha256(real_ip.encode()).hexdigest()

def check_budget_limit(request: Request):
    user_hash = get_client_ip_hash(request)
    today = datetime.now().strftime("%Y-%m-%d")
    
    data = {"date": today, "usage": {}}
    
    if os.path.exists(USAGE_FILE):
        try:
            with open(USAGE_FILE, "r") as f:
                loaded = json.load(f)
                if loaded.get("date") == today:
                    data = loaded
        except: pass 

    user_usage = data["usage"].get(user_hash, 0)
    
    if user_usage >= DAILY_QUOTA:
        return False 

    data["usage"][user_hash] = user_usage + 1
    
    try:
        with open(USAGE_FILE, "w") as f:
            json.dump(data, f)
    except Exception as e:
        log_force(f"⚠️ Erro tracker: {e}")
        
    return True

def get_user_usage_count(request: Request):
    user_hash = get_client_ip_hash(request)
    today = datetime.now().strftime("%Y-%m-%d")
    if os.path.exists(USAGE_FILE):
        try:
            with open(USAGE_FILE, "r") as f:
                data = json.load(f)
                if data.get("date") == today:
                    return data["usage"].get(user_hash, 0)
        except: pass
    return 0

# =============================================================================
# MOTOR PRINCIPAL (FODDACI ENGINE - DUAL BRAIN)
# =============================================================================
class FoddaciEngine:
    def __init__(self):
        self.wiki_db = None
        self.episodes_db = None
        self.style_bank = []
        self.path_wiki = "foddaci_wiki"
        self.path_episodes = "foddaci_db_episodes"
        self.style_filename = "style_bank.json"
        # CAMINHO LOCAL (Da sua pasta no repositório)
        self.local_style_path = "foddaci_db_v2/style_bank.json"

    def load_resources(self):
        log_force("🚀 INICIANDO SISTEMA...")

        # 1. DOWNLOAD DOS BANCOS (EXTERNO)
        try:
            log_force(f"📥 Baixando DBs de: {DATASET_REPO_ID}...")
            snapshot_download(
                repo_id=DATASET_REPO_ID,
                repo_type="dataset",
                local_dir=".", 
                resume_download=True,
                allow_patterns=["foddaci_wiki/*", "foddaci_db_episodes/*"]
            )
            log_force("✅ Download dos bancos concluído.")
        except Exception as e:
            log_force(f"❌ Erro Download DB: {e}")

        # 2. CARREGA ESTILO (DA PASTA LOCAL 'foddaci_db_v2')
        if os.path.exists(self.local_style_path):
            try:
                with open(self.local_style_path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    self.style_bank = data.get("comments", []) if isinstance(data, dict) else data
                log_force(f"🤬 Estilo carregado de '{self.local_style_path}': {len(self.style_bank)} frases.")
            except Exception as e:
                log_force(f"⚠️ Erro ao ler style_bank: {e}")
        else:
            log_force(f"⚠️ AVISO: '{self.local_style_path}' não encontrado. Verifique se a pasta subiu no git.")

        # 3. EMBEDDINGS (E5)
        try:
            embeddings = HuggingFaceEmbeddings(
                model_name=MODEL_EMBEDDING,
                encode_kwargs={'normalize_embeddings': True}
            )
        except Exception as e:
            log_force(f"❌ Erro Embeddings: {e}")
            return

        # 4. CONECTA WIKI
        if os.path.exists(self.path_wiki):
            try:
                self.wiki_db = Chroma(
                    persist_directory=self.path_wiki,
                    embedding_function=embeddings,
                    collection_name="foddaci_wiki"
                )
                log_force("📘 Wiki ON.")
            except: pass
        else:
            log_force("⚠️ Banco Wiki não encontrado após download.")
        
        # 5. CONECTA EPISÓDIOS
        if os.path.exists(self.path_episodes):
            try:
                self.episodes_db = Chroma(
                    persist_directory=self.path_episodes,
                    embedding_function=embeddings,
                    collection_name="foddaci_episodes"
                )
                log_force("📺 Episódios ON.")
            except: pass
        else:
            log_force("⚠️ Banco Episódios não encontrado após download.")

    def get_rag_data(self, query: str):
        formatted_query = f"query: {query}"
        log_force(f"❓ Buscando: '{formatted_query}'")
        
        docs_found = []

        # Busca Wiki
        if self.wiki_db:
            try:
                results = self.wiki_db.similarity_search_with_score(formatted_query, k=5)
                for doc, score in results:
                    if score < 1.45: 
                        clean = doc.page_content.replace("passage: ", "")
                        docs_found.append(f"[FONTE: WIKI/BONDA] {clean}")
            except: pass

        # Busca Episódios
        if self.episodes_db:
            try:
                results = self.episodes_db.similarity_search_with_score(formatted_query, k=5)
                for doc, score in results:
                    if score < 1.45:
                        origem = doc.metadata.get("arquivo", "Episódio Desconhecido")
                        clean = doc.page_content.replace("passage: ", "")
                        docs_found.append(f"[FONTE: VÍDEO '{origem}'] {clean}")
            except: pass

        context = "\n\n".join(docs_found) if docs_found else ""
        
        style = ""
        if self.style_bank:
            style = "\n".join([f"- {s}" for s in random.sample(self.style_bank, min(5, len(self.style_bank)))])
            
        return context, style

engine = FoddaciEngine()

@app.on_event("startup")
async def startup():
    engine.load_resources()

# --- INPUT MODEL ---
class ChatMessage(BaseModel):
    role: str
    content: str

class UserRequest(BaseModel):
    message: str
    history: List[ChatMessage] = []

# --- ROTAS ---
@app.get("/")
@app.head("/")
def home():
    return {"status": "Online", "mode": "Router R1 + Prompt Clássico"}

@app.get("/usage")
def get_usage(request: Request):
    current_count = get_user_usage_count(request)
    return {
        "count": current_count, 
        "limit": DAILY_QUOTA,
        "status": "blocked" if current_count >= DAILY_QUOTA else "open"
    }

@app.post("/chat")
@limiter.limit("10/minute") 
async def chat_handler(req: UserRequest, request: Request):
    
    # 1. COTA
    if not check_budget_limit(request):
        raise HTTPException(status_code=429, detail="COTA DIÁRIA ATINGIDA.")

    if len(req.message) > MAX_INPUT_CHARS:
        raise HTTPException(
            status_code=400, 
            detail=f"TEXTO MUITO LONGO! ({len(req.message)}/{MAX_INPUT_CHARS}). Resuma essa bíblia aí, não sou leitor de TCC."
        )

    user_msg = req.message
    search_query = None 
    
    # 2. R1 ROUTER (Decide SE busca e O QUE busca)
    log_force("🤔 R1 Analisando...")
    try:
        history_txt = "\n".join([f"{m.role}: {m.content}" for m in req.history[-3:]])
        
        prompt = f"""
        Você é um assistente especialista no canal Jogandofoddaci.
        DEFINIÇÕES: "Bonda" = Série de Álbuns Musicais (Lore).
        
        Histórico: {history_txt}
        Usuário: "{user_msg}"
        
        Tarefa: 
        1. Se for conversa fiada, retorne "NO_SEARCH".
        2. Se for dúvida, Reescreva a pergunta para Wiki (Substitua pronomes).
        Retorne APENAS a pergunta reescrita.
        """
        
        resp = client.chat.completions.create(
            model=MODEL_ROUTER, 
            messages=[{"role": "user", "content": prompt}],
            max_tokens=150,
            temperature=0.3
        )
        
        raw_r1 = resp.choices[0].message.content
        clean_decision = re.sub(r'<think>.*?</think>', '', raw_r1, flags=re.DOTALL).strip().replace('"', '')
        
        log_force(f"🧭 Router: {clean_decision}")
        
        if "NO_SEARCH" not in clean_decision and len(clean_decision) > 2:
            search_query = clean_decision
            
    except Exception as e:
        log_force(f"⚠️ Erro Router: {e}")
        search_query = user_msg 

    # 3. BUSCA (Se necessário)
    context = ""
    style_examples = ""
    
    if search_query:
        context, style_examples = engine.get_rag_data(search_query)
    else:
        log_force("⏩ Modo Chat (Sem busca)")
        if engine.style_bank:
            style_examples = "\n".join([f"- {s}" for s in random.sample(engine.style_bank, 3)])

    # 4. GERAÇÃO (PROMPT CLÁSSICO RESTAURADO)
    # Aqui entra exatamente o prompt que você pediu
    system_instruction = f"""
    Você é o FODDACITRON, mascote do canal Jogandofoddaci.
    
    ### HIERARQUIA DE VERDADE:
    1. **BASE DE DADOS (RAG):** É a fonte principal.
    2. **HISTÓRICO:** Se não estiver no banco, use o histórico.
    
    ### SUAS DIRETRIZES:
    - Se a informação não existir, xingue o usuário, mas não invente dados.
    - PERSONALIDADE: Tóxico, ofensivo e engraçado.
    - ESTILO: Integre ofensas na explicação naturalmente.

    ### REGRAS DE FORMATAÇÃO (CRÍTICO PARA O FRONTEND):
    1. **USE MARKDOWN À VONTADE:** O frontend suporta e fica bonito.
       - Use **NEGRITO** (`**palavra**`) para destacar nomes, xingamentos ou ênfases.
       - Use `#` para Títulos/Gritos se quiser impactar.
    
    ### GLOSSÁRIO (Vocabulário - Não copie frases inteiras):
    {style_examples}
    
    ### BASE DE DADOS - SUA MEMÓRIA (Fatos para: "{search_query if search_query else user_msg}"):
    {context if context else "Nada no RAG. Verifique o Histórico."}
    """
    
    msgs = [{"role": "system", "content": system_instruction}]
    for m in req.history[-5:]:
        msgs.append({"role": m.role, "content": m.content})
    msgs.append({"role": "user", "content": user_msg})
    
    return StreamingResponse(stream_deepseek(msgs), media_type="text/plain")

async def stream_deepseek(messages_list):
    try:
        stream = client.chat.completions.create(
            model=MODEL_MAIN, 
            messages=messages_list,
            stream=True,
            temperature=0.7, 
            max_tokens=1500
        )
        for chunk in stream:
            if chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content
    except Exception as e:
        yield f"Erro API: {e}"