#!/usr/bin/env python3
"""
PostgreSQL Analyzer Demo - Practical database analysis with MCP.

Usage:
    # Set your database connection
    export PG_CONNECTION_STRING="postgresql://user:pass@localhost:5432/mydb"
    
    # Run the demo
    python demo.py

This demo connects to a real PostgreSQL database and performs
automated analysis using the MCP server and skill guidance.
"""

import asyncio
import os
import sys
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


async def analyze_database():
    """
    Demonstrate database analysis workflow using MCP.
    """
    # Check connection string
    if not os.environ.get("PG_CONNECTION_STRING"):
        print("""
⚠️  PG_CONNECTION_STRING not set!

Please set your database connection string:
    export PG_CONNECTION_STRING="postgresql://user:password@host:port/database"

For local testing with Docker:
    docker run -d --name pg-demo -e POSTGRES_PASSWORD=demo -p 5432:5432 postgres:15
    export PG_CONNECTION_STRING="postgresql://postgres:demo@localhost:5432/postgres"
        """)
        sys.exit(1)
    
    print("=" * 70)
    print("PostgreSQL Database Analyzer Demo")
    print("=" * 70)
    
    # Configure server
    server_params = StdioServerParameters(
        command="python3",
        args=["pg_mcp_server/server.py"],
        env=os.environ.copy(),
    )
    
    print("\n📡 Connecting to PostgreSQL MCP server...")
    
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            
            await session.initialize()
            print("✅ Connected to database via MCP server\n")
            
            # Step 1: Get database schema
            print("🔍 Step 1: Discovering database schema...")
            print("-" * 70)
            schema_result = await session.call_tool("get_schema", {})
            print(schema_result.content[0].text)
            
            # Step 2: Ask user which table to analyze
            table_name = input("\n📋 Enter table name to analyze (or 'skip' to exit): ").strip()
            
            if table_name.lower() == 'skip':
                print("\nSkipping table analysis.")
                return
            
            # Step 3: Get table statistics
            print(f"\n📊 Step 2: Analyzing table '{table_name}'...")
            print("-" * 70)
            stats_result = await session.call_tool("get_table_stats", {
                "table_name": table_name,
                "sample_size": 5
            })
            print(stats_result.content[0].text)
            
            # Step 4: Custom query
            print("\n📝 Step 3: Running custom analysis queries...")
            print("-" * 70)
            
            # Example: Get column names for suggestions
            print("\nExample queries you can run:")
            print(f"  SELECT COUNT(*) FROM {table_name}")
            print(f"  SELECT * FROM {table_name} LIMIT 10")
            
            custom_query = input("\nEnter a SELECT query (or press Enter to skip): ").strip()
            
            if custom_query:
                query_result = await session.call_tool("execute_query", {
                    "query": custom_query,
                    "limit": 100
                })
                print("\nQuery Results:")
                print(query_result.content[0].text)
            
            # Step 5: Column deep dive
            column_name = input("\n🔎 Enter column name for deep analysis (or press Enter to skip): ").strip()
            
            if column_name:
                print(f"\nAnalyzing column '{column_name}'...")
                print("-" * 70)
                col_result = await session.call_tool("analyze_column", {
                    "table_name": table_name,
                    "column_name": column_name
                })
                print(col_result.content[0].text)
            
            print("\n" + "=" * 70)
            print("Analysis complete!")
            print("=" * 70)


def print_usage_examples():
    """Print usage examples for different scenarios."""
    print("""
🎯 Usage Examples with Kimi Code CLI:

Once configured, you can ask Kimi:

1. Schema Exploration:
   "Show me all tables in my database"
   "What columns does the users table have?"

2. Data Querying:
   "How many orders were placed last month?"
   "Show me the top 10 customers by revenue"

3. Data Analysis:
   "Analyze the orders table"
   "What's the distribution of user signups by month?"

4. Column Investigation:
   "Tell me about the status column in orders"
   "Are there any data quality issues in the email column?"

5. Business Insights:
   "What's our monthly revenue trend?"
   "Which products have the highest return rate?"

The skill will guide Kimi to:
- Use get_schema() to understand structure
- Use get_table_stats() for overview
- Use execute_query() for custom analysis
- Use analyze_column() for deep dives
- Synthesize insights and recommendations
""")


async def main():
    """Main entry point."""
    if len(sys.argv) > 1 and sys.argv[1] == '--examples':
        print_usage_examples()
        return
    
    try:
        await analyze_database()
    except Exception as e:
        print(f"\n❌ Error: {e}")
        print("\nMake sure:")
        print("1. PostgreSQL is running and accessible")
        print("2. PG_CONNECTION_STRING is set correctly")
        print("3. The MCP server dependencies are installed:")
        print("   pip install -r pg_mcp_server/requirements.txt")
        sys.exit(1)


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