#!/usr/bin/env python3
"""
Parent-Child Chunking with Milvus using LangChain ParentDocumentRetriever.

Converts documents (PDF, XLSX, DOCX, PPTX) to markdown using Docling,
applies parent-child chunking via LangChain's ParentDocumentRetriever,
and stores in Milvus.

Chunk sizes are specified in TOKENS (using tiktoken cl100k_base encoding
for text-embedding-3-large model).

Required environment variables:
    AZURE_OPENAI_API_KEY
    AZURE_OPENAI_ENDPOINT
    AZURE_OPENAI_API_VERSION
    AZURE_OPENAI_EMBEDDING_DEPLOYMENT
    MILVUS_HOST
    MILVUS_PORT

Usage:
    python scripts/parent_child_milvus_ingestion.py input.pdf
    python scripts/parent_child_milvus_ingestion.py input.docx --parent-size 2000 --child-size 400
"""

import argparse
import os
import sys
from pathlib import Path

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

import tiktoken
from docling.document_converter import DocumentConverter
from dotenv import load_dotenv
from langchain_classic.retrievers import ParentDocumentRetriever
from langchain_classic.storage import InMemoryStore
from langchain_core.documents import Document
from langchain_milvus import Milvus
from langchain_openai import AzureOpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Load environment variables
load_dotenv()

COLLECTION_NAME = "test_parent_child_chunking"

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


def count_tokens(text: str) -> int:
    """
    Count tokens in text using tiktoken cl100k_base encoding.

    This encoding is used by text-embedding-3-large model.

    Args:
        text: Text to count tokens for

    Returns:
        Number of tokens
    """
    return len(TIKTOKEN_ENCODER.encode(text))


def get_embeddings() -> AzureOpenAIEmbeddings:
    """Create Azure OpenAI embeddings instance."""
    return AzureOpenAIEmbeddings(
        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"
        ),
    )


def get_milvus_connection_args() -> dict:
    """Get Milvus connection arguments from environment."""
    return {
        "host": os.getenv("MILVUS_HOST", "localhost"),
        "port": os.getenv("MILVUS_PORT", "19530"),
    }


def convert_to_markdown(file_path: Path) -> str:
    """
    Convert document to markdown using Docling.

    Args:
        file_path: Path to input document

    Returns:
        Markdown string
    """
    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...")
    converter = DocumentConverter()
    result = converter.convert(file_path)
    markdown = result.document.export_to_markdown()
    token_count = count_tokens(markdown)
    print(f"Converted to {len(markdown):,} characters ({token_count:,} tokens)")

    return markdown


def create_parent_document_retriever(
    parent_chunk_size: int = 2000,
    child_chunk_size: int = 400,
    drop_existing: bool = False,
) -> tuple[ParentDocumentRetriever, Milvus, InMemoryStore]:
    """
    Create a ParentDocumentRetriever with Milvus vector store.

    Args:
        parent_chunk_size: Size of parent chunks in TOKENS (default: 2000)
        child_chunk_size: Size of child chunks in TOKENS (default: 400)
        drop_existing: Whether to drop existing collection

    Returns:
        Tuple of (retriever, vectorstore, docstore)
    """
    embeddings = get_embeddings()
    connection_args = get_milvus_connection_args()

    # Drop existing collection if requested
    if drop_existing:
        from pymilvus import connections, utility

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

    # Create Milvus vector store for child chunks
    vectorstore = Milvus(
        embedding_function=embeddings,
        collection_name=COLLECTION_NAME,
        connection_args=connection_args,
        auto_id=True,
    )

    # InMemoryStore for parent documents
    # In production, consider using a persistent store like Redis
    docstore = InMemoryStore()

    # Parent splitter (larger chunks) - using TOKEN-based length function
    parent_splitter = RecursiveCharacterTextSplitter(
        chunk_size=parent_chunk_size,
        chunk_overlap=int(parent_chunk_size * 0.1),
        separators=["\n\n", "\n", ". ", " ", ""],
        length_function=count_tokens,  # Use tiktoken for token counting
    )

    # Child splitter (smaller chunks for embedding) - using TOKEN-based length function
    child_splitter = RecursiveCharacterTextSplitter(
        chunk_size=child_chunk_size,
        chunk_overlap=int(child_chunk_size * 0.1),
        separators=["\n\n", "\n", ". ", " ", ""],
        length_function=count_tokens,  # Use tiktoken for token counting
    )

    # Create the ParentDocumentRetriever
    retriever = ParentDocumentRetriever(
        vectorstore=vectorstore,
        docstore=docstore,
        child_splitter=child_splitter,
        parent_splitter=parent_splitter,
    )

    return retriever, vectorstore, docstore


