import os
import streamlit as st
import numpy as np
import faiss

from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
from openai import OpenAI

# API Key da OpenAI
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

st.set_page_config(page_title="Professor de Medicina IA")
st.title("🧠 Professor de Medicina IA")

model = SentenceTransformer("all-MiniLM-L6-v2")


# ---------- Ler todos os PDFs ----------

def load_pdfs(folder="pdfs"):

    textos = []

    if not os.path.exists(folder):
        st.error(f"Pasta '{folder}' não encontrada.")
        return textos

    for file in os.listdir(folder):

        if file.endswith(".pdf"):

            caminho = os.path.join(folder, file)

            reader = PdfReader(caminho)

            for page in reader.pages:

                txt = page.extract_text()

                if txt and txt.strip():
                    textos.append(txt.strip())

    return textos


# ---------- Criar índice vetorial ----------

@st.cache_resource
def build_index():

    textos = load_pdfs()

    if not textos:
        return None, []

    embeddings = model.encode(textos)

    embeddings = np.array(embeddings).astype("float32")

    dimension = embeddings.shape[1]

    index = faiss.IndexFlatL2(dimension)

    index.add(embeddings)

    return index, textos


index, textos = build_index()


# ---------- Interface ----------

pergunta = st.text_input("Pergunte sobre medicina:")

if pergunta:

    if index is None:
        st.error("Nenhum texto foi carregado dos PDFs.")
    else:

        q_embed = model.encode([pergunta])
        q_embed = np.array(q_embed).astype("float32")

        D, I = index.search(q_embed, k=3)

        contexto = ""

        for i in I[0]:

            if i < len(textos):
                contexto += textos[i] + "\n\n"

        prompt = f"""
Use apenas o conteúdo abaixo para responder.

Se a resposta não estiver no conteúdo, diga que não encontrou informação suficiente.

Conteúdo:
{contexto}

Pergunta:
{pergunta}
"""

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}]
        )

        st.write(response.choices[0].message.content)