import os
from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
from src.database import get_embeddings
from langchain_postgres import PGVector
from dotenv import load_dotenv

load_dotenv()

def load_and_split_data(file_path='data/reglement.md'):
    """Lit le fichier Markdown et le découpe sémantiquement en chunks."""
    print(f"Lecture du fichier : {file_path}")
    with open(file_path, "r", encoding="utf-8") as f:
        markdown_document = f.read()

    headers_to_split_on = [
        ("#", "Sujet_Principal"),
        ("##", "Type_Regle"),
    ]
    markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
    md_header_splits = markdown_splitter.split_text(markdown_document)

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=700,
        chunk_overlap=200, 
        separators=["\n\n", "\n", ".", " ", ""]
    )
    final_chunks = text_splitter.split_documents(md_header_splits)
    
    print(f"Découpage terminé : {len(final_chunks)} chunks générés.")
    return final_chunks


def build_index(chunks, db_url: str, collection_name: str):
    """Transforme les chunks en vecteurs avec bge-m3 et les stocke dans PostgreSQL."""
    print("Chargement du modèle BAAI/bge-m3...")
    embeddings = get_embeddings()
 
    print(f"Connexion à la base et création de l'index vectoriel...")
    vector_store = PGVector.from_documents(
        embedding=embeddings,
        documents=chunks,
        collection_name=collection_name,
        connection=db_url,
        use_jsonb=True,
    )
    
    print("Indexation réussie dans la base de données PostgreSQL !")
    return vector_store


if __name__ == "__main__":
    FILE_PATH = "data/reglement.md"
    COLLECTION_NAME = "rh_cgi_vectors"
    DATABASE_URL = os.getenv('DATABASE_URL')

    if not DATABASE_URL:
        raise ValueError("Erreur : La variable DATABASE_URL est introuvable.")

    documents = load_and_split_data(FILE_PATH)
    vector_store = build_index(documents, DATABASE_URL, COLLECTION_NAME)