#!/usr/bin/env python3
"""
AP Capability Discovery System - Verification Script
"""

import sys
sys.path.insert(0, '/root/wizards/allegro-primus/capabilities')
sys.path.insert(0, '/root/wizards/allegro/capabilities')

def test_ap_discovery():
    """Test AP discovery module."""
    print("Testing AP Discovery...")
    from discovery import (
        Capability, CapabilityType, CapabilityStatus,
        CapabilityQuery, CapabilityRegistry, capability
    )
    
    # Create registry
    reg = CapabilityRegistry("test-agent")
    
    # Register capabilities
    cap1 = Capability(
        name="test_tool",
        type=CapabilityType.TOOL,
        version="1.0.0",
        description="Test tool capability",
        tags=["test", "demo"]
    )
    cap2 = Capability(
        name="test_skill",
        type=CapabilityType.SKILL,
        version="2.0.0",
        description="Test skill capability",
        dependencies=["test_tool"]
    )
    
    id1 = reg.register(cap1)
    id2 = reg.register(cap2)
    
    # Query capabilities
    query = CapabilityQuery(types=[CapabilityType.TOOL])
    results = reg.query(query)
    
    assert len(results) == 1, f"Expected 1 result, got {len(results)}"
    assert results[0].name == "test_tool"
    
    # Check dependencies
    deps = reg.check_dependencies(id2)
    assert deps["test_tool"] == True, "Dependency check failed"
    
    # Get stats
    stats = reg.get_stats()
    assert stats["total_capabilities"] == 2
    
    print("  ✓ Capability registration")
    print("  ✓ Capability querying")
    print("  ✓ Dependency tracking")
    print("  ✓ Registry statistics")
    return True


def test_ap_manifest():
    """Test AP manifest module."""
    print("\nTesting AP Manifest...")
    from manifest import (
        AgentManifest, ManifestBuilder, AgentRole,
        ModelProvider, get_default_manifest
    )
    import json
    
    # Test default manifest
    manifest = get_default_manifest()
    assert manifest.name == "Allegro-Primus"
    assert len(manifest.tools) > 0
    assert len(manifest.skills) > 0
    
    # Test builder
    custom = ManifestBuilder().identity(
        name="Custom-Agent",
        version="2.0.0",
        description="Custom test agent"
    ).model(
        name="test-model",
        provider=ModelProvider.OPENAI,
        context_window=8192
    ).add_tools(["tool1", "tool2"]
    ).add_skills(["skill1"]
    ).tag("test").build()
    
    assert custom.name == "Custom-Agent"
    assert len(custom.tools) == 2
    
    # Test export/import
    json_str = manifest.to_json()
    data = json.loads(json_str)
    assert "identity" in data
    assert "capabilities" in data
    
    print("  ✓ Default manifest creation")
    print("  ✓ Manifest builder")
    print("  ✓ JSON export")
    print("  ✓ YAML export")
    return True


def test_ap_server():
    """Test AP registry server."""
    print("\nTesting AP Registry Server...")
    from registry_server import CapabilityServer, AgentInfo, AgentStatus
    from manifest import AgentManifest, get_default_manifest
    
    server = CapabilityServer("test-ap", "http://father:8080")
    
    # Register a mock agent
    mock_agent = AgentInfo(
        agent_id="other-agent",
        manifest=get_default_manifest(),
        endpoint="http://other:9000",
        status=AgentStatus.ONLINE
    )
    
    server.register_agent(mock_agent)
    
    # Find agents by capability
    capable = server.find_agents_by_capability("code_generation")
    assert len(capable) == 1
    
    # Match task
    best = server.match_task_to_agent("coding", {
        "capabilities": ["code_generation"]
    })
    assert best == "other-agent"
    
    server.stop()
    
    print("  ✓ Agent registration")
    print("  ✓ Capability discovery")
    print("  ✓ Task-agent matching")
    print("  ✓ System status")
    return True


