#!/usr/bin/env python3
"""
Resource Limits Enforcement Test Script

Tests that container resource limits are properly applied and enforced,
preventing resource exhaustion attacks and ensuring fair resource allocation.
"""

import os
import sys
import asyncio
import json
import time
from pathlib import Path

# Add session-manager to path for imports
sys.path.insert(0, str(Path(__file__).parent))

from resource_manager import (
    get_resource_limits,
    check_system_resources,
    should_throttle_sessions,
    ResourceLimits,
    ResourceValidator,
    ResourceMonitor,
)

# Set up logging
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


async def test_resource_limits_configuration():
    """Test resource limits configuration and validation."""
    print("🧪 Testing Resource Limits Configuration")
    print("=" * 50)

    # Test default configuration
    print("1️⃣ Testing default configuration...")
    try:
        limits = get_resource_limits()
        print(
            f"✅ Default limits loaded: memory={limits.memory_limit}, cpu_quota={limits.cpu_quota}"
        )

        # Validate limits
        valid, message = limits.validate()
        if valid:
            print("✅ Default limits validation passed")
        else:
            print(f"❌ Default limits validation failed: {message}")
            return False
    except Exception as e:
        print(f"❌ Failed to load default configuration: {e}")
        return False

    # Test custom configuration
    print("\n2️⃣ Testing custom configuration...")
    test_configs = [
        {"memory_limit": "2g", "cpu_quota": 50000, "cpu_period": 100000},
        {"memory_limit": "512m", "cpu_quota": 25000, "cpu_period": 100000},
        {"memory_limit": "8g", "cpu_quota": 200000, "cpu_period": 100000},
    ]

    for config in test_configs:
        try:
            valid, message, limits = ResourceValidator.validate_resource_config(config)
            if valid:
                print(f"✅ Config {config} validated successfully")
            else:
                print(f"❌ Config {config} validation failed: {message}")
                return False
        except Exception as e:
            print(f"❌ Config {config} validation error: {e}")
            return False

    # Test invalid configurations
    print("\n3️⃣ Testing invalid configurations...")
    invalid_configs = [
        {
            "memory_limit": "0g",
            "cpu_quota": 100000,
            "cpu_period": 100000,
        },  # Zero memory
        {
            "memory_limit": "1g",
            "cpu_quota": 200000,
            "cpu_period": 100000,
        },  # CPU quota > period
        {
            "memory_limit": "1g",
            "cpu_quota": -1000,
            "cpu_period": 100000,
        },  # Negative CPU
    ]

    for config in invalid_configs:
        valid, message, limits = ResourceValidator.validate_resource_config(config)
        if not valid:
            print(f"✅ Correctly rejected invalid config {config}: {message}")
        else:
            print(f"❌ Incorrectly accepted invalid config {config}")
            return False

    return True


async def test_resource_monitoring():
    """Test system resource monitoring."""
    print("\n🖥️  Testing System Resource Monitoring")
    print("=" * 50)

    # Test resource monitoring
    print("1️⃣ Testing system resource monitoring...")
    try:
        resource_check = check_system_resources()
        print("✅ System resource check completed")

        if isinstance(resource_check, dict):
            system_resources = resource_check.get("system_resources", {})
            alerts = resource_check.get("alerts", [])

            print(f"📊 System resources: {len(system_resources)} metrics collected")
            print(f"🚨 Resource alerts: {len(alerts)} detected")

            for alert in alerts:
                print(
                    f"   {alert.get('level', 'unknown').upper()}: {alert.get('message', 'No message')}"
                )

        # Test throttling logic
        should_throttle, reason = should_throttle_sessions()
        print(f"🎛️  Session throttling: {'YES' if should_throttle else 'NO'} - {reason}")

    except Exception as e:
        print(f"❌ Resource monitoring test failed: {e}")
        return False

    return True


