#!/usr/bin/env python3
"""
Test script for MCP integration with Coherence Gateway.

This script verifies that the MCP adapter can properly communicate
with the gateway and execute tools.
"""

import asyncio
import logging
import os
import sys
from pathlib import Path

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent.parent))

from mcp_adapter.gateway_client import GatewayClient

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Configuration
GATEWAY_URL = os.getenv("GATEWAY_URL", "http://localhost:8000")
JWT_SECRET = os.getenv("JWT_SECRET", "55c121aa48dc1b79e3c04c805aabd2fe7646253ab6edd1c61894f8146ce1ed17")


async def test_health_check():
    """Test gateway health check."""
    logger.info("Testing gateway health check...")

    async with GatewayClient(GATEWAY_URL, JWT_SECRET) as client:
        is_healthy = await client.health_check()

        if is_healthy:
            logger.info("✅ Gateway is healthy")
            return True
        else:
            logger.error("❌ Gateway health check failed")
            return False


async def test_list_pets():
    """Test listing pets."""
    logger.info("\nTesting list pets...")

    async with GatewayClient(GATEWAY_URL, JWT_SECRET) as client:
        result = await client.execute_tool(
            "findPetsByStatus",
            {"status": ["available"]}
        )

        if "error" in result:
            logger.error(f"❌ List pets failed: {result['error']}")
            return False

        # Extract data from response
        data = result
        if "output" in data:
            data = data["output"]
        if isinstance(data, dict) and "data" in data:
            data = data["data"]

        logger.info(f"✅ Found {len(data) if isinstance(data, list) else 1} pets")

        # Show first few pets
        if isinstance(data, list) and data:
            for pet in data[:3]:
                logger.info(f"  - {pet.get('name', 'Unknown')} (ID: {pet.get('id', 'N/A')})")
            if len(data) > 3:
                logger.info(f"  ... and {len(data) - 3} more")

        return True


async def test_get_pet_by_id():
    """Test getting a specific pet."""
    logger.info("\nTesting get pet by ID...")

    async with GatewayClient(GATEWAY_URL, JWT_SECRET) as client:
        result = await client.execute_tool(
            "getPetById",
            {"petId": 1}
        )

        if "error" in result:
            logger.error(f"❌ Get pet failed: {result['error']}")
            return False

        # Extract data from response
        data = result
        if "output" in data:
            data = data["output"]
        if isinstance(data, dict) and "data" in data:
            data = data["data"]

        if isinstance(data, dict) and "name" in data:
            logger.info(f"✅ Found pet: {data['name']} (ID: {data.get('id', 'N/A')})")
            return True
        else:
            logger.error("❌ Unexpected response format")
            return False


async def test_get_inventory():
    """Test getting store inventory."""
    logger.info("\nTesting get inventory...")

    async with GatewayClient(GATEWAY_URL, JWT_SECRET) as client:
        result = await client.execute_tool(
            "getInventory",
            {}
        )

        if "error" in result:
            logger.error(f"❌ Get inventory failed: {result['error']}")
            return False

        # Extract data from response
        data = result
        if "output" in data:
            data = data["output"]
        if isinstance(data, dict) and "data" in data:
            data = data["data"]

        if isinstance(data, dict):
            logger.info("✅ Store inventory:")
            for status, count in data.items():
                logger.info(f"  - {status}: {count}")
            return True
        else:
            logger.error("❌ Unexpected response format")
            return False


async def test_unauthorized_tool():
    """Test executing a tool with wrong JWT capability."""
    logger.info("\nTesting unauthorized tool access...")

    async with GatewayClient(GATEWAY_URL, JWT_SECRET) as client:
        # Manually create a JWT for a different tool
        import time

        from jose import jwt

        wrong_token = jwt.encode(
            {
                "sub": "test-client",
                "tool": "wrongTool",  # Different from actual tool
                "exp": int(time.time()) + 300,
                "iat": int(time.time()),
            },
            JWT_SECRET,
            algorithm="HS256"
        )

        # Try to execute with wrong token
        client.client.headers["Authorization"] = f"Bearer {wrong_token}"

        result = await client.execute_tool(
            "findPetsByStatus",  # Different from JWT capability
            {"status": ["available"]}
        )

        if "error" in result:
            logger.info("✅ Unauthorized access properly rejected")
            return True
        else:
            logger.error("❌ Unauthorized access was not rejected!")
            return False


async def main():
    """Run all tests."""
    logger.info("Starting MCP Integration Tests")
    logger.info(f"Gateway URL: {GATEWAY_URL}")
    logger.info("=" * 50)

    tests = [
        ("Health Check", test_health_check),
        ("List Pets", test_list_pets),
        ("Get Pet by ID", test_get_pet_by_id),
        ("Get Inventory", test_get_inventory),
        ("Unauthorized Access", test_unauthorized_tool),
    ]

    results = []
    for test_name, test_func in tests:
        try:
            success = await test_func()
            results.append((test_name, success))
        except Exception as e:
            logger.error(f"❌ {test_name} failed with exception: {e}")
            results.append((test_name, False))

    # Summary
    logger.info("\n" + "=" * 50)
    logger.info("TEST SUMMARY")
    logger.info("=" * 50)

    passed = sum(1 for _, success in results if success)
    total = len(results)

    for test_name, success in results:
        status = "✅ PASS" if success else "❌ FAIL"
        logger.info(f"{test_name}: {status}")

    logger.info(f"\nTotal: {passed}/{total} tests passed")

    return passed == total


if __name__ == "__main__":
    try:
        success = asyncio.run(main())
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        logger.info("\nTests interrupted by user")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Test suite failed: {e}", exc_info=True)
        sys.exit(1)