#!/usr/bin/env python3
"""
Test script to diagnose Langfuse connectivity issues
"""
import json
import sys

print("=" * 60)
print("Langfuse Connection Diagnostic Tool")
print("=" * 60)

# Test 1: Check if required modules are installed
print("\n1. Checking required modules...")
try:
    import boto3
    print("   ✓ boto3 installed")
except ImportError:
    print("   ✗ boto3 NOT installed")
    sys.exit(1)

try:
    from langfuse import Langfuse
    print("   ✓ langfuse installed")
except ImportError:
    print("   ✗ langfuse NOT installed")
    sys.exit(1)

try:
    from botocore.exceptions import ClientError
    print("   ✓ botocore installed")
except ImportError:
    print("   ✗ botocore NOT installed")
    sys.exit(1)

# Test 2: Fetch secrets from AWS Secrets Manager
print("\n2. Fetching secrets from AWS Secrets Manager...")
secrets = None
try:
    secret_name = "assistente-db-secrets-manager"
    region_name = "us-east-1"

    session = boto3.session.Session()
    client = session.client(
        service_name='secretsmanager',
        region_name=region_name
    )

    get_secret_value_response = client.get_secret_value(SecretId=secret_name)
    secret = get_secret_value_response['SecretString']
    secrets = json.loads(secret)

    print(f"   ✓ Successfully fetched secrets")
    print(f"   ✓ Found keys: {list(secrets.keys())}")

    # Check for Langfuse keys
    has_public_key = 'LANGFUSE-PUBLIC-KEY' in secrets
    has_secret_key = 'LANGFUSE-SECRET-KEY' in secrets

    print(f"   {'✓' if has_public_key else '✗'} LANGFUSE-PUBLIC-KEY {'found' if has_public_key else 'MISSING'}")
    print(f"   {'✓' if has_secret_key else '✗'} LANGFUSE-SECRET-KEY {'found' if has_secret_key else 'MISSING'}")

    if not has_public_key or not has_secret_key:
        print("\n   ⚠ Langfuse credentials not found in Secrets Manager!")
        secrets = None

except ClientError as e:
    print(f"   ⚠ Error fetching secrets: {e}")
    print("   Will try with manual input...")
    secrets = None
except Exception as e:
    print(f"   ⚠ Unexpected error: {e}")
    print("   Will try with manual input...")
    secrets = None

# If AWS Secrets Manager failed, try manual input
if secrets is None:
    print("\n   Please provide Langfuse credentials manually:")
    public_key = input("   Enter LANGFUSE-PUBLIC-KEY: ").strip()
    secret_key = input("   Enter LANGFUSE-SECRET-KEY: ").strip()

    if not public_key or not secret_key:
        print("   ✗ Both keys are required!")
        sys.exit(1)

    secrets = {
        'LANGFUSE-PUBLIC-KEY': public_key,
        'LANGFUSE-SECRET-KEY': secret_key
    }
    print("   ✓ Manual credentials provided")

# Test 3: Initialize Langfuse client
print("\n3. Initializing Langfuse client...")
try:
    host = "http://98.92.98.83:3000"
    langfuse = Langfuse(
        public_key=secrets['LANGFUSE-PUBLIC-KEY'],
        secret_key=secrets['LANGFUSE-SECRET-KEY'],
        host=host
    )
    print(f"   ✓ Langfuse client initialized with host: {host}")
except Exception as e:
    print(f"   ✗ Error initializing Langfuse: {e}")
    sys.exit(1)

# Test 4: Test API connectivity
print("\n4. Testing Langfuse API connectivity...")
try:
    # Create a simple test trace
    trace = langfuse.trace(
        name="connection_test",
        user_id="diagnostic_script"
    )
    print("   ✓ Test trace created")

    # Try to flush
    langfuse.flush()
    print("   ✓ Data flushed to Langfuse")

    print("\n" + "=" * 60)
    print("SUCCESS: Langfuse connection is working properly!")
    print("=" * 60)

except Exception as e:
    print(f"   ✗ Error communicating with Langfuse: {e}")
    print(f"\n   Error details: {type(e).__name__}: {str(e)}")

    # Additional diagnostic info
    print("\n   Troubleshooting tips:")
    print("   - Verify the Langfuse server is running")
    print("   - Check if the host URL is correct")
    print("   - Verify your API keys are valid in Langfuse UI")
    print("   - Check network connectivity to the Langfuse server")
    sys.exit(1)