import os
from typing import List
from llama_index.core import VectorStoreIndex, Document, StorageContext, load_index_from_storage
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from langchain.text_splitter import RecursiveCharacterTextSplitter

EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"  
DOCS_FOLDER = "data/docs"
SAMPLE_FILE = "data/sample.txt"
STORAGE_DIR = "data/storage"

def load_docs_from_folder(folder_path=DOCS_FOLDER) -> List[Document]:
    documents = []
    if not os.path.exists(folder_path):
        return documents

    for fname in os.listdir(folder_path):
        path = os.path.join(folder_path, fname)
        if fname.endswith(".txt"):
            with open(path, "r", encoding="utf-8") as f:
                content = f.read().strip()
                if content:
                    documents.append(Document(text=content))
        elif fname.endswith(".pdf"):
            try:
                from llama_index.readers.pdf import PDFReader
                pdf_docs = PDFReader().load_data(path)
                documents.extend([Document(text=doc.text) for doc in pdf_docs])
            except ImportError:
                pass

    if not documents and os.path.exists(SAMPLE_FILE):
        with open(SAMPLE_FILE, "r", encoding="utf-8") as f:
            sample_text = f.read().strip()
        documents = [Document(text=sample_text)]

    if not documents:
        documents = [Document(text=" ")]

    return documents

def load_vectorstore_from_text(text: str = None, persist: bool = True) -> VectorStoreIndex:
    embed_model = HuggingFaceEmbedding(model_name=EMBED_MODEL)

    if persist and text is None and os.path.exists(STORAGE_DIR):
        storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR)
        try:
            return load_index_from_storage(storage_context)
        except Exception as e:
            print(f"Failed to load index from storage, rebuilding: {e}")

    documents = [Document(text=text.strip())] if text else load_docs_from_folder()

    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
    chunk_docs = []
    for doc in documents:
        for chunk in splitter.split_text(doc.text):
            chunk_docs.append(Document(text=chunk))
            
    index = VectorStoreIndex.from_documents(chunk_docs, embed_model=embed_model)

    if persist:
        os.makedirs(STORAGE_DIR, exist_ok=True)
        index.storage_context.persist(persist_dir=STORAGE_DIR)
    return index

def rebuild_vectorstore(folder_path=DOCS_FOLDER) -> VectorStoreIndex:
    documents = load_docs_from_folder(folder_path)
    if not documents:
        raise ValueError("No documents found in folder!")
    return load_vectorstore_from_text(persist=True)

def get_retriever(text: str = None) -> VectorStoreIndex:
    try:
        return load_vectorstore_from_text(text=text, persist=True)
    except Exception as e:
        print(f"Vectorstore load failed, rebuilding locally: {e}")
        return rebuild_vectorstore()