#!/usr/bin/env python3
"""
Test that execution environment resolution works for both Agno and Claude Code runtimes.

Tests that:
1. Resolved env vars from execution environment are passed to both runtimes
2. MCP servers receive env vars correctly
3. Both runtimes can access secrets and integrations via env vars
"""
import sys
import os
from unittest.mock import Mock, patch

# Add parent directory to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from control_plane_api.worker.runtimes.base import RuntimeExecutionContext, RuntimeType


def print_header(title: str):
    """Print a formatted test section header."""
    print("\n" + "="*80)
    print(f"  {title}")
    print("="*80 + "\n")


def print_test_result(test_name: str, passed: bool, details: str = ""):
    """Print test result with formatting."""
    status = "✅ PASS" if passed else "❌ FAIL"
    print(f"{status} {test_name}")
    if details:
        print(f"     {details}")


def test_runtime_context_env_vars():
    """Test that RuntimeExecutionContext properly stores env vars from runtime_config."""
    print_header("TEST 1: Runtime Context Env Vars Storage")

    test_cases = [
        {
            "name": "Env vars from runtime_config are accessible",
            "runtime_config": {
                "env": {
                    "GH_TOKEN": "ghp_test123",
                    "SLACK_TOKEN": "xoxb-test456",
                    "CUSTOM_VAR": "custom_value"
                }
            },
            "expected_env_keys": ["GH_TOKEN", "SLACK_TOKEN", "CUSTOM_VAR"],
        },
        {
            "name": "Empty runtime_config has no env vars",
            "runtime_config": {},
            "expected_env_keys": [],
        },
        {
            "name": "Runtime config with empty env dict",
            "runtime_config": {"env": {}},
            "expected_env_keys": [],
        },
    ]

    passed = 0
    failed = 0

    for test_case in test_cases:
        test_name = test_case["name"]

        try:
            # Create execution context
            context = Mock(spec=RuntimeExecutionContext)
            context.runtime_type = RuntimeType("default")
            context.model_id = "gpt-4"
            context.runtime_config = test_case["runtime_config"]

            # Get env vars
            env_vars = test_case["runtime_config"].get("env", {})
            env_keys = list(env_vars.keys())

            # Verify
            if env_keys == test_case["expected_env_keys"]:
                print_test_result(
                    test_name,
                    True,
                    f"Env vars: {env_keys}"
                )
                passed += 1
            else:
                print_test_result(
                    test_name,
                    False,
                    f"Expected {test_case['expected_env_keys']}, got {env_keys}"
                )
                failed += 1

        except Exception as e:
            print_test_result(test_name, False, f"Exception: {str(e)}")
            failed += 1

    print(f"\n📊 Test 1 Results: {passed} passed, {failed} failed")
    return passed, failed


def test_mcp_servers_env_injection():
    """Test that MCP servers receive env vars correctly."""
    print_header("TEST 2: MCP Servers Env Vars Injection")

    # Import the env injection function
    from control_plane_api.worker.activities.runtime_activities import inject_env_vars_into_mcp_servers

    test_cases = [
        {
            "name": "MCP server with env vars from agent config",
            "mcp_servers": {
                "github": {
                    "command": "npx",
                    "args": ["-y", "@modelcontextprotocol/server-github"],
                    "env": {
                        "GITHUB_PERSONAL_ACCESS_TOKEN": "{{GH_TOKEN}}"
                    }
                }
            },
            "agent_config": {
                "env_vars": {
                    "GH_TOKEN": "ghp_resolved_token_123"
                }
            },
            "runtime_config": {},
            "expected_token": "ghp_resolved_token_123",
        },
        {
            "name": "MCP server with KUBIYA_API_KEY injection",
            "mcp_servers": {
                "kubiya": {
                    "command": "kubiya-mcp",
                    "env": {
                        "KUBIYA_API_KEY": "{{KUBIYA_API_KEY}}"
                    }
                }
            },
            "agent_config": {},
            "runtime_config": {},
            "expected_has_kubiya_key": True,
        },
    ]

    passed = 0
    failed = 0

    for test_case in test_cases:
        test_name = test_case["name"]

        try:
            result = inject_env_vars_into_mcp_servers(
                mcp_servers=test_case["mcp_servers"],
                agent_config=test_case["agent_config"],
                runtime_config=test_case["runtime_config"],
            )

            # Check if expected values are present
            if "expected_token" in test_case:
                server_name = list(test_case["mcp_servers"].keys())[0]
                actual_token = result[server_name]["env"]["GITHUB_PERSONAL_ACCESS_TOKEN"]

                if actual_token == test_case["expected_token"]:
                    print_test_result(
                        test_name,
                        True,
                        f"Token resolved correctly: {actual_token[:20]}..."
                    )
                    passed += 1
                else:
                    print_test_result(
                        test_name,
                        False,
                        f"Expected {test_case['expected_token']}, got {actual_token}"
                    )
                    failed += 1

            elif "expected_has_kubiya_key" in test_case:
                server_name = list(test_case["mcp_servers"].keys())[0]
                has_key = "KUBIYA_API_KEY" in result[server_name]["env"]

                if has_key == test_case["expected_has_kubiya_key"]:
                    print_test_result(
                        test_name,
                        True,
                        f"KUBIYA_API_KEY present in env"
                    )
                    passed += 1
                else:
                    print_test_result(
                        test_name,
                        False,
                        f"KUBIYA_API_KEY not found in env"
                    )
                    failed += 1

        except Exception as e:
            print_test_result(test_name, False, f"Exception: {str(e)}")
            failed += 1

    print(f"\n📊 Test 2 Results: {passed} passed, {failed} failed")
    return passed, failed


