import os
from dotenv import load_dotenv
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS

def run_embedding_pipeline(folder_path, index_name="faiss_index"):
    load_dotenv()
    api_key = os.getenv("OPENROUTER_API_KEY")
    base_url = "https://openrouter.ai/api/v1"
    
    if not api_key:
        print("❌ Error: OPENROUTER_API_KEY not found in .env file.")
        return

    # 1. Load Documents from Folder
    print(f"📂 Loading .txt files from: {folder_path}...")
    loader = DirectoryLoader(folder_path, glob="*.txt", loader_cls=TextLoader)
    documents = loader.load()
    
    if not documents:
        print("⚠️ No documents found. Check your folder path.")
        return

    # 2. Chunking
    print("✂️ Splitting documents into chunks...")
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
    texts = text_splitter.split_documents(documents)

    # 3. Create Embeddings using Qwen3-8B
    print("🧬 Generating embeddings (this may take a moment)...")
    embeddings = OpenAIEmbeddings(
        model="qwen/qwen3-embedding-8b",
        openai_api_key=api_key,
        openai_api_base=base_url,
        check_embedding_ctx_length=False
    )

    # 4. Save Vector Store Locally
    vectorstore = FAISS.from_documents(texts, embeddings)
    vectorstore.save_local(index_name)
    print(f"✅ Success! Vector database saved to folder: '{index_name}'")

if __name__ == "__main__":
    # Change "pages_text" to your actual folder name
    target_folder = "pages_text" 
    run_embedding_pipeline(target_folder)