#!/usr/bin/env python3
"""
Functional tests for MCP (Model Context Protocol) module.
"""

import asyncio
import os
import sys
import json

sys.path.insert(0, os.path.dirname(__file__))


def get_provider():
    return dict(
        model=os.getenv("OPENAI_COMPAT_MODEL", "gpt-4o-mini"),
        base_url=os.getenv("OPENAI_COMPAT_BASE_URL", "https://api.example.com/v1"),
        api_key=os.getenv("OPENAI_COMPAT_API_KEY", "test-api-key"),
    )


# ============================================================================
# MCP CLIENT TESTS
# ============================================================================


def test_mcp_client_initialization():
    """Test MCPClient initialization."""
    print("\n" + "=" * 60)
    print("TEST 1: MCP Client - Initialization")
    print("=" * 60)

    from logos_adk.mcp import MCPClient

    # Test with command
    client1 = MCPClient(command="npx @modelcontextprotocol/server-github")
    assert client1._command == "npx @modelcontextprotocol/server-github"
    assert client1._url is None
    assert client1._connected is False

    # Test with URL
    client2 = MCPClient(url="http://localhost:3000/sse")
    assert client2._command is None
    assert client2._url == "http://localhost:3000/sse"
    assert client2._connected is False

    # Test env and headers
    client3 = MCPClient(
        command="npx test",
        env={"KEY": "value"},
        headers={"Auth": "Bearer token"},
        timeout=60.0,
    )
    assert client3._env == {"KEY": "value"}
    assert client3._headers == {"Auth": "Bearer token"}
    assert client3._timeout == 60.0

    print(f"  Client 1: command={client1._command}")
    print(f"  Client 2: url={client2._url}")
    print(f"  Client 3: env={client3._env}, timeout={client3._timeout}")

    print("✓ TEST 1 PASSED\n")
    return True


def test_mcp_client_validation():
    """Test MCPClient parameter validation."""
    print("\n" + "=" * 60)
    print("TEST 2: MCP Client - Validation")
    print("=" * 60)

    from logos_adk.mcp import MCPClient

    # Missing both command and url
    try:
        MCPClient()
        print("  ERROR: Should have raised ValueError")
        return False
    except ValueError as e:
        assert "requires either" in str(e)
        print(f"  Correctly rejected missing params: {e}")

    # Both command and url
    try:
        MCPClient(command="test", url="http://localhost")
        print("  ERROR: Should have raised ValueError")
        return False
    except ValueError as e:
        assert "not both" in str(e)
        print(f"  Correctly rejected both params: {e}")

    print("✓ TEST 2 PASSED\n")
    return True


def test_mcp_client_from_string():
    """Test MCPClient.from_string factory method."""
    print("\n" + "=" * 60)
    print("TEST 3: MCP Client - From String")
    print("=" * 60)

    from logos_adk.mcp import MCPClient

    # HTTP URL
    client1 = MCPClient.from_string("http://localhost:3000/sse")
    assert client1._url == "http://localhost:3000/sse"
    assert client1._command is None
    print(f"  HTTP URL: {client1._url}")

    # HTTPS URL
    client2 = MCPClient.from_string("https://localhost:3000/sse")
    assert client2._url == "https://localhost:3000/sse"
    print(f"  HTTPS URL: {client2._url}")

    # Command (non-URL)
    client3 = MCPClient.from_string("npx @modelcontextprotocol/server-github")
    assert client3._command == "npx @modelcontextprotocol/server-github"
    assert client3._url is None
    print(f"  Command: {client3._command}")

    # With extra kwargs
    client4 = MCPClient.from_string("npx test", timeout=60.0)
    assert client4._timeout == 60.0
    print(f"  With kwargs: timeout={client4._timeout}")

    print("✓ TEST 3 PASSED\n")
    return True


def test_mcp_client_repr():
    """Test MCPClient string representation."""
    print("\n" + "=" * 60)
    print("TEST 4: MCP Client - Repr")
    print("=" * 60)

    from logos_adk.mcp import MCPClient

    client = MCPClient(command="npx test")

    # Not connected
    repr1 = repr(client)
    assert "disconnected" in repr1
    assert "npx test" in repr1
    print(f"  Not connected: {repr1}")

    # Connected (mock)
    client._connected = True
    from logos_adk.tools import Tool

    async def mock_fn1():
        return "result1"

    async def mock_fn2():
        return "result2"

    async def mock_fn3():
        return "result3"

    client._tools = [
        Tool(fn=mock_fn1, name="tool1", description="desc1"),
        Tool(fn=mock_fn2, name="tool2", description="desc2"),
        Tool(fn=mock_fn3, name="tool3", description="desc3"),
    ]
    repr2 = repr(client)
    assert "connected" in repr2
    assert "tools=3" in repr2
    print(f"  Connected: {repr2}")

    print("✓ TEST 4 PASSED\n")
    return True


def test_mcp_client_get_tools_not_connected():
    """Test get_tools raises error when not connected."""
    print("\n" + "=" * 60)
    print("TEST 5: MCP Client - Get Tools Error")
    print("=" * 60)

    from logos_adk.mcp import MCPClient

    client = MCPClient(command="npx test")

    try:
        client.get_tools()
        print("  ERROR: Should have raised RuntimeError")
        return False
    except RuntimeError as e:
        assert "not connected" in str(e).lower()
        print(f"  Correctly raised: {e}")

    print("✓ TEST 5 PASSED\n")
    return True


