import streamlit as st
from streamlit_pdf_viewer import pdf_viewer
import base64
import os
from io import BytesIO
from pypdf import PdfReader
from langchain.schema import Document
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_pinecone import PineconeVectorStore
from pinecone import Pinecone as PineconeClient, ServerlessSpec

# Configuración de entorno
hf_token = os.getenv("HUGGINGFACE_TOKEN")
pinecone_api_key = os.getenv("PINECONE_API_KEY")

st.set_page_config(page_title="Clasificador de CVs", layout="wide")
st.title("🎯 Clasificador de CVs por Puesto de Trabajo")

# Cargar modelo de embeddings
@st.cache_resource(show_spinner="⏳ Cargando modelo de embeddings...")
def load_embeddings():
    model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    _ = model.embed_query("test")  # Calentamiento
    return model

embedding = load_embeddings()

# Inputs
titulo_puesto = st.text_input("🧑‍💼 Título del puesto", placeholder="Ej: Desarrollador Backend Senior")
descripcion_puesto = st.text_area("📝 Descripción del puesto", height=200)

# Subida de archivos
uploaded_files = st.file_uploader("📎 Subir CVs (PDF)", type="pdf", accept_multiple_files=True)

if uploaded_files:
    if "cv_data" not in st.session_state:
        st.session_state["cv_data"] = {}

    cv_data = st.session_state["cv_data"]

    for file in uploaded_files:
        if file.name in cv_data:
            st.info(f"ℹ️ El archivo `{file.name}` ya fue subido. Ignorando duplicado.")
            continue

        buffer = BytesIO(file.read())
        buffer.seek(0)
        try:
            reader = PdfReader(buffer)
            text = ""
            for page in reader.pages:
                page_text = page.extract_text()
                if page_text:
                    text += page_text + "\n"

            if text.strip():
                cv_data[file.name] = {
                    "text": text.strip(),
                    "pdf": buffer
                }
                st.success(f"✅ Procesado `{file.name}` correctamente.")
            else:
                st.warning(f"⚠️ No se pudo extraer texto de `{file.name}`.")
        except Exception as e:
            st.error(f"❌ Error procesando `{file.name}`: {e}")

# Procesamiento de CVs
if st.button("📊 Procesar CVs"):
    cv_data = st.session_state.get("cv_data", {})

    if not cv_data:
        st.warning("No hay CVs para procesar.")
        st.stop()

    if not descripcion_puesto.strip():
        st.warning("Debes ingresar una descripción del puesto.")
        st.stop()

    # Inicializar Pinecone
    pc = PineconeClient(api_key=pinecone_api_key)
    index_name = "cv-index"

    if index_name not in pc.list_indexes().names():
        pc.create_index(
            name=index_name,
            dimension=384,
            metric='cosine',
            spec=ServerlessSpec(cloud='aws', region='us-east-1')
        )
    index = pc.Index(index_name)
    index.delete(delete_all=True)

    vector_store = PineconeVectorStore(index=index, embedding=embedding)

    # Crear documentos
    documents = []
    for filename, data in cv_data.items():
        doc = Document(
            page_content=data["text"],
            metadata={"filename": filename, "titulo_puesto": titulo_puesto}
        )
        documents.append(doc)

    # Subir a Pinecone
    vector_store.add_documents(documents)

    # Búsqueda por similitud
    results = vector_store.similarity_search_with_score(descripcion_puesto, k=len(documents))
    st.success(f"{len(results)} CV(s) procesado(s).")

    # Mostrar resultados
    for doc, score in results:
        filename = doc.metadata["filename"]
        data = cv_data[filename]

        st.markdown("---")
        col1, col2 = st.columns([2, 1])

        with col1:
            data["pdf"].seek(0)
            pdf_bytes = data["pdf"].read()
            st.markdown(f"#### 👀 Visualizador PDF: `{filename}`")
            if pdf_bytes:
                pdf_viewer(pdf_bytes)
            else:
                st.error("⚠️ El archivo PDF está vacío o no se pudo leer.")


        with col2:
            st.markdown("#### 📄 Detalles")
            st.write(f"**Nombre del archivo:** `{filename}`")
            st.write(f"**Similitud con descripción:** `{score * 100:.2f}%`")

    # Opcional: eliminar vectores del índice (descomenta si deseas limpiar después)