#!/usr/bin/env python3
"""
Query script for Parent-Child Chunking with Milvus.

Queries the Milvus collection using LangChain ParentDocumentRetriever
and saves results to a JSON file.

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_query.py "your search query"
    python scripts/parent_child_query.py "選定療養" -o results.json -k 5
"""

import argparse
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
from langchain_milvus import Milvus
from langchain_openai import AzureOpenAIEmbeddings

load_dotenv()

COLLECTION_NAME = "test_parent_child_chunking"


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 query_collection(
    query: str,
    top_k: int = 5,
    score_threshold: float | None = None,
) -> list[dict]:
    """
    Query the Milvus collection for similar documents.

    Args:
        query: Search query text
        top_k: Number of results to return
        score_threshold: Optional minimum similarity score (0-1 for cosine)

    Returns:
        List of result dictionaries with text, metadata, and score
    """
    embeddings = get_embeddings()
    connection_args = get_milvus_connection_args()

    # Connect to existing Milvus collection
    vectorstore = Milvus(
        embedding_function=embeddings,
        collection_name=COLLECTION_NAME,
        connection_args=connection_args,
    )

    # Perform similarity search with scores
    results = vectorstore.similarity_search_with_score(
        query=query,
        k=top_k,
    )

    # Format results
    formatted_results = []
    for doc, score in results:
        result = {
            "text": doc.page_content,
            "metadata": doc.metadata,
            "similarity_score": float(score),
        }

        # Apply score threshold if specified
        if score_threshold is None or score >= score_threshold:
            formatted_results.append(result)

    return formatted_results


def save_results(
    query: str,
    results: list[dict],
    output_path: Path,
):
    """
    Save query results to a JSON file.

    Args:
        query: Original search query
        results: List of result dictionaries
        output_path: Path to output JSON file
    """
    output = {
        "query": query,
        "timestamp": datetime.now().isoformat(),
        "collection": COLLECTION_NAME,
        "result_count": len(results),
        "results": results,
    }

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(output, f, indent=2, ensure_ascii=False)

    print(f"Results saved to: {output_path}")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Query Milvus parent-child collection and save results to JSON.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python scripts/parent_child_query.py "選定療養"
    python scripts/parent_child_query.py "診療報酬" -k 10 -o results.json
    python scripts/parent_child_query.py "search query" --threshold 0.7
        """,
    )
    parser.add_argument("query", type=str, help="Search query text")
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=Path("query_results.json"),
        help="Output JSON file (default: query_results.json)",
    )
    parser.add_argument(
        "-k",
        "--top-k",
        type=int,
        default=5,
        help="Number of results to return (default: 5)",
    )
    parser.add_argument(
        "--threshold",
        type=float,
        default=None,
        help="Minimum similarity score threshold (0-1)",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print results to console",
    )

    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)

    print(f"Query: {args.query}")
    print(f"Collection: {COLLECTION_NAME}")
    print(f"Top-K: {args.top_k}")
    print("-" * 60)

    # Query collection
    results = query_collection(
        query=args.query,
        top_k=args.top_k,
        score_threshold=args.threshold,
    )

    print(f"Found {len(results)} results")

    # Print results if verbose
    if args.verbose:
        print("\n" + "=" * 60)
        for i, result in enumerate(results, 1):
            print(f"\n--- Result {i} (score: {result['similarity_score']:.4f}) ---")
            print(f"Metadata: {result['metadata']}")
            preview = result["text"][:500] + "..." if len(result["text"]) > 500 else result["text"]
            print(f"Text:\n{preview}")

    # Save to JSON
    save_results(args.query, results, args.output)


if __name__ == "__main__":
    main()