import os
import re
import glob
from llama_index.core import Document
from llama_index.core.node_parser import MarkdownNodeParser

import chromadb
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 1. Load all markdown files from markdown_data/
md_files = glob.glob("markdown_data/*.md")
parser = MarkdownNodeParser()
final_chunks = []

for filepath in md_files:
    with open(filepath, "r", encoding="utf-8") as f:
        content = f.read()

    filename = os.path.basename(filepath)
    doc_name_with_ext = os.path.splitext(filename)[0]

    title_match = re.match(r"^#\s+(.+)$", content, re.MULTILINE)
    doc_name = title_match.group(1).strip() if title_match else doc_name_with_ext

    year_match = re.match(r"(\d{4})\s", doc_name_with_ext)
    if not year_match:
        year_match = re.search(r"(\d{4})", doc_name)
    doc_year = int(year_match.group(1)) if year_match else 0

    doc = Document(
        text=content,
        metadata={
            "doc_year": doc_year,
            "document_name": doc_name,
        }
    )

    nodes = parser.get_nodes_from_documents([doc])

    for node in nodes:
        if not node.text.strip():
            continue

        header_path_str = node.metadata.get("header_path", "")
        headers = [h.strip() for h in header_path_str.split("/") if h.strip()]

        enriched_text_lines = []
        enriched_text_lines.append(f"Document: {node.metadata['document_name']}")
        sub_headers = headers[1:]
        if sub_headers:
            enriched_text_lines.append(f"Section: {' > '.join(sub_headers)}")
        enriched_text_lines.append("")
        enriched_text_lines.append(node.text.lstrip('#').strip())

        final_chunk_text = "\n".join(enriched_text_lines)

        node.text = final_chunk_text
        node.metadata = {
            "doc_year": node.metadata["doc_year"],
            "doc_name": node.metadata["document_name"],
        }

        final_chunks.append(node)

# --- Verification ---
print(f"Generated {len(final_chunks)} chunks from {len(md_files)} files:")
for chunk in final_chunks:
    print(f"  - [{chunk.metadata['doc_year']}] {chunk.text[:80]}...")

# Create vector database
# 1. Create a local Vector Database "Table" (Collection)
db = chromadb.PersistentClient(path="./chroma_db")
try:
    db.delete_collection("clinical_guidelines")
except Exception:
    pass
chroma_collection = db.create_collection("clinical_guidelines")

# 2. Assign Chroma to LlamaIndex as our Vector Storage layer
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 3. Input your "Rows" (The final_chunks list we processed in the previous step)
# This step automatically runs the text through an embedding model, 
# generates the vectors, and inserts them as rows into the database.
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
index = VectorStoreIndex(final_chunks, storage_context=storage_context, embed_model=embed_model)

print(f"Successfully inserted {len(final_chunks)} rows into the vector table.")