import os
from pathlib import Path

from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings

from dotenv import load_dotenv

load_dotenv(override=True)

# -----------------------------
# CONFIG
# -----------------------------
DB_NAME = str(Path(__file__).parent / "vector_db")
FILE_PATH = str( "about-us.md")

# -----------------------------
# EMBEDDINGS
# -----------------------------
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)

# -----------------------------
# LOAD & CHUNK DOCUMENT
# -----------------------------
def fetch_and_chunk_document():
    loader = TextLoader(FILE_PATH, encoding="utf-8")
    raw_document = loader.load()[0]

    headers_to_split_on = [
        ("#", "Section"),
        ("##", "Header2"),
        ("###", "Header3"),
    ]

    markdown_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=headers_to_split_on,
        strip_headers=False
    )

    md_chunks = markdown_splitter.split_text(raw_document.page_content)

    # Safety net: only splits chunks that are too large
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50
    )

    chunks = text_splitter.split_documents(md_chunks)

    for chunk in chunks:
        chunk.metadata["doc_type"] = "about-us"

    return chunks


# -----------------------------
# CREATE VECTOR STORE
# -----------------------------
def create_embeddings(chunks):
    # reset DB if exists
    if os.path.exists(DB_NAME):
        Chroma(
            persist_directory=DB_NAME,
            embedding_function=embeddings
        ).delete_collection()

    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=DB_NAME
    )

    # debug info
    collection = vectorstore._collection
    count = collection.count()

    sample_embedding = collection.get(limit=1, include=["embeddings"])["embeddings"][0]
    dimensions = len(sample_embedding)

    print(f"There are {count:,} vectors with {dimensions:,} dimensions in the vector store")

    return vectorstore


# -----------------------------
# MAIN
# -----------------------------
if __name__ == "__main__":
    # Ensure folder structure exists for safety
    if not os.path.exists(FILE_PATH):
        raise FileNotFoundError(f"Could not find your markdown file at: {FILE_PATH}")

    chunks = fetch_and_chunk_document()
    create_embeddings(chunks)

    print("Ingestion complete 🚀")