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
from config_loader import load_config


CFG = load_config()
DATA_DIR = CFG["ingestion"]["data_dir"]

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=CFG["ingestion"]["chunk_size"], 
        chunk_overlap=CFG["ingestion"]["chunk_overlap"], 
        separators=["\n\n", "\n", ". ", " ", ""]
    )
    return text_splitter.split_documents(docs)

def enrich_chunks(chunks: list) -> list:
    MODEL_BASE = CFG["models"]["inference"]
    API_BASE = CFG["api"]["base_url"]
    API_VERSION = CFG["api"]["api_version"]
    
    def process_single_chunk(indexed_chunk):
        idx, chunk = indexed_chunk
        lm_index = idx % 8
        
        try:
            with dspy.context(lm=dspy.LM(f"{MODEL_BASE}:{lm_index}", api_base=API_BASE+API_VERSION)):
                response = IngestionAgent().forward(note=chunk.page_content)
                
                # This is now an object, not a string!
                metadata = response.answer.dict() 

        except Exception as e:
            print(f"⚠️ Failed for chunk {idx}: {e}")
            metadata = {"synopsis": "Summary failed", "tags": ["error"], "entities": []}

            chunk.metadata.update(metadata)
            return chunk
            

    enriched_results = []
    with ThreadPoolExecutor(max_workers=CFG["ingestion"]["max_workers"]) 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()