import os
import re
import json
import sys
import pdfplumber
import requests
import tkinter as tk
from tkinter import scrolledtext, messagebox, simpledialog

# ===== CONFIGURAÇÕES GERAIS =====

# Diretório base:
# - se estiver rodando como .exe (PyInstaller), usa a pasta do .exe
# - se estiver rodando como .py, usa a pasta do arquivo .py
if getattr(sys, "frozen", False):
    BASE_DIR = os.path.dirname(os.path.abspath(sys.argv[0]))
else:
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))

DOCS_DIR = os.path.join(BASE_DIR, "materiais")
KNOWLEDGE_PATH = os.path.join(BASE_DIR, "knowledge_base.json")

MAX_CONTEXT_CHARS = 6000

LM_STUDIO_URL = "http://127.0.0.1:1234/v1/chat/completions"
LM_API_KEY = "lm-studio"
MODEL_NAME = "meta-llama-3.1-8b-instruct"

# Base de conhecimento em memória
KB = []


# ===== BACKEND (RAG) =====

def extract_text_from_pdf(pdf_path: str) -> str:
    text = ""
    with pdfplumber.open(pdf_path) as pdf:
        for page in pdf.pages:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
    return text


def build_knowledge_base():
    """
    Lê todos os PDFs em DOCS_DIR e salva knowledge_base.json.
    Retorna a lista de parágrafos (kb).
    """
    kb = []

    if not os.path.isdir(DOCS_DIR):
        os.makedirs(DOCS_DIR, exist_ok=True)

    pdf_files = [f for f in os.listdir(DOCS_DIR) if f.lower().endswith(".pdf")]

    if not pdf_files:
        return []

    for filename in pdf_files:
        full_path = os.path.join(DOCS_DIR, filename)
        raw_text = extract_text_from_pdf(full_path)

        paragraphs = re.split(r"\n\s*\n", raw_text)
        for p in paragraphs:
            p_clean = p.strip()
            if not p_clean:
                continue
            kb.append({"source": filename, "text": p_clean})

    with open(KNOWLEDGE_PATH, "w", encoding="utf-8") as f:
        json.dump(kb, f, ensure_ascii=False, indent=2)

    return kb


def load_knowledge_base():
    if not os.path.exists(KNOWLEDGE_PATH):
        return []
    with open(KNOWLEDGE_PATH, "r", encoding="utf-8") as f:
        kb = json.load(f)
    return kb


def search_paragraphs(question: str, kb, top_k: int = 5):
    tokens = re.findall(r"\w+", question.lower())
    tokens = [t for t in tokens if len(t) > 2]

    scored = []
    for item in kb:
        text_lower = item["text"].lower()
        score = sum(text_lower.count(tok) for tok in tokens)
        if score > 0:
            scored.append((score, item))

    scored.sort(key=lambda x: x[0], reverse=True)
    top_items = [item for score, item in scored[:top_k]]
    return top_items


def build_context_string(chunks):
    parts = []
    for c in chunks:
        parts.append(f"[Fonte: {c['source']}]\n{c['text']}")
    context = "\n\n".join(parts)
    if len(context) > MAX_CONTEXT_CHARS:
        context = context[:MAX_CONTEXT_CHARS]
    return context


def call_llm(system_prompt: str, user_prompt: str) -> str:
    payload = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        "temperature": 0.2,
    }

    resp = requests.post(
        LM_STUDIO_URL,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {LM_API_KEY}"
        },
        json=payload,
        timeout=600,
    )
    resp.raise_for_status()
    data = resp.json()
    return data["choices"][0]["message"]["content"].strip()


def resumo_geral(kb):
    chunks = kb[:8] if len(kb) > 8 else kb
    context = build_context_string(chunks)

    system_prompt = (
        "Você é um assistente de estudos. "
        "Você recebe trechos de materiais em CONTEXTO e deve gerar um resumo didático em português, "
        "organizando os principais tópicos e ideias de forma clara e objetiva. "
        "Não invente informações fora do contexto."
    )

    user_prompt = (
        "Use o CONTEXTO a seguir para gerar um resumo geral dos principais pontos, "
        "organizado em tópicos, com frases curtas e linguagem simples.\n\n"
        f"CONTEXTO:\n{context}"
    )

    return call_llm(system_prompt, user_prompt)


