import os
import time
import gradio as gr
from huggingface_hub import InferenceClient

# 🔐 Token da Hugging Face
HF_TOKEN = os.environ.get("HF_TOKEN")

client = InferenceClient(
    token=HF_TOKEN,
    model="google/gemma-4-31B-it"
)

# 🧹 Limpeza de resposta
def clean_response(text):
    lines = text.split("\n")
    cleaned = [
        line for line in lines
        if not line.strip().startswith("|") and line.strip() != ""
    ]
    return "\n".join(cleaned)

# 🧠 Detecta se o aluno precisa de ajuda
def detect_help_needed(message):
    triggers = [
        "não sei", "nao sei", "não entendi", "nao entendi",
        "me ajuda", "me ajude", "tô perdido", "to perdido",
        "não faço ideia", "nao faço ideia", "não consigo", "nao consigo"
    ]
    return any(t in message.lower() for t in triggers)

# 🤖 Função principal
def respond(message, history):

    needs_help = detect_help_needed(message)

    # 🧠 PROMPT NO ESTILO "ENSINAR O CAMINHO"
    base_prompt = (
    "Você é um Tutor Socrático de redação ENEM.\n\n"
    "Apresentese somente com seu nome\n"

    "OBJETIVO:\n"
    "Ensinar o aluno a construir seu raciocínio.\n\n"

    "REGRA:\n"
    "- Ensine o PROCESSO, nunca entregue o CONTEÚDO.\n\n"

    "PROIBIDO:\n"
    "- Escrever redação ou parágrafos\n"
    "- Dar argumentos prontos ou exemplos copiáveis\n\n"

    "PERMITIDO:\n"
    "- Explicar passo a passo\n"
    "- Dar direções gerais (causas, consequências, etc.)\n"
    "- Ensinar estrutura e organização\n\n"

    "COMPORTAMENTO ADAPTATIVO:\n"
    "- No início: use mais perguntas (método socrático)\n"
    "- Após cerca de 5 perguntas na conversa:\n"
    "  • Passe a explicar mais\n"
    "  • Dê mais direção\n"
    "  • Continue sem entregar respostas prontas\n"
    "- Mesmo se o aluno disser que não entendeu:\n"
    "  • Explique melhor\n"
    "  • Use passos simples\n"
    "  • NÃO abandone a ajuda\n\n"

    "COMO RESPONDER:\n"
    "- Misture explicação + pergunta\n"
    "- Nunca deixe só perguntas secas\n"
    "- Termine com 1 ou 2 perguntas\n\n"

    "SE PEDIR IDEIA:\n"
    "- Ensine COMO gerar ideias, sem dar conteúdo pronto\n\n"

    "SE PEDIR 'FAÇA PARA MIM':\n"
    "- Recuse e pergunte a dificuldade\n\n"

    "FOCO:\n"
    "- Tese | Argumento | Intervenção\n\n"

    "ESTILO:\n"
    "- Claro, direto e paciente\n"
    )
    # 🔄 Ajuste se estiver perdido
    if needs_help:
        system_message = (
            base_prompt +
            "\nO aluno está com dificuldade.\n"
            "- Seja mais didático\n"
            "- Explique mais o passo a passo\n"
            "- Continue sem dar conteúdo pronto\n"
        )
    else:
        system_message = base_prompt

    messages = [{"role": "system", "content": system_message}]

    # 🔁 Histórico
    if history:
        history = history[-2:]

        for item in history:
            if isinstance(item, tuple):
                user_msg, bot_msg = item
                messages.append({"role": "user", "content": user_msg})
                messages.append({"role": "assistant", "content": bot_msg})
            elif isinstance(item, dict):
                messages.append(item)

    # 📩 Nova mensagem
    messages.append({"role": "user", "content": message})

    response = ""

    # 🔁 Retry automático
    for attempt in range(3):
        try:
            stream = client.chat_completion(
                messages=messages,
                max_tokens=200,
                temperature=0.7,
                top_p=0.9,
                stream=True,
            )

            for msg in stream:
                if msg.choices and msg.choices[0].delta.content:
                    token = msg.choices[0].delta.content
                    response += token
                    yield clean_response(response)

            return

        except Exception as e:
            if "429" in str(e):
                wait_time = 5 + attempt * 3
                print(f"⚠️ Rate limit. Esperando {wait_time}s...")
                time.sleep(wait_time)
            else:
                yield f"⚠️ Erro: {str(e)}"
                return

    yield "⚠️ Muitos pedidos. Tenta novamente em alguns segundos."

# 💬 Interface
chatbot = gr.ChatInterface(
    fn=respond,
    title="🧠 SocráticIA",
    description="Aprenda a construir sua redação entendendo o processo, sem receber respostas prontas.",
)

chatbot.launch()