#!/usr/bin/env python3
"""
Hierarchical Chunking with LlamaIndex and Milvus.

Converts documents (PDF, XLSX, DOCX, PPTX) to markdown using Docling,
applies hierarchical chunking via LlamaIndex's HierarchicalNodeParser,
and stores in Milvus with AutoMergingRetriever support.

Features:
- Multi-level hierarchical chunking (configurable levels)
- Automatic parent-child node relationships
- AutoMergingRetriever for query-time merging
- Page/sheet-level metadata extraction (page_number, sheet_number, slide_number)
- DOCX/PPTX to PDF conversion via LibreOffice for reliable page numbers

Required environment variables:
    AZURE_OPENAI_API_KEY
    AZURE_OPENAI_ENDPOINT
    AZURE_OPENAI_API_VERSION
    AZURE_OPENAI_EMBEDDING_DEPLOYMENT
    MILVUS_HOST
    MILVUS_PORT
    REDIS_URL (optional, for Redis docstore)
    POSTGRES_URL (optional, for PostgreSQL docstore)

Optional dependencies:
    LibreOffice - Required for DOCX/PPTX files (converts to PDF for page extraction)
        macOS: brew install --cask libreoffice
        Ubuntu: sudo apt install libreoffice
        Windows: https://www.libreoffice.org/

Usage:
    python scripts/hierarchical_llamaindex_milvus.py input.pdf             # page_number metadata
    python scripts/hierarchical_llamaindex_milvus.py input.xlsx            # sheet_number metadata
    python scripts/hierarchical_llamaindex_milvus.py input.docx            # converts to PDF, page_number
    python scripts/hierarchical_llamaindex_milvus.py input.pptx            # converts to PDF, slide_number
    python scripts/hierarchical_llamaindex_milvus.py doc.pdf --query "search query"
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Optional

sys.path.insert(0, str(Path(__file__).parent.parent))

from dotenv import load_dotenv

load_dotenv()

import asyncio

# LlamaIndex imports
from llama_index.core import Document, Settings, StorageContext, VectorStoreIndex
from llama_index.core.node_parser import HierarchicalNodeParser, SentenceSplitter, get_leaf_nodes
import tiktoken

# Initialize tiktoken encoder for text-embedding-3-large (uses cl100k_base)
TIKTOKEN_ENCODER = tiktoken.get_encoding("cl100k_base")


def tiktoken_tokenizer(text: str) -> list[int]:
    """Tokenize text using tiktoken cl100k_base encoding."""
    return TIKTOKEN_ENCODER.encode(text)


def count_tokens(text: str) -> int:
    """Count tokens in text using tiktoken."""
    return len(TIKTOKEN_ENCODER.encode(text))


def find_libreoffice() -> Optional[str]:
    """Find LibreOffice executable path."""
    # Common LibreOffice paths
    candidates = [
        "libreoffice",  # Linux/macOS with PATH
        "soffice",  # Alternative name
        "/Applications/LibreOffice.app/Contents/MacOS/soffice",  # macOS
        "/usr/bin/libreoffice",  # Linux
        "/usr/local/bin/libreoffice",  # Linux alternative
        "C:\\Program Files\\LibreOffice\\program\\soffice.exe",  # Windows
        "C:\\Program Files (x86)\\LibreOffice\\program\\soffice.exe",  # Windows 32-bit
    ]

    for candidate in candidates:
        if shutil.which(candidate):
            return candidate

    return None


def convert_to_pdf_with_libreoffice(file_path: Path, output_dir: Optional[Path] = None) -> Path:
    """
    Convert DOCX/PPTX to PDF using LibreOffice CLI.

    Args:
        file_path: Path to input document (DOCX or PPTX)
        output_dir: Directory for output PDF (defaults to temp directory)

    Returns:
        Path to converted PDF file

    Raises:
        RuntimeError: If LibreOffice is not found or conversion fails
    """
    libreoffice_path = find_libreoffice()
    if not libreoffice_path:
        raise RuntimeError(
            "LibreOffice not found. Install LibreOffice for DOCX/PPTX to PDF conversion.\n"
            "  macOS: brew install --cask libreoffice\n"
            "  Ubuntu: sudo apt install libreoffice\n"
            "  Windows: Download from https://www.libreoffice.org/"
        )

    if output_dir is None:
        output_dir = Path(tempfile.mkdtemp(prefix="libreoffice_convert_"))

    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Converting {file_path.name} to PDF using LibreOffice...")

    # LibreOffice CLI command for PDF conversion
    cmd = [
        libreoffice_path,
        "--headless",
        "--convert-to", "pdf",
        "--outdir", str(output_dir),
        str(file_path),
    ]

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=300,  # 5 minute timeout
        )

        if result.returncode != 0:
            raise RuntimeError(f"LibreOffice conversion failed: {result.stderr}")

        # Find the output PDF
        pdf_name = file_path.stem + ".pdf"
        pdf_path = output_dir / pdf_name

        if not pdf_path.exists():
            raise RuntimeError(f"PDF not created: {pdf_path}")

        print(f"Converted to PDF: {pdf_path}")
        return pdf_path

    except subprocess.TimeoutExpired:
        raise RuntimeError(f"LibreOffice conversion timed out for {file_path.name}")


from llama_index.core.retrievers import AutoMergingRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
from llama_index.vector_stores.milvus import MilvusVectorStore
from llama_index.storage.docstore.redis import RedisDocumentStore
from llama_index.storage.docstore.postgres import PostgresDocumentStore

# Docling for document conversion
from docling.document_converter import DocumentConverter

# Milvus for collection management
from pymilvus import connections, utility

COLLECTION_NAME = "hierarchical_chunking"


def get_azure_embedding():
    """Create Azure OpenAI embedding model."""
    return AzureOpenAIEmbedding(
        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
        api_key=os.getenv("AZURE_OPENAI_API_KEY"),
        api_version=os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-15-preview"),
        azure_deployment=os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT", "text-embedding-3-large"),
        model="text-embedding-3-large",
        embed_batch_size=16,
    )


def get_milvus_connection():
    """Get Milvus connection parameters."""
    return {
        "host": "localhost",
        "port": int(os.getenv("MILVUS_PORT", "19530")),
    }


def convert_to_markdown(file_path: Path) -> tuple[str, dict]:
    """
    Convert document to markdown using Docling.

    Args:
        file_path: Path to input document

    Returns:
        Tuple of (markdown_text, metadata)
    """
    supported = {".pdf", ".xlsx", ".xls", ".docx", ".pptx"}
    suffix = file_path.suffix.lower()

    if suffix not in supported:
        raise ValueError(f"Unsupported file type: {suffix}. Supported: {supported}")

    print(f"Converting {file_path.name} to markdown using Docling...")
    converter = DocumentConverter()
    result = converter.convert(file_path)
    markdown = result.document.export_to_markdown()

    token_count = count_tokens(markdown)
    metadata = {
        "file_name": file_path.name,
        "file_path": str(file_path),
        "file_type": suffix,
        "char_count": len(markdown),
        "token_count": token_count,
    }

    print(f"Converted to {len(markdown):,} characters ({token_count:,} tokens)")
    return markdown, metadata


def convert_to_markdown_with_pages(file_path: Path) -> list[tuple[str, dict]]:
    """
    Convert document to markdown with page/sheet-level metadata using Docling.

    For PDFs: extracts content per page with page_number metadata.
    For Excel: extracts content per sheet with sheet_number and sheet_name metadata.
    For DOCX/PPTX: converts to PDF first using LibreOffice for reliable page numbers.

    Args:
        file_path: Path to input document

    Returns:
        List of tuples: [(markdown_text, metadata), ...] one per page/sheet
    """
    supported = {".pdf", ".xlsx", ".xls", ".docx", ".pptx"}
    suffix = file_path.suffix.lower()

    if suffix not in supported:
        raise ValueError(f"Unsupported file type: {suffix}. Supported: {supported}")

    # Track if we created a temp PDF that needs cleanup
    temp_pdf_dir = None
    processing_path = file_path

    # Convert DOCX/PPTX to PDF first for reliable page numbers
    if suffix in {".docx", ".pptx"}:
        print(f"DOCX/PPTX detected - converting to PDF for reliable page metadata...")
        temp_pdf_dir = Path(tempfile.mkdtemp(prefix="hierarchical_pdf_"))
        processing_path = convert_to_pdf_with_libreoffice(file_path, temp_pdf_dir)

    try:
        print(f"Converting {processing_path.name} to markdown with page/sheet metadata...")
        converter = DocumentConverter()
        result = converter.convert(processing_path)
        doc = result.document

        pages_data = []
        # Always use original file info in metadata
        base_metadata = {
            "file_name": file_path.name,
            "file_path": str(file_path),
            "file_type": suffix,
        }

        # Track if converted from another format
        if suffix in {".docx", ".pptx"}:
            base_metadata["converted_from"] = suffix
            base_metadata["converted_to"] = ".pdf"

        # Get number of pages/sheets
        num_pages = doc.num_pages() if hasattr(doc, 'num_pages') and callable(doc.num_pages) else 0

        if num_pages > 0:
            # Document has pages (PDF) or sheets (Excel)
            print(f"Document has {num_pages} pages/sheets")

            for page_no in range(1, num_pages + 1):
                try:
                    # Export markdown for specific page (1-based index)
                    page_markdown = doc.export_to_markdown(page_no=page_no)

                    if not page_markdown or not page_markdown.strip():
                        print(f"  Page {page_no}: empty, skipping")
                        continue

                    token_count = count_tokens(page_markdown)

                    page_metadata = {
                        **base_metadata,
                        "total_pages": num_pages,
                    }

                    # For Excel files, use sheet_number only
                    if suffix in {".xlsx", ".xls"}:
                        page_metadata["sheet_number"] = page_no
                        # Try to get sheet name from document structure
                        sheet_name = _get_sheet_name(doc, page_no)
                        if sheet_name:
                            page_metadata["sheet_name"] = sheet_name
                    else:
                        # For PDF, DOCX, PPTX: use start_page/end_page
                        page_metadata["start_page"] = page_no
                        page_metadata["end_page"] = page_no

                    pages_data.append((page_markdown, page_metadata))
                    print(f"  Page {page_no}: {len(page_markdown):,} chars ({token_count:,} tokens)")

                except Exception as e:
                    print(f"  Page {page_no}: error extracting - {e}")
                    continue

        if not pages_data:
            # Fallback: no page structure or extraction failed, use full document
            print("No page structure found or extraction failed, using full document")
            markdown = doc.export_to_markdown()
            token_count = count_tokens(markdown)
            metadata = {
                **base_metadata,
                "total_pages": 1,
            }
            # For Excel, use sheet_number; for others use start_page/end_page
            if suffix in {".xlsx", ".xls"}:
                metadata["sheet_number"] = 1
            else:
                metadata["start_page"] = 1
                metadata["end_page"] = 1
            pages_data.append((markdown, metadata))

        total_chars = sum(len(text) for text, _ in pages_data)
        total_tokens = sum(count_tokens(text) for text, _ in pages_data)
        print(f"Total: {len(pages_data)} pages, {total_chars:,} chars ({total_tokens:,} tokens)")

        return pages_data

    finally:
        # Cleanup temp PDF directory if created
        if temp_pdf_dir and temp_pdf_dir.exists():
            shutil.rmtree(temp_pdf_dir, ignore_errors=True)
            print(f"Cleaned up temp directory: {temp_pdf_dir}")


def _get_sheet_name(doc, page_no: int) -> Optional[str]:
    """
    Try to extract sheet name from DoclingDocument for Excel files.

    Args:
        doc: DoclingDocument instance
        page_no: 1-based page/sheet number

    Returns:
        Sheet name if found, None otherwise
    """
    try:
        # Docling stores sheets as SECTION groups
        # Try to get the group label for this page/sheet
        from docling_core.types.doc import GroupLabel

        section_groups = [g for g in doc.groups if g.label == GroupLabel.SECTION]
        if section_groups and page_no <= len(section_groups):
            group = section_groups[page_no - 1]
            # The group might have a name or we can derive it
            if hasattr(group, 'name') and group.name:
                name = group.name
                # Docling adds "sheet: " prefix, remove it
                if name.lower().startswith("sheet: "):
                    name = name[7:]  # Remove "sheet: " prefix
                return name
    except Exception:
        pass

    return None


async def create_hierarchical_index(
    markdown: str,
    metadata: dict,
    chunk_sizes: list[int],
    drop_existing: bool = False,
    docstore_type: str = "memory",
) -> tuple[VectorStoreIndex, StorageContext, list]:
    """
    Create hierarchical index with Milvus vector store.

    Args:
        markdown: Document content in markdown format
        metadata: Document metadata
        chunk_sizes: List of chunk sizes for each hierarchy level
        drop_existing: Whether to drop existing collection
        docstore_type: Docstore backend - "memory", "redis", or "postgres"

    Returns:
        Tuple of (index, storage_context, all_nodes)
    """
    # Convert single document to pages format for unified processing
    pages_data = [(markdown, metadata)]
    return await create_hierarchical_index_from_pages(
        pages_data=pages_data,
        chunk_sizes=chunk_sizes,
        drop_existing=drop_existing,
        docstore_type=docstore_type,
    )


async def create_hierarchical_index_from_pages(
    pages_data: list[tuple[str, dict]],
    chunk_sizes: list[int],
    drop_existing: bool = False,
    docstore_type: str = "memory",
) -> tuple[VectorStoreIndex, StorageContext, list]:
    """
    Create hierarchical index with Milvus vector store from page-level documents.

    Args:
        pages_data: List of (markdown, metadata) tuples, one per page/sheet
        chunk_sizes: List of chunk sizes for each hierarchy level
        drop_existing: Whether to drop existing collection
        docstore_type: Docstore backend - "memory", "redis", or "postgres"

    Returns:
        Tuple of (index, storage_context, all_nodes)
    """
    milvus_conn = get_milvus_connection()

    # Drop existing collection if requested
    if drop_existing:
        connections.connect(
            alias="default",
            host=milvus_conn["host"],
            port=milvus_conn["port"],
        )
        if utility.has_collection(COLLECTION_NAME):
            print(f"Dropping existing collection: {COLLECTION_NAME}")
            utility.drop_collection(COLLECTION_NAME)
        connections.disconnect("default")

    # Configure embedding model
    embed_model = get_azure_embedding()
    Settings.embed_model = embed_model

    # Create LlamaIndex Documents with page/sheet metadata
    documents = []
    for page_markdown, page_metadata in pages_data:
        doc = Document(
            text=page_markdown,
            metadata=page_metadata,
            metadata_seperator="\n",
            metadata_template="{key}: {value}",
            text_template="Metadata:\n{metadata_str}\n\nContent:\n{content}",
        )
        documents.append(doc)

    print(f"Created {len(documents)} LlamaIndex Documents with page/sheet metadata")

    # Create HierarchicalNodeParser with token-based chunking (tiktoken cl100k_base)
    print(f"Creating hierarchical nodes with chunk sizes (tokens): {chunk_sizes}")

    # Create custom SentenceSplitters with tiktoken tokenizer for each level
    # Using 20% overlap and markdown-aware paragraph separator to avoid mid-sentence cuts
    node_parser_map = {}
    for chunk_size in chunk_sizes:
        splitter = SentenceSplitter(
            chunk_size=chunk_size,
            chunk_overlap=int(chunk_size * 0.2),  # 20% overlap for better context
            tokenizer=tiktoken_tokenizer,
            paragraph_separator="\n\n",  # Respect markdown paragraph breaks
            secondary_chunking_regex="[^,.;!?]+[,.;!?]?",  # Split on sentence boundaries
        )
        node_parser_map[f"chunk_size_{chunk_size}"] = splitter

    node_parser = HierarchicalNodeParser(
        node_parser_ids=[f"chunk_size_{s}" for s in chunk_sizes],
        node_parser_map=node_parser_map,
    )

    # Parse documents into hierarchical nodes (metadata is automatically inherited)
    nodes = node_parser.get_nodes_from_documents(documents)
    print(f"Created {len(nodes)} total nodes")

    # Add chunk-level char_count and token_count to each node's metadata
    for node in nodes:
        node.metadata["char_count"] = len(node.text)
        node.metadata["token_count"] = count_tokens(node.text)

    # Get leaf nodes (smallest chunks) for vector indexing
    leaf_nodes = get_leaf_nodes(nodes)
    print(f"Leaf nodes for indexing: {len(leaf_nodes)}")

    # Count nodes by level
    level_counts = {}
    for node in nodes:
        # Estimate level based on token count
        token_count = count_tokens(node.text)
        for i, size in enumerate(chunk_sizes):
            if token_count > size * 0.7:  # Approximate
                level = i
                break
        else:
            level = len(chunk_sizes) - 1
        level_counts[level] = level_counts.get(level, 0) + 1

    print("Nodes by level:")
    for level, count in sorted(level_counts.items()):
        print(f"  Level {level}: {count} nodes")

    # Create Milvus vector store
    vector_store = MilvusVectorStore(
        uri=f"http://{milvus_conn['host']}:{milvus_conn['port']}",
        collection_name=COLLECTION_NAME,
        dim=3072,  # text-embedding-3-large dimension
        overwrite=drop_existing,
    )

    # Create docstore based on type
    if docstore_type == "redis":
        redis_url = "redis://localhost:6379"
        print(f"Using Redis docstore: {redis_url}")
        from urllib.parse import urlparse
        parsed = urlparse(redis_url)
        redis_host = parsed.hostname or "localhost"
        redis_port = parsed.port or 6379
        docstore = RedisDocumentStore.from_host_and_port(
            host=redis_host,
            port=redis_port,
            namespace="hierarchical_nodes",
        )
    elif docstore_type == "postgres":
        postgres_url = "postgresql+asyncpg://textiq:textiq_dev_password@localhost:5432/textiq"
        if not postgres_url:
            raise ValueError("POSTGRES_URL environment variable required for postgres docstore")
        print(f"Using PostgreSQL docstore")
        docstore = PostgresDocumentStore.from_uri(
            uri=postgres_url,
            table_name="hierarchical_nodes",
            namespace="hierarchical",
            perform_setup=True,
            use_jsonb=True,
        )
    else:
        print("Using in-memory docstore")
        docstore = SimpleDocumentStore()

    # Create storage context with docstore for all nodes
    storage_context = StorageContext.from_defaults(
        vector_store=vector_store,
        docstore=docstore,
    )

    # Add ALL nodes to docstore (for parent lookup during retrieval)
    storage_context.docstore.add_documents(nodes)

    # Create index with ONLY leaf nodes (for vector search)
    print("Creating vector index with leaf nodes...")
    index = VectorStoreIndex(
        leaf_nodes,
        storage_context=storage_context,
        show_progress=True,
    )

    print(f"Index created successfully in collection: {COLLECTION_NAME}")

    return index, storage_context, nodes


def query_with_auto_merging(
    index: VectorStoreIndex,
    storage_context: StorageContext,
    query: str,
    similarity_top_k: int = 6,
    merge_threshold: float = 0.5,
) -> dict:
    """
    Query the index with AutoMergingRetriever.

    Args:
        index: VectorStoreIndex
        storage_context: StorageContext with docstore
        query: Search query
        similarity_top_k: Number of leaf nodes to retrieve
        merge_threshold: Threshold for merging children to parent

    Returns:
        Query results with merged nodes
    """
    print(f"\nQuerying: {query}")
    print(f"Settings: top_k={similarity_top_k}, merge_threshold={merge_threshold}")
    print("-" * 60)

    # Create base retriever
    base_retriever = index.as_retriever(similarity_top_k=similarity_top_k)

    # Create AutoMergingRetriever
    retriever = AutoMergingRetriever(
        base_retriever,
        storage_context,
        simple_ratio_thresh=merge_threshold,
        verbose=True,
    )

    # Retrieve nodes
    retrieved_nodes = retriever.retrieve(query)

    print(f"\nRetrieved {len(retrieved_nodes)} nodes after merging:")

    results = []
    for i, node_with_score in enumerate(retrieved_nodes, 1):
        node = node_with_score.node
        score = node_with_score.score

        result = {
            "rank": i,
            "score": float(score) if score else None,
            "text": node.text,
            "char_count": len(node.text),
            "metadata": node.metadata,
            "node_id": node.node_id,
        }
        results.append(result)

        print(f"\n--- Result {i} (score: {f'{score:.4f}' if score else 'N/A'}) ---")
        print(f"Chars: {len(node.text)}, Metadata: {node.metadata}")
        preview = node.text[:200] + "..." if len(node.text) > 200 else node.text
        print(f"Preview: {preview}")

    return {
        "query": query,
        "settings": {
            "similarity_top_k": similarity_top_k,
            "merge_threshold": merge_threshold,
        },
        "result_count": len(results),
        "results": results,
    }


async def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Hierarchical chunking with LlamaIndex and Milvus.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python scripts/hierarchical_llamaindex_milvus.py document.pdf      # Extracts page_number metadata
    python scripts/hierarchical_llamaindex_milvus.py doc.xlsx          # Extracts sheet_number/sheet_name
    python scripts/hierarchical_llamaindex_milvus.py doc.docx          # Converts to PDF, extracts page_number
    python scripts/hierarchical_llamaindex_milvus.py doc.pptx          # Converts to PDF, extracts slide_number
    python scripts/hierarchical_llamaindex_milvus.py doc.pdf --chunk-sizes 2048 512 256
    python scripts/hierarchical_llamaindex_milvus.py doc.pdf --drop-collection
    python scripts/hierarchical_llamaindex_milvus.py doc.pdf --query "search query"
    python scripts/hierarchical_llamaindex_milvus.py doc.pdf --query "query" --output results.json
        """,
    )
    parser.add_argument(
        "input_file", type=Path, help="Input document (PDF, XLSX, DOCX, PPTX)"
    )
    parser.add_argument(
        "--chunk-sizes",
        type=int,
        nargs="+",
        default=[2048, 512, 256],
        help="Chunk sizes in TOKENS for hierarchy levels (default: 2048 512 256)",
    )
    parser.add_argument(
        "--drop-collection",
        action="store_true",
        help="Drop existing collection before indexing",
    )
    parser.add_argument(
        "--docstore",
        type=str,
        choices=["memory", "redis", "postgres"],
        default="memory",
        help="Docstore backend: memory (default), redis (REDIS_URL), postgres (POSTGRES_URL)",
    )
    parser.add_argument(
        "--query",
        type=str,
        help="Query to test after indexing",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=6,
        help="Number of leaf nodes to retrieve (default: 6)",
    )
    parser.add_argument(
        "--merge-threshold",
        type=float,
        default=0.5,
        help="Threshold for auto-merging (default: 0.5)",
    )
    parser.add_argument(
        "-o", "--output",
        type=Path,
        help="Output JSON file for query results",
    )

    args = parser.parse_args()

    # Validate input file
    if not args.input_file.exists():
        print(f"Error: Input file not found: {args.input_file}")
        sys.exit(1)

    # Validate environment
    required_env = [
        "AZURE_OPENAI_API_KEY",
        "AZURE_OPENAI_ENDPOINT",
    ]
    missing = [e for e in required_env if not os.getenv(e)]
    if missing:
        print(f"Error: Missing environment variables: {', '.join(missing)}")
        sys.exit(1)

    # Validate chunk sizes
    if len(args.chunk_sizes) < 2:
        print("Error: At least 2 chunk sizes required for hierarchy")
        sys.exit(1)

    # Sort chunk sizes descending
    chunk_sizes = sorted(args.chunk_sizes, reverse=True)

    print(f"\nProcessing: {args.input_file}")
    print(f"Chunk sizes (hierarchy): {chunk_sizes}")
    print(f"Collection: {COLLECTION_NAME}")
    print(f"Docstore: {args.docstore}")
    print("=" * 60)

    # Convert to markdown with page/sheet-level metadata
    pages_data = convert_to_markdown_with_pages(args.input_file)

    # Create hierarchical index from pages
    index, storage_context, nodes = await create_hierarchical_index_from_pages(
        pages_data=pages_data,
        chunk_sizes=chunk_sizes,
        drop_existing=args.drop_collection,
        docstore_type=args.docstore,
    )

    print("=" * 60)
    print("Summary:")
    print(f"  Collection: {COLLECTION_NAME}")
    print(f"  Total nodes: {len(nodes)}")
    print(f"  Leaf nodes (indexed): {len(get_leaf_nodes(nodes))}")
    print(f"  Hierarchy levels: {len(chunk_sizes)}")

    # Query if requested
    if args.query:
        print("\n" + "=" * 60)
        results = query_with_auto_merging(
            index=index,
            storage_context=storage_context,
            query=args.query,
            similarity_top_k=args.top_k,
            merge_threshold=args.merge_threshold,
        )

        # Save results if output specified
        if args.output:
            with open(args.output, "w", encoding="utf-8") as f:
                json.dump(results, f, indent=2, ensure_ascii=False)
            print(f"\nResults saved to: {args.output}")


if __name__ == "__main__":
    asyncio.run(main())