#!/usr/bin/env python3
"""
FinBot Ingestion Script
========================
Runs the full ingestion pipeline on a PDF document:
  1. Parse PDF → text + tables (PyMuPDF + pdfplumber)
  2. Chunk text → semantic chunks
  3. Embed chunks → local BGE embeddings → ChromaDB
  4. Index tables → JSON table store
  5. Render pages → high-quality slide PNG images for display + ColPali
  6. ColPali indexing → multi-vector patch embeddings per page
  7. Save all metadata

Usage:
  python ingest.py --pdf "../IR Chatbot/emiratesnbd_investor_presentation_2026_q1.pdf"
  python ingest.py --pdf path/to/doc.pdf --doc-id my_doc_id

ColPali note:
  First run downloads vidore/colpali-v1.2-merged (~7GB).
  Subsequent runs are fast (model cached by HuggingFace).
  Runs on Mac MPS (M-series) or CPU.
"""

import argparse
import asyncio
import logging
import json
import sys
import time
from pathlib import Path

# Add backend to path
sys.path.insert(0, str(Path(__file__).parent))

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(name)s  %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("finbot.ingest")

BASE_DIR    = Path(__file__).parent
DATA_DIR    = BASE_DIR / "data"
PAGES_DIR   = DATA_DIR / "pages"
COLPALI_DIR = DATA_DIR / "colpali_index"
CHROMA_DIR  = DATA_DIR / "chroma"
TABLES_FILE = DATA_DIR / "tables.json"


def ingest_document_sync(doc_id: str, pdf_path: str):
    """
    Synchronous ingestion pipeline.
    Called from CLI or background task.
    """
    pdf_path = Path(pdf_path)
    if not pdf_path.exists():
        raise FileNotFoundError(f"PDF not found: {pdf_path}")

    doc_name = pdf_path.stem.replace("_", " ").replace("-", " ").title()
    total_start = time.time()

    logger.info("=" * 60)
    logger.info(f"FinBot Ingestion Pipeline")
    logger.info(f"Document : {pdf_path.name}")
    logger.info(f"Doc ID   : {doc_id}")
    logger.info("=" * 60)

    # ── Step 1: Parse PDF ────────────────────────────────────────────
    logger.info("\n[1/6] Parsing PDF (text + tables)…")
    from services.ingestion.pdf_parser import parse_pdf, save_extraction
    t0 = time.time()
    extracted = parse_pdf(pdf_path, doc_id, doc_name)
    logger.info(
        f"      ✓ {extracted.total_pages} pages parsed in {time.time()-t0:.1f}s"
        f" | Tables on {sum(1 for p in extracted.pages if p.has_table)} pages"
    )

    # Save extraction JSON for debugging
    save_extraction(extracted, DATA_DIR / "processed")

    # ── Step 2: Chunk text ───────────────────────────────────────────
    logger.info("\n[2/6] Chunking text…")
    from services.ingestion.text_chunker import TextChunker
    t0 = time.time()
    chunker = TextChunker(chunk_size=300, chunk_overlap=60)
    page_texts = [
        {
            "page_number": p.page_number,
            "text": p.text,
            "section_heading": getattr(p, "section_heading", ""),
            "slide_number": getattr(p, "slide_number", p.page_number),
        }
        for p in extracted.pages if p.text.strip()
    ]
    chunks = chunker.chunk_document(page_texts, doc_id)
    logger.info(f"      ✓ {len(chunks)} chunks in {time.time()-t0:.1f}s")

    # ── Step 3: Embed + index text in ChromaDB ───────────────────────
    logger.info("\n[3/6] Embedding text chunks (local BGE model)…")
    from services.retrieval.text_retriever import TextRetriever
    t0 = time.time()
    chunks = chunker.embed_chunks(chunks)
    text_retriever = TextRetriever(persist_dir=CHROMA_DIR)
    n_indexed = text_retriever.index_chunks(chunks)
    logger.info(f"      ✓ {n_indexed} chunks indexed in ChromaDB | {time.time()-t0:.1f}s")

    # ── Step 4: Index tables ─────────────────────────────────────────
    logger.info("\n[4/6] Indexing tables…")
    from services.retrieval.table_retriever import TableRetriever
    t0 = time.time()
    table_retriever = TableRetriever(TABLES_FILE)
    n_tables = table_retriever.index_tables(extracted)
    logger.info(f"      ✓ {n_tables} tables indexed in {time.time()-t0:.1f}s")

    # ── Step 5: Render PDF pages to images ───────────────────────────
    logger.info("\n[5/6] Rendering PDF pages as high-quality slide images…")
    from services.ingestion.page_renderer import render_pdf_pages
    t0 = time.time()
    PAGES_DIR.mkdir(parents=True, exist_ok=True)
    page_records = render_pdf_pages(
        pdf_path=pdf_path,
        output_dir=PAGES_DIR,
        doc_id=doc_id,
    )
    logger.info(
        f"      ✓ {len(page_records)} pages rendered in {time.time()-t0:.1f}s"
        f" → {PAGES_DIR / doc_id / 'pages'}"
    )

    # ── Step 6: ColPali indexing ──────────────────────────────────────────
    logger.info("\n[6/6] Generating ColPali patch embeddings…")
    logger.info("      (First run downloads ~7GB model — cached after that)")
    logger.info("      Running on: Mac MPS or CPU (batch_size=1 to avoid OOM)")
    import os, torch
    if torch.backends.mps.is_available():
        os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0"
        logger.info("      MPS detected — disabled memory watermark limit")
    from services.ingestion.colpali_indexer import ColPaliIndexer
    t0 = time.time()
    colpali_indexer = ColPaliIndexer(store_dir=COLPALI_DIR)
    page_records = colpali_indexer.index_pages(
        page_records=page_records,
        doc_id=doc_id,
        batch_size=1,   # Always 1 on MPS — OOM otherwise
    )
    elapsed = time.time() - t0
    logger.info(
        f"      ✓ {len(page_records)} pages ColPali-indexed in {elapsed:.1f}s"
        f" (~{elapsed/len(page_records):.1f}s/page)"
    )

    # ── Save document metadata ────────────────────────────────────────
    docs_file = DATA_DIR / "documents.json"
    if docs_file.exists():
        with open(docs_file) as f:
            docs = json.load(f)
    else:
        docs = []

    # Remove old entry for this doc_id
    docs = [d for d in docs if d["doc_id"] != doc_id]
    docs.append({
        "doc_id": doc_id,
        "name": doc_name,
        "doc_type": extracted.metadata.get("doc_type", "Financial Document"),
        "period": extracted.metadata.get("period", "Unknown"),
        "institution": extracted.metadata.get("institution", "Emirates NBD"),
        "total_pages": extracted.total_pages,
        "status": "indexed",
        "filename": pdf_path.name,
        "chunks_indexed": n_indexed,
        "tables_indexed": n_tables,
        "colpali_pages": len(page_records),
    })

    DATA_DIR.mkdir(parents=True, exist_ok=True)
    with open(docs_file, "w") as f:
        json.dump(docs, f, indent=2)

    total_elapsed = time.time() - total_start
    logger.info("\n" + "=" * 60)
    logger.info(f"✅  Ingestion complete in {total_elapsed:.1f}s")
    logger.info(f"    Text chunks  : {n_indexed}")
    logger.info(f"    Tables       : {n_tables}")
    logger.info(f"    Pages imaged : {len(page_records)}")
    logger.info(f"    ColPali pages: {len(page_records)}")
    logger.info("=" * 60)
    logger.info("\nStart the backend: uvicorn app.main:app --reload --port 8000")
    logger.info("Start the frontend: npm run dev  (in finbot-ir-platform/)")