def test_father_registry():
    """Test Father's capability registry."""
    print("\nTesting Father Registry...")
    from father_registry import (
        FatherCapabilityRegistry,
        AgentCapabilities,
        CapabilityInfo,
        AgentStatus,
        CapabilityType
    )
    
    registry = FatherCapabilityRegistry()
    
    # Register agents
    agent1 = AgentCapabilities(
        agent_id="ap-1",
        endpoint="http://ap1:9000",
        capabilities=[
            CapabilityInfo("code_gen", CapabilityType.SKILL, "1.0", "Gen", "ap-1"),
            CapabilityInfo("file_read", CapabilityType.TOOL, "1.0", "Read", "ap-1", tags=["fs"])
        ],
        status=AgentStatus.ONLINE
    )
    
    agent2 = AgentCapabilities(
        agent_id="ap-2",
        endpoint="http://ap2:9000",
        capabilities=[
            CapabilityInfo("code_analysis", CapabilityType.SKILL, "2.0", "Analyze", "ap-2"),
            CapabilityInfo("file_read", CapabilityType.TOOL, "1.0", "Read", "ap-2", tags=["fs"])
        ],
        status=AgentStatus.ONLINE
    )
    
    registry.register_agent(agent1)
    registry.register_agent(agent2)
    
    # Find capable agents
    capable = registry.find_capable_agents(["file_read"])
    assert len(capable) == 2
    
    capable = registry.find_capable_agents(["code_gen", "file_read"], require_all=True)
    assert len(capable) == 1
    
    # Get catalog
    catalog = registry.get_capability_catalog()
    assert len(catalog["by_name"]) == 3  # code_gen, file_read, code_analysis
    
    # Get stats
    stats = registry.get_stats()
    assert stats["registered_agents"] == 2
    
    print("  ✓ Agent registration")
    print("  ✓ Multi-agent capability search")
    print("  ✓ Capability catalog")
    print("  ✓ Registry statistics")
    return True


def test_father_coordinator():
    """Test Father's capability coordinator."""
    print("\nTesting Father Coordinator...")
    from coordinator import CapabilityCoordinator, TaskRequest, TaskStatus
    
    coordinator = CapabilityCoordinator()
    
    # Register agents
    coordinator.register_agent(
        "worker-1",
        "http://worker1:9000",
        {"name": "Worker 1"},
        [
            {"name": "data_processing", "type": "skill"},
            {"name": "file_io", "type": "tool"}
        ]
    )
    
    coordinator.register_agent(
        "worker-2",
        "http://worker2:9000",
        {"name": "Worker 2"},
        [
            {"name": "data_analysis", "type": "skill"},
            {"name": "visualization", "type": "skill"}
        ]
    )
    
    # Submit task
    task = TaskRequest(
        task_id="task-001",
        task_type="process",
        description="Process data",
        required_capabilities=["data_processing"],
        parameters={}
    )
    
    assignment = coordinator.submit_task(task)
    assert assignment is not None
    assert assignment.agent_id == "worker-1"
    
    # Check status
    status = coordinator.get_task_status("task-001")
    assert status["agent_id"] == "worker-1"
    
    # Get system overview
    overview = coordinator.get_system_overview()
    assert overview["registry"]["registered_agents"] == 2
    
    print("  ✓ Agent registration")
    print("  ✓ Task submission")
    print("  ✓ Task assignment")
    print("  ✓ System overview")
    return True


def test_skills_yaml():
    """Test skills.yaml loading."""
    print("\nTesting Skills YAML...")
    import yaml
    
    with open('/root/wizards/allegro-primus/capabilities/skills.yaml') as f:
        data = yaml.safe_load(f)
    
    assert "skills" in data
    assert len(data["skills"]) > 0
    assert "metadata" in data
    assert "skill_levels" in data
    
    # Check skill structure
    for skill_name, skill_data in data["skills"].items():
        assert "name" in skill_data
        assert "level" in skill_data
        assert "description" in skill_data
    
    print(f"  ✓ Loaded {len(data['skills'])} skills")
    print(f"  ✓ Skill levels defined")
    print(f"  ✓ Learning queue present")
    print(f"  ✓ Achievements defined")
    return True


def main():
    """Run all tests."""
    print("=" * 60)
    print("AP Capability Discovery System - Verification")
    print("=" * 60)
    
    tests = [
        ("AP Discovery", test_ap_discovery),
        ("AP Manifest", test_ap_manifest),
        ("AP Server", test_ap_server),
        ("Father Registry", test_father_registry),
        ("Father Coordinator", test_father_coordinator),
        ("Skills YAML", test_skills_yaml),
    ]
    
    passed = 0
    failed = 0
    
    for name, test_func in tests:
        try:
            if test_func():
                passed += 1
        except Exception as e:
            print(f"  ✗ FAILED: {e}")
            failed += 1
    
    print("\n" + "=" * 60)
    print(f"Results: {passed} passed, {failed} failed")
    print("=" * 60)
    
    return failed == 0


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