# -*- coding: utf-8 -*-
import os
import torch
import gradio as gr

from pypdf import PdfReader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.schema import Document
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate

from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer


CAMINHO_PDF = "Regimento.pdf"
CAMINHO_DB = "db"


# =====================================
# 1. Carregar PDF e criar Vetor Store
# =====================================

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)
    chunks = splitter.split_documents(documentos)
    print(f"Dividido em {len(chunks)} chunks")
    return chunks


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

    db = Chroma(
        collection_name="regimento",
        embedding_function=embeddings,
        persist_directory=CAMINHO_DB
    )

    db.add_documents(chunks)
    return db


def criar_db(caminho_pdf):
    documentos = carregar_pdf(caminho_pdf)
    chunks = dividir_em_chunks(documentos)
    db = vetorizar(chunks)
    return db


# =====================================
# 2. Criar DB se não existir
# =====================================

if not os.path.exists(CAMINHO_DB):
    print("DB não encontrado. Criando...")
    criar_db(CAMINHO_PDF)
else:
    print("DB encontrado. Usando existente.")


# =====================================
# 3. Carregar Vetor Store + Modelo LLM
# =====================================

_emb = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-small")

_db = Chroma(
    collection_name="regimento",
    persist_directory=CAMINHO_DB,
    embedding_function=_emb
)

MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(MODEL)

mdl = AutoModelForCausalLM.from_pretrained(
    MODEL,
    device_map="auto",
    torch_dtype=torch.float32,
)

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


prompt_template = """
Primeiramente, inicie a resposta com "Oi, querido!".
Depois responda a pergunta do usuário:
{pergunta}

Com base SOMENTE nas informações abaixo:
{base_conhecimento}

Caso não encontre a resposta, diga: "não sei te dizer isso".
"""

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


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 encontrei informação relevante para responder."
        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 sempre fiel ao documento."},
        {"role": "user", "content": mensagem},
    ]

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

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

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

    return history + [(user_msg, out)]


# =====================================
# 4. Interface Gradio
# =====================================

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

    gr.Markdown("<h1>CHAT IEPG</h1>")
    gr.Markdown("Chatbot usando RAG + Qwen 2.5 Instruct")

    chat = gr.Chatbot(height=450)
    txt = gr.Textbox(label="Pergunta")

    def responder(msg, history):
        return rag_chat(msg, history)

    txt.submit(responder, [txt, chat], chat)


# =====================================
# 5. Launch
# =====================================

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