async def ingest_document(doc_id: str, pdf_path: str):
    """Async wrapper for background task usage."""
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, ingest_document_sync, doc_id, pdf_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="IRIS PDF Ingestion — text + tables + ColPali visual indexing"
    )
    parser.add_argument(
        "--pdf",
        default=None,
        help="Path to the PDF file to ingest (if omitted, scans the documents/ folder)",
    )
    parser.add_argument(
        "--doc-id",
        default=None,
        help="Document ID (default: derived from filename)",
    )
    parser.add_argument(
        "--skip-colpali",
        action="store_true",
        help="Skip ColPali indexing (faster, but no visual retrieval)",
    )
    args = parser.parse_args()

    if args.pdf:
        pdf = Path(args.pdf)
        doc_id = args.doc_id or pdf.stem.lower().replace(" ", "_").replace("-", "_")
        logger.info(f"Ingesting specific file: {pdf.name} with doc_id: {doc_id}")
        ingest_document_sync(doc_id, str(pdf))
    else:
        # Scan the documents/ directory for PDF files
        documents_dir = BASE_DIR.parent / "documents"
        if not documents_dir.exists():
            documents_dir.mkdir(parents=True, exist_ok=True)
            logger.info(f"Created documents/ folder. Please place PDF files in: {documents_dir}")
            sys.exit(0)

        pdf_files = list(documents_dir.glob("*.pdf"))
        if not pdf_files:
            logger.info(f"No PDF files found in documents/ folder: {documents_dir}")
            sys.exit(0)

        # Check existing metadata to skip already indexed files
        docs_file = DATA_DIR / "documents.json"
        indexed_filenames = set()
        if docs_file.exists():
            try:
                with open(docs_file) as f:
                    docs = json.load(f)
                    for d in docs:
                        if d.get("status") == "indexed":
                            indexed_filenames.add(d.get("filename"))
            except Exception as e:
                logger.warning(f"Could not load documents.json: {e}")

        logger.info(f"Scanning documents/ folder. Found {len(pdf_files)} PDF files.")
        new_files = [f for f in pdf_files if f.name not in indexed_filenames]

        if not new_files:
            logger.info("All PDF files in documents/ are already indexed. Nothing to do!")
            sys.exit(0)

        logger.info(f"Found {len(new_files)} new files to ingest.")
        for pdf in new_files:
            doc_id = pdf.stem.lower().replace(" ", "_").replace("-", "_")
            logger.info(f"\n>>> Starting ingestion for new file: {pdf.name} (ID: {doc_id})")
            try:
                ingest_document_sync(doc_id, str(pdf))
            except Exception as e:
                logger.error(f"❌ Failed to ingest {pdf.name}: {e}")