# -*- coding: utf-8 -*-
"""app.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/18vjUd8TiNpmeTVPmcgYEYrMyQlozkkzT
"""

import os
import torch
import gradio as gr

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain_core.prompts import PromptTemplate
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from pypdf import PdfReader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document


CAMINHO_DB = "db"


# =====================================================
#                 PREPARAÇÃO DO PDF
# =====================================================

def carregar_pdf(caminho):
    reader = PdfReader(caminho)
    textos = []
    for i, pagina in enumerate(reader.pages):
        texto = pagina.extract_text()
        if texto:
            textos.append(
                Document(page_content=texto, metadata={"page": i + 1})
            )
    return textos


def dividir_em_chunks(documentos):
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    return splitter.split_documents(documentos)


def vetorizar_chunks(chunks):
    embeddings = HuggingFaceEmbeddings(
        model_name="intfloat/multilingual-e5-small",
        model_kwargs={"device": "cpu"},
        encode_kwargs={"batch_size": 32},
    )

    db = Chroma(
        embedding_function=embeddings,
        persist_directory=CAMINHO_DB,
    )

    db.add_documents(chunks)
    db.persist()
    return db


def criar_db(caminho_pdf="Regimento.pdf"):
    documentos = carregar_pdf(caminho_pdf)
    chunks = dividir_em_chunks(documentos)
    return vetorizar_chunks(chunks)


# =====================================================
#     CRIAR DB AUTOMATICAMENTE SE NÃO EXISTIR
# =====================================================

if not os.path.exists(CAMINHO_DB):
    print("🔧 Criando base vetorial...")
    criar_db()
else:
    print("📚 Base vetorial encontrada. Carregando...")


# =====================================================
#                   PROMPT TEMPLATE
# =====================================================

prompt_template = """
Primeiramente, inicie a resposta sempre com "Oi, querido!"

E depois responda a pergunta:
{pergunta}

Com base nessas informações:
{base_conhecimento}

Se não houver resposta na base, diga apenas: não sei te dizer isso.
"""

_prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["pergunta", "base_conhecimento"],
)


# =====================================================
#                   EMBEDDINGS + CHROMA
# =====================================================

emb = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-small")
db = Chroma(persist_directory=CAMINHO_DB, embedding_function=emb)


# =====================================================
#                  CARREGAR MODELO LLM
# =====================================================

def carregar_modelo():
    MODEL = "Qwen/Qwen2.5-1.5B-Instruct"

    tok = AutoTokenizer.from_pretrained(MODEL)

    mdl = AutoModelForCausalLM.from_pretrained(
        MODEL,
        device_map="auto",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    )

    gen = pipeline(
        "text-generation",
        model=mdl,
        tokenizer=tok,
        max_new_tokens=1000,
        temperature=0.2,
        do_sample=False,
        pad_token_id=tok.eos_token_id,
        return_full_text=False,
    )

    return tok, gen


tok, generator = carregar_modelo()


# =====================================================
#                    FUNÇÕES DO RAG
# =====================================================

def listar_fontes(resultados):
    pags = []
    for doc, score in resultados:
        p = doc.metadata.get("page")
        if p and p not in pags:
            pags.append(p)
    return ", ".join([f"p.{p}" for p in pags])


def rag_chat(user_msg, history):
    resultados = db.similarity_search_with_relevance_scores(user_msg, k=3)

    if not resultados or resultados[0][1] < 0.7:
        resp = "Oi, querido! Não consegui encontrar algo relevante na base para responder com segurança."
        return history + [(user_msg, resp)]

    textos = [
        f"(p.{doc.metadata.get('page')}) {doc.page_content}"
        for doc, score in resultados
    ]

    base_conhecimento = "\n\n----\n\n".join(textos)

    mensagem = _prompt.format(
        pergunta=user_msg,
        base_conhecimento=base_conhecimento,
    )

    messages = [
        {"role": "system", "content": "Você responde em PT-BR de forma objetiva e educada."},
        {"role": "user", "content": mensagem},
    ]

    prompt_chat = tok.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    out = generator(prompt_chat, return_full_text=False)[0]["generated_text"].strip()

    fontes = listar_fontes(resultados)
    if fontes:
        out += f"\n\nFontes: {fontes}"

    return history + [(user_msg, out)]


# =====================================================
#                   INTERFACE GRADIO
# =====================================================

with gr.Blocks(title="CHAT IEPG") as demo:

    gr.Markdown("<h1 style='text-align:center;'>CHAT IEPG</h1>")
    gr.Markdown(
        "<p style='text-align:center;'>Faça perguntas sobre o Regimento. "
        "O chatbot usa RAG (Chroma + E5) com o modelo Qwen.</p>"
    )

    chat = gr.Chatbot(height=450)
    txt = gr.Textbox(label="Pergunta", placeholder="Digite sua pergunta...")
    btn = gr.Button("Enviar")
    clear = gr.Button("Limpar")

    def responder(user_msg, history):
        return rag_chat(user_msg, history), ""

    txt.submit(responder, [txt, chat], [chat, txt])
    btn.click(responder, [txt, chat], [chat, txt])
    clear.click(lambda: ([], ""), None, [chat, txt])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))