"""
Quick validation script for orchestrator_two_teams.py

This script validates that:
1. All imports work correctly
2. All agents can be registered
3. Workspace structure can be created
4. No obvious errors in the setup

Run this to verify the example is ready before full execution.
"""
import asyncio
import os
import sys
from pathlib import Path

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))

try:
    import synqed
    from synqed.workspace_manager import WorkspaceManager, AgentRuntimeRegistry
    from synqed.planner import PlannerLLM, TaskTreeNode
    print("✅ All imports successful")
except ImportError as e:
    print(f"❌ Import error: {e}")
    sys.exit(1)


async def validate():
    """Validate the orchestrator setup without executing agents."""
    print("\n" + "="*70)
    print("  ORCHESTRATOR VALIDATION")
    print("="*70 + "\n")
    
    # ========================================================================
    # Step 1: Validate agent functions exist
    # ========================================================================
    print("📋 Step 1: Validating agent logic functions...")
    
    # Import agent functions from the example
    try:
        from orchestrator_two_teams import (
            project_manager_logic,
            research_lead_logic,
            data_analyst_logic,
            report_writer_logic,
            tech_lead_logic,
            backend_dev_logic,
            frontend_dev_logic
        )
        print("   ✅ All 7 agent logic functions imported successfully")
    except ImportError as e:
        print(f"   ❌ Failed to import agent functions: {e}")
        return False
    
    # ========================================================================
    # Step 2: Create and register agents
    # ========================================================================
    print("\n📋 Step 2: Creating and registering agents...")
    
    # Clear any existing registrations
    AgentRuntimeRegistry.clear()
    
    agents_to_register = [
        ("Project Manager", project_manager_logic, "USER"),
        ("Research Lead", research_lead_logic, "Data Analyst"),
        ("Data Analyst", data_analyst_logic, "Research Lead"),
        ("Report Writer", report_writer_logic, "Research Lead"),
        ("Tech Lead", tech_lead_logic, "Backend Dev"),
        ("Backend Dev", backend_dev_logic, "Tech Lead"),
        ("Frontend Dev", frontend_dev_logic, "Tech Lead"),
    ]
    
    for name, logic, default_target in agents_to_register:
        agent = synqed.Agent(
            name=name,
            description=f"{name} agent",
            logic=logic,
            default_target=default_target
        )
        AgentRuntimeRegistry.register(name, agent)
        print(f"   ✅ Registered: {name}")
    
    registered_roles = AgentRuntimeRegistry.list_roles()
    print(f"\n   Total registered: {len(registered_roles)} agents")
    
    # ========================================================================
    # Step 3: Create task tree structure
    # ========================================================================
    print("\n📋 Step 3: Creating hierarchical task tree...")
    
    root_task_node = TaskTreeNode(
        id="root",
        description="Orchestrate project with specialized teams",
        required_agents=["Project Manager"],
        may_need_subteams=True,
        children=[
            TaskTreeNode(
                id="research-team",
                description="Research and analysis tasks",
                required_agents=["Research Lead", "Data Analyst", "Report Writer"],
                may_need_subteams=False,
                children=[]
            ),
            TaskTreeNode(
                id="dev-team",
                description="Development and implementation tasks",
                required_agents=["Tech Lead", "Backend Dev", "Frontend Dev"],
                may_need_subteams=False,
                children=[]
            )
        ]
    )
    
    print(f"   ✅ Root node: {root_task_node.id}")
    print(f"      - Agents: {root_task_node.required_agents}")
    print(f"      - Children: {len(root_task_node.children)}")
    
    for i, child in enumerate(root_task_node.children):
        print(f"   ✅ Child {i+1}: {child.id}")
        print(f"      - Agents: {child.required_agents}")
    
    # ========================================================================
    # Step 4: Create workspace manager (no API key needed for this)
    # ========================================================================
    print("\n📋 Step 4: Creating workspace manager...")
    
    workspace_manager = WorkspaceManager(
        workspaces_root=Path("/tmp/synqed_orchestrator_validation")
    )
    print(f"   ✅ WorkspaceManager created")
    print(f"      - Root: {workspace_manager.workspaces_root}")
    
    # ========================================================================
    # Step 5: Create workspaces (structure only, no execution)
    # ========================================================================
    print("\n📋 Step 5: Creating workspace structure...")
    
    try:
        # Create root workspace
        root_workspace = await workspace_manager.create_workspace(
            task_tree_node=root_task_node,
            parent_workspace_id=None
        )
        print(f"   ✅ Root workspace: {root_workspace.workspace_id}")
        print(f"      - Agents: {list(root_workspace.agents.keys())}")
        
        # Create Research Team workspace
        research_team_node = root_task_node.children[0]
        research_workspace = await workspace_manager.create_workspace(
            task_tree_node=research_team_node,
            parent_workspace_id=root_workspace.workspace_id
        )
        print(f"   ✅ Research Team: {research_workspace.workspace_id}")
        print(f"      - Agents: {list(research_workspace.agents.keys())}")
        print(f"      - Parent: {research_workspace.parent_id}")
        
        # Create Development Team workspace
        dev_team_node = root_task_node.children[1]
        dev_workspace = await workspace_manager.create_workspace(
            task_tree_node=dev_team_node,
            parent_workspace_id=root_workspace.workspace_id
        )
        print(f"   ✅ Development Team: {dev_workspace.workspace_id}")
        print(f"      - Agents: {list(dev_workspace.agents.keys())}")
        print(f"      - Parent: {dev_workspace.parent_id}")
        
        # ====================================================================
        # Step 6: Validate workspace hierarchy
        # ====================================================================
        print("\n📋 Step 6: Validating workspace hierarchy...")
        
        # Check parent-child relationships
        assert research_workspace.workspace_id in root_workspace.children
        assert dev_workspace.workspace_id in root_workspace.children
        print(f"   ✅ Parent-child relationships correct")
        
        # Check agent counts
        assert len(root_workspace.agents) == 1  # Project Manager only
        assert len(research_workspace.agents) == 3  # 3 research agents
        assert len(dev_workspace.agents) == 3  # 3 dev agents
        print(f"   ✅ Agent counts correct (1 + 3 + 3 = 7 total)")
        
        # ====================================================================
        # Step 7: Clean up
        # ====================================================================
        print("\n📋 Step 7: Cleaning up...")
        await workspace_manager.destroy_workspace(root_workspace.workspace_id)
        print(f"   ✅ Workspaces cleaned up")
        
        print("\n" + "="*70)
        print("  ✅ VALIDATION SUCCESSFUL!")
        print("="*70)
        print("\nThe orchestrator example is ready to run.")
        print("Next step: Run orchestrator_two_teams.py with a valid API key")
        print()
        
        return True
        
    except Exception as e:
        print(f"\n❌ Validation failed: {e}")
        import traceback
        traceback.print_exc()
        return False


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