def test_both_runtimes_receive_env_vars():
    """Test that both Agno and Claude Code runtimes receive env vars."""
    print_header("TEST 3: Both Runtimes Receive Env Vars")

    print("Testing that runtime_config['env'] is properly passed to both runtimes:\n")

    test_cases = [
        {
            "name": "Agno runtime receives env vars",
            "runtime_type": "default",
        },
        {
            "name": "Claude Code runtime receives env vars",
            "runtime_type": "claude_code",
        },
    ]

    passed = 0
    failed = 0

    for test_case in test_cases:
        test_name = test_case["name"]

        try:
            # Create execution context with env vars
            context = Mock(spec=RuntimeExecutionContext)
            context.runtime_type = RuntimeType(test_case["runtime_type"])
            context.model_id = "claude-sonnet-4"
            context.runtime_config = {
                "env": {
                    "GH_TOKEN": "ghp_test",
                    "SLACK_TOKEN": "xoxb_test",
                    "CUSTOM_SECRET": "secret_value"
                }
            }

            # Verify env vars are present
            env_vars = context.runtime_config.get("env", {})

            if len(env_vars) == 3:
                print_test_result(
                    test_name,
                    True,
                    f"Runtime: {test_case['runtime_type']}, Env vars: {list(env_vars.keys())}"
                )
                passed += 1
            else:
                print_test_result(
                    test_name,
                    False,
                    f"Expected 3 env vars, got {len(env_vars)}"
                )
                failed += 1

        except Exception as e:
            print_test_result(test_name, False, f"Exception: {str(e)}")
            failed += 1

    print(f"\n📊 Test 3 Results: {passed} passed, {failed} failed")
    return passed, failed


def test_execution_environment_flow():
    """Test the complete execution environment resolution flow."""
    print_header("TEST 4: Complete Execution Environment Flow")

    print("Simulating complete flow: API resolution → runtime_config → runtimes\n")

    # Simulate resolved env vars from API
    resolved_env_vars = {
        "GH_TOKEN": "ghp_api_resolved",
        "SLACK_TOKEN": "xoxb_api_resolved",
        "AWS_ACCESS_KEY_ID": "AKIA_api_resolved",
    }

    # Simulate agent config env vars (override)
    agent_env_vars = {
        "CUSTOM_VAR": "custom_value",
    }

    # Merge (as done in runtime_activities.py)
    merged_env = {**resolved_env_vars, **agent_env_vars}

    # Build runtime_config (as done in runtime_activities.py)
    runtime_config = {
        "env": merged_env
    }

    passed = 0
    failed = 0

    # Test 1: Check merge
    if len(merged_env) == 4:
        print_test_result(
            "Env vars merged correctly",
            True,
            f"4 env vars: {list(merged_env.keys())}"
        )
        passed += 1
    else:
        print_test_result(
            "Env vars merged correctly",
            False,
            f"Expected 4, got {len(merged_env)}"
        )
        failed += 1

    # Test 2: Check runtime_config
    if runtime_config.get("env") == merged_env:
        print_test_result(
            "Runtime config contains merged env",
            True,
            "runtime_config['env'] matches merged_env"
        )
        passed += 1
    else:
        print_test_result(
            "Runtime config contains merged env",
            False,
            "Mismatch between runtime_config and merged_env"
        )
        failed += 1

    # Test 3: Simulate passing to runtime context
    context = Mock(spec=RuntimeExecutionContext)
    context.runtime_config = runtime_config

    if context.runtime_config.get("env") == merged_env:
        print_test_result(
            "Execution context receives env vars",
            True,
            "Context has all 4 env vars"
        )
        passed += 1
    else:
        print_test_result(
            "Execution context receives env vars",
            False,
            "Context missing env vars"
        )
        failed += 1

    print(f"\n📊 Test 4 Results: {passed} passed, {failed} failed")
    return passed, failed


def run_all_tests():
    """Run all execution environment integration tests."""
    print("\n" + "🔐"*40)
    print("  EXECUTION ENVIRONMENT INTEGRATION TESTS")
    print("🔐"*40)

    total_passed = 0
    total_failed = 0

    # Test 1: Runtime context storage
    passed, failed = test_runtime_context_env_vars()
    total_passed += passed
    total_failed += failed

    # Test 2: MCP servers env injection
    passed, failed = test_mcp_servers_env_injection()
    total_passed += passed
    total_failed += failed

    # Test 3: Both runtimes receive env vars
    passed, failed = test_both_runtimes_receive_env_vars()
    total_passed += passed
    total_failed += failed

    # Test 4: Complete flow
    passed, failed = test_execution_environment_flow()
    total_passed += passed
    total_failed += failed

    # Final summary
    print("\n" + "="*80)
    print("  FINAL RESULTS")
    print("="*80)
    print(f"\n✅ Total Passed: {total_passed}")
    print(f"❌ Total Failed: {total_failed}")
    print(f"📊 Success Rate: {(total_passed / (total_passed + total_failed) * 100) if (total_passed + total_failed) > 0 else 0:.1f}%\n")

    if total_failed == 0:
        print("🎉 ALL TESTS PASSED! 🎉\n")
        print("✅ Execution environment resolution works for both runtimes")
        print("✅ Env vars/secrets are properly passed through the system")
        print("✅ MCP servers receive resolved environment variables\n")
        return 0
    else:
        print("⚠️  SOME TESTS FAILED ⚠️\n")
        return 1


if __name__ == "__main__":
    exit_code = run_all_tests()
    sys.exit(exit_code)