async def test_memory_limit_parsing():
    """Test memory limit parsing functionality."""
    print("\n💾 Testing Memory Limit Parsing")
    print("=" * 50)

    test_cases = [
        ("4g", (4 * 1024**3, "GB")),
        ("512m", (512 * 1024**2, "MB")),
        ("256k", (256 * 1024, "KB")),
        ("1073741824", (1073741824, "bytes")),  # 1GB in bytes
    ]

    for memory_str, expected in test_cases:
        try:
            bytes_val, unit = ResourceValidator.parse_memory_limit(memory_str)
            if bytes_val == expected[0] and unit == expected[1]:
                print(f"✅ Parsed {memory_str} -> {bytes_val} {unit}")
            else:
                print(
                    f"❌ Failed to parse {memory_str}: got {bytes_val} {unit}, expected {expected}"
                )
                return False
        except Exception as e:
            print(f"❌ Error parsing {memory_str}: {e}")
            return False

    return True


async def test_docker_limits_conversion():
    """Test Docker limits conversion."""
    print("\n🐳 Testing Docker Limits Conversion")
    print("=" * 50)

    limits = ResourceLimits(
        memory_limit="2g",
        cpu_quota=75000,
        cpu_period=100000,
    )

    docker_limits = limits.to_docker_limits()
    expected = {
        "mem_limit": "2g",
        "cpu_quota": 75000,
        "cpu_period": 100000,
    }

    if docker_limits == expected:
        print("✅ Docker limits conversion correct")
        return True
    else:
        print(f"❌ Docker limits conversion failed: {docker_limits} != {expected}")
        return False


async def test_environment_variables():
    """Test environment variable configuration."""
    print("\n🌍 Testing Environment Variable Configuration")
    print("=" * 50)

    # Save original environment
    original_env = {}
    test_vars = [
        "CONTAINER_MEMORY_LIMIT",
        "CONTAINER_CPU_QUOTA",
        "CONTAINER_CPU_PERIOD",
        "MAX_CONCURRENT_SESSIONS",
        "MEMORY_WARNING_THRESHOLD",
        "CPU_WARNING_THRESHOLD",
    ]

    for var in test_vars:
        original_env[var] = os.environ.get(var)

    try:
        # Test custom environment configuration
        os.environ["CONTAINER_MEMORY_LIMIT"] = "3g"
        os.environ["CONTAINER_CPU_QUOTA"] = "80000"
        os.environ["CONTAINER_CPU_PERIOD"] = "100000"
        os.environ["MAX_CONCURRENT_SESSIONS"] = "5"

        # Force reload of configuration
        import importlib
        import resource_manager

        importlib.reload(resource_manager)

        limits = resource_manager.get_resource_limits()

        if limits.memory_limit == "3g" and limits.cpu_quota == 80000:
            print("✅ Environment variable configuration working")
            return True
        else:
            print(f"❌ Environment variable configuration failed: got {limits}")
            return False

    finally:
        # Restore original environment
        for var, value in original_env.items():
            if value is not None:
                os.environ[var] = value
            elif var in os.environ:
                del os.environ[var]


async def run_all_tests():
    """Run all resource limit tests."""
    print("🚀 Resource Limits Enforcement Test Suite")
    print("=" * 60)

    tests = [
        ("Resource Limits Configuration", test_resource_limits_configuration),
        ("System Resource Monitoring", test_resource_monitoring),
        ("Memory Limit Parsing", test_memory_limit_parsing),
        ("Docker Limits Conversion", test_docker_limits_conversion),
        ("Environment Variables", test_environment_variables),
    ]

    results = []
    for test_name, test_func in tests:
        print(f"\n{'=' * 20} {test_name} {'=' * 20}")
        try:
            result = await test_func()
            results.append(result)
            status = "✅ PASSED" if result else "❌ FAILED"
            print(f"\n{status}: {test_name}")
        except Exception as e:
            print(f"\n❌ ERROR in {test_name}: {e}")
            results.append(False)

    # Summary
    print(f"\n{'=' * 60}")
    passed = sum(results)
    total = len(results)
    print(f"📊 Test Results: {passed}/{total} tests passed")

    if passed == total:
        print("🎉 All resource limit tests completed successfully!")
        print("💪 Container resource limits are properly configured and enforced.")
    else:
        print("⚠️  Some tests failed. Check the output above for details.")
        sys.exit(1)


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