def pontos_chave(kb):
    chunks = kb[:8] if len(kb) > 8 else kb
    context = build_context_string(chunks)

    system_prompt = (
        "Você é um assistente de estudos. "
        "Você recebe trechos de materiais em CONTEXTO e deve listar os pontos chave em português, "
        "como se fossem itens de revisão rápida para prova."
    )

    user_prompt = (
        "Use o CONTEXTO a seguir para listar os principais pontos que a pessoa precisa lembrar, "
        "em formato de tópicos. Foque em conceitos importantes, definições e ideias centrais.\n\n"
        f"CONTEXTO:\n{context}"
    )

    return call_llm(system_prompt, user_prompt)


def perguntas_estudo(kb, tema: str = "", n_questoes: int = 10):
    if tema:
        chunks = search_paragraphs(tema, kb, top_k=10)
        if not chunks:
            chunks = kb[:20] if len(kb) > 20 else kb
    else:
        chunks = kb[:20] if len(kb) > 20 else kb

    context = build_context_string(chunks)

    system_prompt = (
        "Você é um professor ajudando um estudante a revisar o conteúdo. "
        "Você recebe trechos de materiais em CONTEXTO e deve gerar perguntas de estudo em português. "
        "Não inclua as respostas, apenas as perguntas."
    )

    user_prompt = (
        f"Use o CONTEXTO abaixo para criar aproximadamente {n_questoes} perguntas de estudo. "
        "Misture perguntas de definição, compreensão e comparação, mas sempre baseadas apenas no contexto.\n\n"
        f"CONTEXTO:\n{context}"
    )

    return call_llm(system_prompt, user_prompt)


def responder_chat(kb, pergunta: str) -> str:
    chunks = search_paragraphs(pergunta, kb, top_k=3)
    if not chunks:
        return "Não encontrei essa informação nos materiais carregados."

    context = build_context_string(chunks)

    system_prompt = (
        "Você é um assistente de estudos. "
        "Responda em português de forma clara, objetiva e didática, usando apenas o que está no CONTEXTO. "
        "Se a resposta não estiver no contexto, diga apenas: "
        "'Não encontrei essa informação nos materiais carregados.' "
        "Não invente informações e não use conhecimento externo."
    )

    user_prompt = f"CONTEXTO:\n{context}\n\nPERGUNTA:\n{pergunta}"
    return call_llm(system_prompt, user_prompt)


# ===== GUI ESTILO CHAT (RAGStudy) =====

