#!/usr/bin/python3
"""
Test script for resource allocation in Debian Forge build orchestrator

This script tests the resource management and allocation functionality.
"""

import sys
import time
from build_orchestrator import BuildOrchestrator, ResourceManager


def test_resource_manager():
    """Test the ResourceManager class"""
    print("Testing ResourceManager...")
    
    rm = ResourceManager()
    
    # Test resource availability
    available = rm.get_available_resources()
    print(f"Available resources: CPU {available['cpu_percent']:.1f}%, "
          f"Memory {available['memory_gb']:.1f}GB, "
          f"Storage {available['storage_gb']:.1f}GB")
    
    # Test resource allocation
    test_reqs = {
        "cpu_percent": 50,
        "memory_gb": 2,
        "storage_gb": 5
    }
    
    can_allocate = rm.can_allocate_resources(test_reqs)
    print(f"Can allocate resources for {test_reqs}: {can_allocate}")
    
    return True


def test_build_orchestrator():
    """Test the BuildOrchestrator with resource management"""
    print("Testing BuildOrchestrator with resource management...")
    
    orchestrator = BuildOrchestrator()
    
    # Submit builds with different resource requirements
    build1 = orchestrator.submit_build(
        "test-manifest.json",
        priority=5,
        resource_requirements={"cpu_percent": 30, "memory_gb": 1, "storage_gb": 2}
    )
    
    build2 = orchestrator.submit_build(
        "test-manifest.json",
        priority=3,
        resource_requirements={"cpu_percent": 60, "memory_gb": 3, "storage_gb": 5}
    )
    
    print(f"Submitted builds: {build1}, {build2}")
    
    # Check resource status
    resource_status = orchestrator.get_resource_status()
    print(f"Resource status: {resource_status}")
    
    # List builds
    builds = orchestrator.list_builds()
    print(f"Pending builds: {len(builds['pending'])}")
    print(f"Running builds: {len(builds['running'])}")
    print(f"Completed builds: {len(builds['completed'])}")
    
    return True


def test_concurrent_builds():
    """Test concurrent build handling with resource constraints"""
    print("Testing concurrent build handling...")
    
    orchestrator = BuildOrchestrator()
    
    # Submit multiple builds with resource requirements
    builds = []
    for i in range(5):
        build_id = orchestrator.submit_build(
            f"test-manifest-{i}.json",
            priority=5-i,  # Higher priority for lower i
            resource_requirements={
                "cpu_percent": 20 + (i * 10),
                "memory_gb": 1 + i,
                "storage_gb": 2 + i
            }
        )
        builds.append(build_id)
        print(f"Submitted build {build_id}")
    
    # Start orchestrator
    orchestrator.start()
    
    # Monitor for a short time
    try:
        for _ in range(10):
            resource_status = orchestrator.get_resource_status()
            print(f"Resources: CPU {resource_status['available_resources']['cpu_percent']:.1f}% free, "
                  f"Memory {resource_status['available_resources']['memory_gb']:.1f}GB free, "
                  f"Queue: {resource_status['queue_length']} pending")
            time.sleep(1)
    except KeyboardInterrupt:
        pass
    
    # Stop orchestrator
    orchestrator.stop()
    
    return True


def main():
    """Run all tests"""
    print("Resource Allocation Tests")
    print("=" * 40)
    
    tests = [
        test_resource_manager,
        test_build_orchestrator,
        test_concurrent_builds
    ]
    
    passed = 0
    total = len(tests)
    
    for test in tests:
        try:
            print(f"\nRunning {test.__name__}...")
            if test():
                print(f"✅ {test.__name__} passed")
                passed += 1
            else:
                print(f"❌ {test.__name__} failed")
        except Exception as e:
            print(f"❌ {test.__name__} failed with exception: {e}")
        
        print()
    
    print("=" * 40)
    print(f"Test Results: {passed}/{total} passed")
    
    if passed == total:
        print("🎉 All resource allocation tests passed!")
        return 0
    else:
        print("⚠️ Some tests failed")
        return 1


if __name__ == "__main__":
    sys.exit(main())