#!/usr/bin/env python3
"""
Query script for Hierarchical Chunking with Milvus and persistent docstore.

Queries the Milvus collection using LlamaIndex AutoMergingRetriever
with parent nodes stored in PostgreSQL or Redis.

Required environment variables:
    AZURE_OPENAI_API_KEY
    AZURE_OPENAI_ENDPOINT
    AZURE_OPENAI_API_VERSION
    AZURE_OPENAI_EMBEDDING_DEPLOYMENT
    MILVUS_HOST
    MILVUS_PORT
    POSTGRES_URL (for postgres docstore)
    REDIS_URL (for redis docstore)

Usage:
    python scripts/hierarchical_query.py "search query" --docstore postgres
    python scripts/hierarchical_query.py "検索クエリ" --docstore postgres -o results.json
    python scripts/hierarchical_query.py "query" --docstore redis --top-k 10
"""

import argparse
import asyncio
import json
import os
import sys
from datetime import datetime
from pathlib import Path

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

from dotenv import load_dotenv

load_dotenv()

# LlamaIndex imports
from llama_index.core import Settings, StorageContext, VectorStoreIndex
from llama_index.core.retrievers import AutoMergingRetriever
from llama_index.core.schema import NodeRelationship
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.postgres import PostgresDocumentStore
from llama_index.storage.docstore.redis import RedisDocumentStore

from src.core.config import settings

# Use project settings for collection name
COLLECTION_NAME = settings.milvus_collection


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 create_docstore(docstore_type: str):
    """Create docstore based on type."""
    if docstore_type == "postgres":
        # Use project settings for database URL (same as Airflow tasks)
        postgres_url = "postgresql+asyncpg://textiq:textiq_dev_password@localhost:5432/textiq"
        print(f"Connecting to PostgreSQL docstore...")
        return PostgresDocumentStore.from_uri(
            uri=postgres_url,
            table_name="hierarchical_nodes",
            namespace="hierarchical",
            perform_setup=False,  # Don't create table, it should exist
            use_jsonb=True,
        )
    elif docstore_type == "redis":
        redis_url = os.getenv("REDIS_URL", "redis://localhost:6379")
        print(f"Connecting to 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
        return RedisDocumentStore.from_host_and_port(
            host=redis_host,
            port=redis_port,
            namespace="hierarchical_nodes",
        )
    else:
        raise ValueError(f"Docstore type '{docstore_type}' requires persistent storage. Use 'postgres' or 'redis'.")


async def query_hierarchical(
    query: str,
    docstore_type: str,
    similarity_top_k: int = 6,
    merge_threshold: float = 0.5,
) -> dict:
    """
    Query hierarchical index with AutoMergingRetriever.

    Args:
        query: Search query text
        docstore_type: Docstore backend - "postgres" or "redis"
        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"\nQuery: {query}")
    print(f"Settings: top_k={similarity_top_k}, merge_threshold={merge_threshold}")
    print("-" * 60)

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

    # Connect to Milvus
    milvus_conn = get_milvus_connection()
    print(f"Connecting to Milvus at {milvus_conn['host']}:{milvus_conn['port']}...")

    vector_store = MilvusVectorStore(
        uri=f"http://{milvus_conn['host']}:{milvus_conn['port']}",
        collection_name=COLLECTION_NAME,
        dim=3072,
    )

    # Connect to docstore
    docstore = create_docstore(docstore_type)

    # Create storage context
    storage_context = StorageContext.from_defaults(
        vector_store=vector_store,
        docstore=docstore,
    )

    # Create index from existing vector store
    print("Loading index from Milvus...")
    index = VectorStoreIndex.from_vector_store(
        vector_store=vector_store,
        storage_context=storage_context,
    )

    # 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
    print("\nSearching...")
    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

        # Extract children information from node relationships
        children_info = []
        if NodeRelationship.CHILD in node.relationships:
            child_refs = node.relationships[NodeRelationship.CHILD]
            # child_refs can be a single RelatedNodeInfo or a list
            if not isinstance(child_refs, list):
                child_refs = [child_refs]
            for child_ref in child_refs:
                child_data = {
                    "node_id": child_ref.node_id,
                }
                # Try to fetch child node text from docstore
                try:
                    child_node = docstore.get_document(child_ref.node_id)
                    if child_node:
                        child_data["text"] = child_node.text
                        child_data["char_count"] = len(child_node.text)
                        if child_node.metadata:
                            child_data["metadata"] = child_node.metadata
                except Exception:
                    # Child node may not be in docstore (leaf nodes are in vector store)
                    pass
                children_info.append(child_data)

        children_count = len(children_info) if children_info else 0

        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,
            "children": children_info if children_info else None,
            "children_count": children_count,
        }
        results.append(result)

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

    return {
        "query": query,
        "timestamp": datetime.now().isoformat(),
        "collection": COLLECTION_NAME,
        "docstore": docstore_type,
        "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="Query hierarchical chunks from Milvus with AutoMergingRetriever.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python scripts/hierarchical_query.py "search query" --docstore postgres
    python scripts/hierarchical_query.py "検索クエリ" --docstore postgres -o results.json
    python scripts/hierarchical_query.py "query" --docstore redis --top-k 10
    python scripts/hierarchical_query.py "query" --docstore postgres --merge-threshold 0.3
        """,
    )
    parser.add_argument("query", type=str, help="Search query text")
    parser.add_argument(
        "--docstore",
        type=str,
        choices=["postgres", "redis"],
        required=True,
        help="Docstore backend (postgres or redis)",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=10,
        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",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print full result text (not just preview)",
    )

    args = parser.parse_args()

    # 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)

    # Query
    results = await query_hierarchical(
        query=args.query,
        docstore_type=args.docstore,
        similarity_top_k=args.top_k,
        merge_threshold=args.merge_threshold,
    )

    # Save to JSON 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())