# ============================================================================
# MCP SERVER TESTS
# ============================================================================


def test_mcp_server_validation():
    """Test serve_as_mcp parameter validation."""
    print("\n" + "=" * 60)
    print("TEST 6: MCP Server - Validation")
    print("=" * 60)

    from logos_adk.mcp import serve_as_mcp
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider

    # Reset runtime
    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    provider = OpenAICompatibleProvider(**get_provider())
    agent = Agent("test_agent", model=provider)

    # Invalid transport
    try:
        serve_as_mcp(agent, transport="invalid")
        print("  ERROR: Should have raised ValueError")
        return False
    except ValueError as e:
        assert "Unknown transport" in str(e)
        print(f"  Correctly rejected invalid transport: {e}")

    print("✓ TEST 6 PASSED\n")
    return True


async def test_mcp_server_async_validation():
    """Test serve_as_mcp_async validation."""
    print("\n" + "=" * 60)
    print("TEST 7: MCP Server Async - Validation")
    print("=" * 60)

    from logos_adk.mcp import serve_as_mcp_async
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider

    # Reset runtime
    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    provider = OpenAICompatibleProvider(**get_provider())
    agent = Agent("test_agent", model=provider)

    # Invalid transport
    try:
        await serve_as_mcp_async(agent, transport="sse")
        print("  ERROR: Should have raised ValueError")
        return False
    except ValueError as e:
        assert "only supports" in str(e)
        print(f"  Correctly rejected invalid transport: {e}")

    print("✓ TEST 7 PASSED\n")
    return True


def test_mcp_server_create():
    """Test _create_mcp_server function."""
    print("\n" + "=" * 60)
    print("TEST 8: MCP Server - Create")
    print("=" * 60)

    from logos_adk.mcp.server import _create_mcp_server
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider

    # Reset runtime
    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    provider = OpenAICompatibleProvider(**get_provider())
    agent = Agent("test_agent", model=provider)

    # Create server
    server = _create_mcp_server(agent, name="custom-name")
    assert server is not None
    print(f"  Created server: {server}")

    # Create with default name
    server2 = _create_mcp_server(agent)
    assert server2 is not None
    print(f"  Created server (default name): {server2}")

    print("✓ TEST 8 PASSED\n")
    return True


# ============================================================================
# INTEGRATION TEST - Mock Connection
# ============================================================================


async def test_mcp_client_mock_connection():
    """Test MCPClient with mocked connection."""
    print("\n" + "=" * 60)
    print("TEST 9: Integration - Mock Connection")
    print("=" * 60)

    from logos_adk.mcp import MCPClient

    # Create a client but don't actually connect
    client = MCPClient(command="npx test-server")

    # Manually set connected state and mock tools
    client._connected = True

    # Create a mock tool
    from logos_adk.tools import Tool

    async def mock_fn():
        return "result"

    mock_tool = Tool(fn=mock_fn, name="mock_tool", description="A mock tool")
    client._tools = [mock_tool]

    # Get tools
    tools = client.get_tools()
    assert len(tools) == 1
    assert tools[0].name == "mock_tool"
    print(f"  Got tools: {[t.name for t in tools]}")

    # Disconnect
    await client.disconnect()
    assert client._connected is False
    assert len(client._tools) == 0
    print(f"  Disconnected: {client.connected}")

    print("✓ TEST 9 PASSED\n")
    return True


# ============================================================================
# MAIN
# ============================================================================


async def main():
    print("=" * 60)
    print("MCP (Model Context Protocol) TESTS")
    print("=" * 60)

    results = []

    # Sync tests
    sync_tests = [
        ("Client: Initialization", test_mcp_client_initialization),
        ("Client: Validation", test_mcp_client_validation),
        ("Client: From String", test_mcp_client_from_string),
        ("Client: Repr", test_mcp_client_repr),
        ("Client: Get Tools Error", test_mcp_client_get_tools_not_connected),
        ("Server: Validation", test_mcp_server_validation),
        ("Server: Create", test_mcp_server_create),
    ]

    for name, test_func in sync_tests:
        try:
            result = test_func()
            results.append((name, result))
        except Exception as e:
            print(f"✗ {name}: {e}")
            import traceback

            traceback.print_exc()
            results.append((name, False))

    # Async tests
    async_tests = [
        ("Server: Async Validation", test_mcp_server_async_validation),
        ("Integration: Mock Connection", test_mcp_client_mock_connection),
    ]

    for name, test_func in async_tests:
        try:
            result = await test_func()
            results.append((name, result))
        except Exception as e:
            print(f"✗ {name}: {e}")
            import traceback

            traceback.print_exc()
            results.append((name, False))

    print("=" * 60)
    print("RESULTS")
    print("=" * 60)

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

    for name, ok in results:
        status = "✓ PASS" if ok else "✗ FAIL"
        print(f"{status}: {name}")

    print(f"\n{passed}/{total} tests passed")
    return passed == total


if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)
