import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import dspy
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from tqdm import tqdm

from embedding import LocalLMEmbeddings
from experts.ingestion_agent import IngestionAgent

DATA_DIR = "/home/cosmic/DnD"

def load_documents():
    docs = []
    data_path = Path(DATA_DIR)

    if not data_path.exists() or not data_path.is_dir():
        print(f"⚠️ Data directory '{DATA_DIR}' does not exist.")
        return docs

    for file_path in data_path.rglob("*.md"):
        try:
            loader = TextLoader(str(file_path))
            loaded_docs = loader.load()

            for doc in loaded_docs:
                # Ensure these keys are set before splitting
                doc.metadata["source"] = file_path.name
                doc.metadata["full_path"] = str(file_path.absolute())

            docs.extend(loaded_docs)
            print(f"✅ Loaded: {file_path.name}")
        except Exception as e:
            print(f"❌ Failed to load {file_path}: {e}")

    return docs

def chunk_documents(docs):
    # LangChain preserves metadata during splitting automatically
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=800, 
        chunk_overlap=100, 
        separators=["\n\n", "\n", ". ", " ", ""]
    )
    return text_splitter.split_documents(docs)

def enrich_chunks(chunks: list) -> list:
    MODEL_BASE = "lm_studio/qwen/qwen3-8b"
    API_BASE = "http://192.168.0.49:1234/v1/"
    
    def process_single_chunk(indexed_chunk):
        idx, chunk = indexed_chunk
        lm_index = idx % 8
        
        try:
            # Configure context for this specific thread
            with dspy.context(lm=dspy.LM(f"{MODEL_BASE}:{lm_index}", api_base=API_BASE)):
                # Pass the text, but we will update the original chunk object
                response = IngestionAgent().ingest(note=chunk.page_content) 
                
                answer = response.answer
                start = answer.find("{")
                end = answer.rfind("}") + 1
                metadata_extracted = json.loads(answer[start:end])
                
                # UPDATE: Put AI data in a sub-key to avoid overwriting 'source'
                chunk.metadata["enrichment"] = metadata_extracted
                # Also flatten tags for easier searching if needed
                if "tags" in metadata_extracted:
                    chunk.metadata["tags"] = metadata_extracted["tags"]

        except Exception as e:
            # If enrichment fails, we KEEP the chunk but flag the error
            # This ensures 'source' and 'full_path' are NEVER lost
            chunk.metadata["enrichment_error"] = str(e)
            chunk.metadata["tags"] = ["error"]

        return idx, chunk

    enriched_results = []
    with ThreadPoolExecutor(max_workers=8) as executor:
        # Wrap chunks in enumerate to keep track of order
        futures = [executor.submit(process_single_chunk, (i, c)) for i, c in enumerate(chunks)]

        for future in tqdm(as_completed(futures), total=len(chunks), desc="Enriching chunks"):
            enriched_results.append(future.result())

    # Sort by the index (first element of tuple) and return only the chunk
    enriched_results.sort(key=lambda x: x[0])
    return [item[1] for item in enriched_results]

def store_chunks_locally(chunks, db_path="./local_faiss_db"):
    embeddings_model = LocalLMEmbeddings(
        model="text-embedding-qwen3-embedding-8b",
        base_url="http://192.168.0.49:1234",
        batch_size=32,
    )

    print(f"Index creation started for {len(chunks)} chunks...")
    # FAISS.from_documents extracts metadata directly from the Document objects
    vectorstore = FAISS.from_documents(documents=chunks, embedding=embeddings_model)
    vectorstore.save_local(db_path)
    print(f"✅ Successfully stored in FAISS at '{db_path}'")
    return vectorstore

def main():
    docs = load_documents()
    if not docs: return
    
    chunks = chunk_documents(docs)
    enriched_chunks = enrich_chunks(chunks)
    store_chunks_locally(enriched_chunks)
    print("🎉 Ingestion complete!")

if __name__ == "__main__":
    main()