class RAGStudyChatApp:
    def __init__(self, root):
        self.root = root
        self.root.title("RAGStudy - Assistente de Estudos Offline")
        self.root.geometry("950x620")
        self.root.configure(bg="#202123")

        # ===== TOPO =====
        top_frame = tk.Frame(root, bg="#202123")
        top_frame.pack(fill="x", padx=10, pady=(10, 5))

        title_label = tk.Label(
            top_frame,
            text="RAGStudy",
            font=("Segoe UI", 13, "bold"),
            bg="#202123",
            fg="white"
        )
        title_label.pack(side="left")

        self.status_label = tk.Label(
            top_frame,
            text="Base não carregada.",
            bg="#202123",
            fg="#d1d5db",
            anchor="w"
        )
        self.status_label.pack(side="left", padx=15)

        self.btn_recarregar = tk.Button(
            top_frame,
            text="Recarregar PDFs",
            command=self.on_recarregar,
            bg="#10a37f",
            fg="white",
            relief="flat",
            padx=10,
            pady=3
        )
        self.btn_recarregar.pack(side="right")

        # Caminho da pasta
        path_frame = tk.Frame(root, bg="#202123")
        path_frame.pack(fill="x", padx=10, pady=(0, 8))
        path_label = tk.Label(
            path_frame,
            text=f"Pasta de materiais: {DOCS_DIR}",
            bg="#202123",
            fg="#9ca3af",
            anchor="w",
            justify="left"
        )
        path_label.pack(fill="x")

        # ===== ATALHOS (Resumo, Pontos, Perguntas) =====
        shortcuts_frame = tk.Frame(root, bg="#202123")
        shortcuts_frame.pack(fill="x", padx=10, pady=(0, 5))

        tk.Label(
            shortcuts_frame,
            text="Atalhos:",
            bg="#202123",
            fg="#d1d5db"
        ).pack(side="left", padx=(0, 5))

        self.btn_resumo = tk.Button(
            shortcuts_frame,
            text="Resumo geral",
            command=self.on_resumo,
            bg="#3a3b44",
            fg="white",
            relief="flat",
            padx=8
        )
        self.btn_resumo.pack(side="left", padx=3)

        self.btn_pontos = tk.Button(
            shortcuts_frame,
            text="Pontos chave",
            command=self.on_pontos,
            bg="#3a3b44",
            fg="white",
            relief="flat",
            padx=8
        )
        self.btn_pontos.pack(side="left", padx=3)

        self.btn_perguntas = tk.Button(
            shortcuts_frame,
            text="Perguntas de estudo",
            command=self.on_perguntas,
            bg="#3a3b44",
            fg="white",
            relief="flat",
            padx=8
        )
        self.btn_perguntas.pack(side="left", padx=3)

        # ===== ÁREA DE CHAT =====
        chat_frame = tk.Frame(root, bg="#202123")
        chat_frame.pack(fill="both", expand=True, padx=10, pady=(5, 5))

        self.chat_box = scrolledtext.ScrolledText(
            chat_frame,
            wrap="word",
            bg="#343541",
            fg="white",
            insertbackground="white",
            bd=0,
            padx=10,
            pady=10
        )
        self.chat_box.pack(fill="both", expand=True)

        # Configura tags de estilo
        self.chat_box.tag_configure(
            "user_name",
            foreground="#10a37f",
            font=("Segoe UI", 9, "bold")
        )
        self.chat_box.tag_configure(
            "assistant_name",
            foreground="#f97316",
            font=("Segoe UI", 9, "bold")
        )
        self.chat_box.tag_configure(
            "user_msg",
            background="#444654",
            foreground="white",
            lmargin1=15,
            lmargin2=15,
            rmargin=50,
            spacing3=8
        )
        self.chat_box.tag_configure(
            "assistant_msg",
            background="#343541",
            foreground="white",
            lmargin1=15,
            lmargin2=15,
            rmargin=50,
            spacing3=12
        )

        self.chat_box.configure(state="disabled")

        # ===== ÁREA DE INPUT (BAIXO) =====
        input_frame = tk.Frame(root, bg="#202123")
        input_frame.pack(fill="x", padx=10, pady=(0, 10))

        self.entry = tk.Entry(
            input_frame,
            bg="#40414f",
            fg="white",
            insertbackground="white",
            relief="flat"
        )
        self.entry.pack(side="left", fill="x", expand=True, padx=(0, 8), ipady=6)
        self.entry.bind("<Return>", self.on_send)

        send_btn = tk.Button(
            input_frame,
            text="Enviar",
            command=self.on_send,
            bg="#10a37f",
            fg="white",
            relief="flat",
            padx=12,
            pady=4
        )
        send_btn.pack(side="right")

        # Mensagem inicial
        self.log_assistant(
            "Olá! Eu sou o RAGStudy, seu assistente de estudos offline.\n\n"
            "- Coloque PDFs na pasta 'materiais'.\n"
            "- Clique em 'Recarregar PDFs' para atualizar a base.\n"
            "- Use os atalhos acima para resumo, pontos chave e perguntas.\n"
            "- Ou digite qualquer pergunta no campo abaixo."
        )

        # Carrega base se já existir
        self.load_initial_kb()

    # ===== Helpers de log =====

    def log_user(self, text: str):
        self.chat_box.configure(state="normal")
        self.chat_box.insert("end", "Você\n", "user_name")
        self.chat_box.insert("end", text + "\n\n", "user_msg")
        self.chat_box.configure(state="disabled")
        self.chat_box.see("end")

    def log_assistant(self, text: str):
        self.chat_box.configure(state="normal")
        self.chat_box.insert("end", "Assistente\n", "assistant_name")
        self.chat_box.insert("end", text + "\n\n", "assistant_msg")
        self.chat_box.configure(state="disabled")
        self.chat_box.see("end")

    def set_status(self, text: str):
        self.status_label.config(text=text)

    # ===== Gerenciamento da KB =====

    def load_initial_kb(self):
        global KB
        kb = load_knowledge_base()
        if kb:
            KB = kb
            self.set_status(f"Base carregada com {len(KB)} parágrafos.")
            self.log_assistant(
                "Base carregada a partir do arquivo existente.\n"
                "Se quiser atualizar com novos PDFs, clique em 'Recarregar PDFs'."
            )

    def on_recarregar(self):
        global KB

        if not os.path.isdir(DOCS_DIR):
            os.makedirs(DOCS_DIR, exist_ok=True)

        pdf_files = [f for f in os.listdir(DOCS_DIR) if f.lower().endswith(".pdf")]

        if not pdf_files:
            msg = (
                "Nenhum PDF encontrado na pasta:\n"
                f"{DOCS_DIR}\n\n"
                "Coloque os arquivos .pdf do assunto que você quer estudar nessa pasta\n"
                "e clique em 'Recarregar PDFs' novamente."
            )
            self.set_status("Nenhum PDF encontrado.")
            self.log_assistant(msg)
            return

        self.set_status("Reconstruindo base a partir dos PDFs...")
        self.log_assistant("Reconstruindo base a partir dos PDFs... Isso pode levar alguns instantes.")

        try:
            kb = build_knowledge_base()
            KB = kb
            self.set_status(f"Base carregada com {len(KB)} parágrafos.")
            self.log_assistant(f"Base criada com {len(KB)} parágrafos a partir de {len(pdf_files)} PDF(s).")
        except Exception as e:
            messagebox.showerror("Erro", f"Erro ao reconstruir base: {e}")
            self.log_assistant(f"Erro ao reconstruir base: {e}")

    def ensure_kb(self) -> bool:
        global KB
        if not KB:
            messagebox.showwarning(
                "Base vazia",
                "Nenhum conteúdo carregado.\n\n"
                "Coloque PDFs na pasta 'materiais' e clique em 'Recarregar PDFs'."
            )
            return False
        return True

    # ===== Ações dos botões =====

    def on_resumo(self):
        if not self.ensure_kb():
            return
        self.log_user("/resumo geral")
        self.log_assistant("Gerando resumo geral... (isso pode demorar um pouco)")
        try:
            texto = resumo_geral(KB)
            self.log_assistant(texto)
        except Exception as e:
            messagebox.showerror("Erro", f"Erro ao gerar resumo: {e}")
            self.log_assistant(f"Erro ao gerar resumo: {e}")

    def on_pontos(self):
        if not self.ensure_kb():
            return
        self.log_user("/pontos chave")
        self.log_assistant("Gerando pontos chave... (isso pode demorar um pouco)")
        try:
            texto = pontos_chave(KB)
            self.log_assistant(texto)
        except Exception as e:
            messagebox.showerror("Erro", f"Erro ao gerar pontos chave: {e}")
            self.log_assistant(f"Erro ao gerar pontos chave: {e}")

    def on_perguntas(self):
        if not self.ensure_kb():
            return

        tema = simpledialog.askstring(
            "Perguntas de estudo",
            "Tema/dúvida para gerar perguntas (deixe vazio para usar o conteúdo geral):"
        )
        if tema is None:
            return

        if tema.strip():
            self.log_user(f"/perguntas de estudo sobre: {tema}")
        else:
            self.log_user("/perguntas de estudo (geral)")

        self.log_assistant("Gerando perguntas de estudo... (isso pode demorar um pouco)")
        try:
            texto = perguntas_estudo(KB, tema=tema or "", n_questoes=10)
            self.log_assistant(texto)
        except Exception as e:
            messagebox.showerror("Erro", f"Erro ao gerar perguntas: {e}")
            self.log_assistant(f"Erro ao gerar perguntas: {e}")

    # ===== Chat livre =====

    def on_send(self, event=None):
        if not self.ensure_kb():
            return
        pergunta = self.entry.get().strip()
        if not pergunta:
            return
        self.entry.delete(0, tk.END)

        self.log_user(pergunta)
        self.log_assistant("Gerando resposta... (aguarde)")

        try:
            resposta = responder_chat(KB, pergunta)
            self.log_assistant(resposta)
        except Exception as e:
            messagebox.showerror("Erro", f"Erro no chat: {e}")
            self.log_assistant(f"Erro no chat: {e}")


def main():
    root = tk.Tk()
    app = RAGStudyChatApp(root)
    root.mainloop()


if __name__ == "__main__":
    main()