"""
PDF Ingestion & Vector Store Builder for Study Ocean RAG.

This module:
1. Recursively scans the PDF directory for all .pdf files
2. Extracts text from each PDF using PyMuPDF (faster & more reliable)
3. Splits text into overlapping chunks
4. Generates embeddings using a local sentence-transformer model
5. Stores embeddings in a FAISS vector store for fast retrieval
"""

import os
import sys
import time
import fitz  # PyMuPDF
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings

import config


def collect_pdf_paths(pdf_dir: str) -> list[str]:
    """Recursively collect all PDF file paths from the given directory."""
    pdf_paths = []
    for root, _, files in os.walk(pdf_dir):
        for file in files:
            if file.lower().endswith(".pdf"):
                pdf_paths.append(os.path.join(root, file))
    return sorted(pdf_paths)


def extract_text_from_pdf(pdf_path: str) -> str:
    """Extract all text from a PDF using PyMuPDF."""
    text = ""
    try:
        doc = fitz.open(pdf_path)
        for page_num in range(len(doc)):
            page = doc[page_num]
            text += page.get_text()
        doc.close()
    except Exception as e:
        print(f"  ⚠ Error reading {os.path.basename(pdf_path)}: {e}")
    return text


def build_documents(pdf_paths: list[str]) -> list:
    """
    Extract text from all PDFs and split into chunks.
    Each chunk carries metadata about its source PDF and page info.
    """
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=config.CHUNK_SIZE,
        chunk_overlap=config.CHUNK_OVERLAP,
        separators=["\n\n", "\n", ". ", " ", ""],
        length_function=len,
    )

    all_documents = []
    for i, pdf_path in enumerate(pdf_paths, 1):
        filename = os.path.basename(pdf_path)
        relative_path = os.path.relpath(pdf_path, config.PDF_DIR)
        # Derive subject from folder structure
        parts = relative_path.replace("\\", "/").split("/")
        subject = parts[0] if len(parts) > 1 else "General"

        print(f"  [{i}/{len(pdf_paths)}] Processing: {relative_path}")
        
        text = extract_text_from_pdf(pdf_path)
        if not text.strip():
            print(f"    ⚠ No text extracted (might be scanned/image PDF)")
            continue

        # Create chunks with rich metadata
        chunks = text_splitter.create_documents(
            texts=[text],
            metadatas=[{
                "source": relative_path,
                "filename": filename,
                "subject": subject,
            }]
        )
        all_documents.extend(chunks)

    return all_documents


def create_vector_store(documents: list, embeddings) -> FAISS:
    """Create a FAISS vector store from documents."""
    print(f"\n📊 Creating vector store with {len(documents)} chunks...")
    
    # Process in batches to avoid memory issues
    batch_size = 500
    vector_store = None
    
    for i in range(0, len(documents), batch_size):
        batch = documents[i:i + batch_size]
        batch_num = (i // batch_size) + 1
        total_batches = (len(documents) + batch_size - 1) // batch_size
        print(f"  Embedding batch {batch_num}/{total_batches} ({len(batch)} chunks)...")
        
        if vector_store is None:
            vector_store = FAISS.from_documents(batch, embeddings)
        else:
            batch_store = FAISS.from_documents(batch, embeddings)
            vector_store.merge_from(batch_store)
    
    return vector_store


def run_ingestion():
    """Main ingestion pipeline."""
    print("=" * 60)
    print("🎓 Study Ocean RAG — PDF Ingestion Pipeline")
    print("=" * 60)

    # 1. Collect PDFs
    print(f"\n📁 Scanning for PDFs in: {config.PDF_DIR}")
    pdf_paths = collect_pdf_paths(config.PDF_DIR)
    
    if not pdf_paths:
        print("❌ No PDF files found! Make sure your PDFs are in the data/pdfs directory.")
        print("   Extract StudyOceanPDFs.zip into data/pdfs/ first.")
        sys.exit(1)
    
    print(f"   Found {len(pdf_paths)} PDF files\n")

    # 2. Extract and chunk
    print("📄 Extracting text and creating chunks...")
    start_time = time.time()
    documents = build_documents(pdf_paths)
    extract_time = time.time() - start_time
    print(f"\n   ✅ Created {len(documents)} text chunks in {extract_time:.1f}s")

    if not documents:
        print("❌ No text could be extracted from any PDF. Exiting.")
        sys.exit(1)

    # 3. Initialize embeddings
    print(f"\n🧠 Loading embedding model: {config.EMBEDDING_MODEL}")
    embeddings = HuggingFaceEmbeddings(
        model_name=config.EMBEDDING_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True, "batch_size": 64},
    )

    # 4. Build vector store
    start_time = time.time()
    vector_store = create_vector_store(documents, embeddings)
    embed_time = time.time() - start_time
    print(f"   ✅ Vector store created in {embed_time:.1f}s")

    # 5. Save to disk
    os.makedirs(config.VECTOR_STORE_PATH, exist_ok=True)
    print(f"\n💾 Saving vector store to: {config.VECTOR_STORE_PATH}")
    vector_store.save_local(config.VECTOR_STORE_PATH)
    print("   ✅ Vector store saved successfully!")

    # Summary
    print("\n" + "=" * 60)
    print("✅ INGESTION COMPLETE")
    print(f"   PDFs processed:  {len(pdf_paths)}")
    print(f"   Chunks created:  {len(documents)}")
    print(f"   Extract time:    {extract_time:.1f}s")
    print(f"   Embedding time:  {embed_time:.1f}s")
    print(f"   Store location:  {config.VECTOR_STORE_PATH}")
    print("=" * 60)


if __name__ == "__main__":
    run_ingestion()