#!/usr/bin/env python3
"""Manual test script for Q&A endpoint.

This script demonstrates how to query documents using the Q&A system:
    - Semantic retrieval from Milvus vector store
    - Answer generation using Azure OpenAI
    - Source references with exact locations
"""

import hashlib
import asyncio
from uuid import UUID

from src.db.session import get_sync_session
from src.db.models import ApiKey, Document
from src.retrieval.semantic_retriever import SemanticRetriever
from src.retrieval.answer_generator import AnswerGenerator


def setup_test_api_key() -> str:
    """Create or get test API key.

    Returns:
        The plain text API key
    """
    test_key = "test-api-key-12345"
    key_hash = hashlib.sha256(test_key.encode()).hexdigest()

    with get_sync_session() as session:
        # Check if key exists
        existing = session.query(ApiKey).filter(ApiKey.key_hash == key_hash).first()

        if not existing:
            # Create new API key
            api_key = ApiKey(
                name="Test Key",
                key_hash=key_hash,
                is_active=True,
                rate_limit_per_minute=100
            )
            session.add(api_key)
            session.commit()
            print(f"[OK] Created API key: {test_key}")
        else:
            print(f"[OK] Using existing API key: {test_key}")

    return test_key


async def test_query(query: str, document_ids: list[str] = None):
    """Test a Q&A query.

    Args:
        query: Natural language question
        document_ids: Optional list of document IDs to filter
    """
    print(f"\n{'='*70}")
    print(f"QUERY: {query}")
    print(f"{'='*70}")

    if document_ids:
        print(f"Filtering to documents: {document_ids}")

    try:
        # Step 1: Retrieve relevant sources
        print(f"\n[1/2] Retrieving relevant sources from Milvus...")
        retriever = SemanticRetriever(top_k=5)
        sources, top_similarity = await retriever.retrieve(
            query=query,
            document_ids=[UUID(did) for did in document_ids] if document_ids else None
        )

        print(f"  ✓ Found {len(sources)} relevant sources")
        print(f"  Top similarity score: {top_similarity:.4f}")

        # Show sources
        if sources:
            print(f"\n  Sources:")
            for i, source in enumerate(sources, 1):
                print(f"    {i}. {source.file} - {source.sheet} @ {source.location}")
                print(f"       Context: {source.context[:100]}...")

        # Step 2: Generate answer
        print(f"\n[2/2] Generating answer using Azure OpenAI GPT-4...")
        generator = AnswerGenerator()
        answer, confidence = await generator.generate(
            query=query,
            sources=sources,
            top_similarity=top_similarity
        )

        print(f"  ✓ Answer generated (confidence: {confidence:.2f})")

        # Display results
        print(f"\n{'='*70}")
        print(f"ANSWER:")
        print(f"{'='*70}")
        print(f"{answer}")
        print(f"\nConfidence: {confidence:.2f}")
        print(f"Sources: {len(sources)}")

        if sources:
            print(f"\n{'='*70}")
            print(f"SOURCE REFERENCES:")
            print(f"{'='*70}")
            for i, source in enumerate(sources, 1):
                print(f"\n{i}. File: {source.file}")
                print(f"   Sheet: {source.sheet}")
                print(f"   Location: {source.location}")
                print(f"   Context: {source.context}")

        return answer, sources, confidence

    except Exception as e:
        print(f"\n✗ Query failed: {e}")
        import traceback
        traceback.print_exc()
        return None, [], 0.0


