import os
from dotenv import load_dotenv
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_community.vectorstores import Chroma

# Cargamos el .env desde la misma carpeta
load_dotenv()

# Validamos que tengamos la clave
if not os.getenv("GOOGLE_API_KEY"):
    print("ERROR: No se encontró GOOGLE_API_KEY en el .env")
    exit()

pdf_path = "Guia_DerechoTrabajoRural (2025) (2).pdf"

print("--- Iniciando ingesta en la carpeta ia_consultor ---")

# Cargamos el PDF
loader = PyPDFLoader(pdf_path)
data = loader.load()

# Dividimos en bloques
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
chunks = text_splitter.split_documents(data)

# --- USO MODELO CORRECTO ---
# Usamos text-embedding-004 sin el prefijo models/
model_name = "text-embedding-004"
print(f"DEBUG: Usando modelo: {model_name}")

# Quitamos el argumento 'model' para que use el default de la librería
embeddings = GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001")

# Guardamos
print("Guardando en ChromaDB...")
vector_db = Chroma.from_documents(
    documents=chunks, 
    embedding=embeddings,
    persist_directory="./chroma_db"
)

print(f"¡Éxito! Base creada con {len(chunks)} fragmentos.")