def ingest_document(
    file_path: Path,
    parent_chunk_size: int = 2000,
    child_chunk_size: int = 400,
    drop_existing: bool = False,
) -> tuple[ParentDocumentRetriever, int, int]:
    """
    Ingest a document using parent-child chunking strategy.

    Args:
        file_path: Path to the document
        parent_chunk_size: Size of parent chunks in TOKENS
        child_chunk_size: Size of child chunks in TOKENS
        drop_existing: Whether to drop existing collection

    Returns:
        Tuple of (retriever, parent_count, child_count)
    """
    # Convert to markdown
    markdown = convert_to_markdown(file_path)

    # Create the retriever
    retriever, vectorstore, docstore = create_parent_document_retriever(
        parent_chunk_size=parent_chunk_size,
        child_chunk_size=child_chunk_size,
        drop_existing=drop_existing,
    )

    # Create a LangChain Document with metadata
    doc = Document(
        page_content=markdown,
        metadata={
            "file_name": file_path.name,
            "source": str(file_path),
        },
    )

    # Add document to retriever (this splits into parent/child chunks and indexes)
    print("Adding document to retriever (splitting and embedding)...")
    retriever.add_documents([doc])

    # Get counts
    # Child count from vector store
    child_count = vectorstore.col.num_entities if hasattr(vectorstore, "col") else "unknown"

    # Parent count from docstore
    parent_keys = list(docstore.yield_keys())
    parent_count = len(parent_keys)

    print(f"Indexed {parent_count} parent chunks and {child_count} child chunks")

    return retriever, parent_count, child_count


def test_retrieval(retriever: ParentDocumentRetriever, query: str, k: int = 3):
    """
    Test retrieval with the parent document retriever.

    Args:
        retriever: The configured retriever
        query: Search query
        k: Number of results
    """
    print(f"\nSearching for: {query}")
    print("-" * 60)

    # Get relevant documents (returns parent chunks based on child chunk similarity)
    docs = retriever.invoke(query)

    for i, doc in enumerate(docs[:k], 1):
        content_tokens = count_tokens(doc.page_content)
        print(f"\n--- Result {i} ---")
        print(f"Metadata: {doc.metadata}")
        print(f"Content: {len(doc.page_content)} chars, {content_tokens} tokens")
        preview = doc.page_content[:500] + "..." if len(doc.page_content) > 500 else doc.page_content
        print(f"Preview:\n{preview}")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Ingest documents into Milvus using LangChain ParentDocumentRetriever.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python scripts/parent_child_milvus_ingestion.py document.pdf
    python scripts/parent_child_milvus_ingestion.py report.docx --parent-size 2000
    python scripts/parent_child_milvus_ingestion.py data.xlsx --drop-collection
    python scripts/parent_child_milvus_ingestion.py doc.pdf --test-search "search query"

Note: Chunk sizes are in TOKENS (using tiktoken cl100k_base for text-embedding-3-large)
        """,
    )
    parser.add_argument(
        "input_file", type=Path, help="Input document (PDF, XLSX, DOCX, PPTX)"
    )
    parser.add_argument(
        "--parent-size",
        type=int,
        default=2000,
        help="Parent chunk size in TOKENS (default: 2000)",
    )
    parser.add_argument(
        "--child-size",
        type=int,
        default=400,
        help="Child chunk size in TOKENS (default: 400)",
    )
    parser.add_argument(
        "--drop-collection",
        action="store_true",
        help="Drop existing collection before ingestion",
    )
    parser.add_argument(
        "--test-search",
        type=str,
        help="Test search query after ingestion",
    )

    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)

    print(f"\nProcessing: {args.input_file}")
    print(f"Config: parent={args.parent_size} tokens, child={args.child_size} tokens")
    print(f"Collection: {COLLECTION_NAME}")
    print(f"Tokenizer: tiktoken cl100k_base (text-embedding-3-large)")
    print("-" * 60)

    # Ingest document
    retriever, parent_count, child_count = ingest_document(
        file_path=args.input_file,
        parent_chunk_size=args.parent_size,
        child_chunk_size=args.child_size,
        drop_existing=args.drop_collection,
    )

    print("-" * 60)
    print("Summary:")
    print(f"  Collection: {COLLECTION_NAME}")
    print(f"  Parent chunks: {parent_count}")
    print(f"  Child chunks (in Milvus): {child_count}")

    # Test search if requested
    if args.test_search:
        print(f"\n{'=' * 60}")
        test_retrieval(retriever, args.test_search)


if __name__ == "__main__":
    main()