async def main():
    """Run Q&A tests."""
    import sys

    print("="*70)
    print("Q&A SYSTEM TEST")
    print("="*70)

    # Setup API key
    api_key = setup_test_api_key()

    # Get completed documents
    with get_sync_session() as session:
        docs = session.query(Document).filter(
            Document.status == 'completed'
        ).order_by(Document.created_at.desc()).all()

        if not docs:
            print("\n✗ No completed documents found.")
            print("  Please upload and process a document first using:")
            print("    python test_pipeline_manual.py")
            return

        print(f"\n[OK] Found {len(docs)} completed document(s)")
        print(f"\nAvailable documents:")
        for i, doc in enumerate(docs, 1):
            print(f"  {i}. {doc.filename}")
            print(f"       ID: {doc.id}")
            print(f"       Records: {doc.record_count}, Sheets: {doc.sheet_count}")

    # Select document or all
    print(f"\nTest options:")
    print(f"  1. Query all documents")
    print(f"  2. Query specific document")

    try:
        choice = input(f"\nSelect option (1-2, default=1): ").strip() or "1"
    except (EOFError, KeyboardInterrupt):
        print("\n\nUsing default option: 1 (Query all documents)")
        choice = "1"

    document_ids = None
    if choice == "2" and docs:
        print(f"\nSelect document:")
        for i, doc in enumerate(docs, 1):
            print(f"  {i}. {doc.filename} ({doc.record_count} records)")

        try:
            doc_choice = input(f"\nDocument (1-{len(docs)}, default=1): ").strip() or "1"
        except (EOFError, KeyboardInterrupt):
            print("\nUsing default: 1")
            doc_choice = "1"

        try:
            doc_idx = int(doc_choice) - 1
            if 0 <= doc_idx < len(docs):
                document_ids = [str(docs[doc_idx].id)]
                print(f"[OK] Filtering to: {docs[doc_idx].filename}")
        except ValueError:
            print(f"[WARN] Invalid choice, using all documents")

    # Predefined test queries
    test_queries = [
        "What is in this document?",
        "What are the main sections?",
        "List all項番 (item numbers)",
        "What parameters have 繰り返し有無 = 1?",
        "What is IFパラメータKey名?",
    ]

    print(f"\nPredefined queries:")
    for i, q in enumerate(test_queries, 1):
        print(f"  {i}. {q}")
    print(f"  {len(test_queries) + 1}. Custom query")

    try:
        query_choice = input(f"\nSelect query (1-{len(test_queries) + 1}, default=1): ").strip() or "1"
    except (EOFError, KeyboardInterrupt):
        print("\n\nUsing default query: 1")
        query_choice = "1"

    try:
        query_idx = int(query_choice) - 1
        if 0 <= query_idx < len(test_queries):
            query = test_queries[query_idx]
        else:
            try:
                query = input("\nEnter your custom query: ").strip()
            except (EOFError, KeyboardInterrupt):
                query = ""
            if not query:
                query = test_queries[0]
    except ValueError:
        query = test_queries[0]

    # Run the query
    print(f"\n{'='*70}")
    print(f"RUNNING QUERY")
    print(f"{'='*70}")

    answer, sources, confidence = await test_query(query, document_ids)

    print(f"\n{'='*70}")
    print(f"TEST COMPLETE")
    print(f"{'='*70}")

    if answer:
        print(f"\n✓ Query succeeded")
        print(f"  Confidence: {confidence:.2f}")
        print(f"  Sources: {len(sources)}")
    else:
        print(f"\n✗ Query failed")

    # Show API usage instructions
    print(f"\n{'='*70}")
    print(f"API USAGE")
    print(f"{'='*70}")
    print(f"\nTo test via API:")
    print(f"1. Start API server:")
    print(f"     uv run uvicorn src.api.main:app --reload --host 0.0.0.0 --port 8000")
    print(f"\n2. Use curl:")
    print(f'''
curl -X POST http://localhost:8000/query \\
  -H "X-API-Key: {api_key}" \\
  -H "Content-Type: application/json" \\
  -d '{{
    "query": "{query}",
    "include_context": true
  }}'
''')

    print(f"\n3. Or use Swagger UI: http://localhost:8000/docs")

    print(f"\n{'='*70}")
    print(f"MILVUS WEB UI")
    print(f"{'='*70}")
    print(f"\nView indexed vectors at: http://localhost:3000")
    print(f"  Collection: dsol_records")
    print(f"  Vectors: {len(sources) if sources else 238}")


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