#!/usr/bin/env python3
"""
Demonstration of the User Node Communication System.

This example shows how agents can communicate with users through the User node,
supporting both synchronous (terminal) and asynchronous (web) communication patterns.
"""

import asyncio
import logging
from pathlib import Path
import sys

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from marsys.agents import Agent
from marsys.agents.registry import AgentRegistry
from marsys.coordination import Orchestra
from marsys.coordination.communication import (
    CommunicationManager,
    CommunicationMode
)
from marsys.coordination.communication.channels import TerminalChannel
from marsys.models import ModelConfig

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


def create_agents():
    """Create demonstration agents."""
    
    # Market Research Coordinator
    coordinator_config = ModelConfig(
        type="api",
        provider="openai",
        name="gpt-4-mini",
        temperature=0.7,
        max_tokens=1000
    )
    
    coordinator = Agent(
        model_config=coordinator_config,
        name="Coordinator",
        description="""You are a Market Research Coordinator.

Your role is to:
1. Understand the user's research needs
2. Identify areas that need clarification
3. Return to the User node to ask clarifying questions
4. Use the clarified information to create a research plan

When you need clarification, respond with:
{
    "next_action": "invoke_agent",
    "action_input": "User",
    "message": "Your clarifying question here"
}

After receiving clarification, create a research plan and pass it to the Researcher.""",
        allowed_peers=["User", "Researcher"]
    )
    
    # Researcher Agent
    researcher_config = ModelConfig(
        type="api",
        provider="openai", 
        name="gpt-4-mini",
        temperature=0.5,
        max_tokens=1500
    )
    
    researcher = Agent(
        model_config=researcher_config,
        name="Researcher",
        description="""You are a Market Researcher.

Your role is to:
1. Execute research based on the coordinator's plan
2. If you encounter ambiguities during research, return to User for clarification
3. Compile findings and pass to the Synthesizer

When you need user input, respond with:
{
    "next_action": "invoke_agent",
    "action_input": "User",
    "message": "Your question about the research"
}""",
        allowed_peers=["User", "Synthesizer"]
    )
    
    # Synthesizer Agent
    synthesizer_config = ModelConfig(
        type="api",
        provider="openai",
        name="gpt-4-mini",
        temperature=0.3,
        max_tokens=2000
    )
    
    synthesizer = Agent(
        model_config=synthesizer_config,
        name="Synthesizer",
        description="""You are a Research Synthesizer.

Your role is to:
1. Compile research findings into a coherent report
2. Ask the user for preferences on report format if needed
3. Generate the final report

When asking about preferences, respond with:
{
    "next_action": "invoke_agent", 
    "action_input": "User",
    "message": "Question about report preferences"
}

For the final report, respond with:
{
    "next_action": "final_response",
    "content": "Your synthesized report here"
}""",
        allowed_peers=["User"]
    )
    
    return {
        "Coordinator": coordinator,
        "Researcher": researcher,
        "Synthesizer": synthesizer
    }


async def demo_sync_terminal():
    """Demonstrate synchronous terminal-based communication."""
    print("\n" + "="*80)
    print("DEMO 1: Synchronous Terminal Communication")
    print("="*80)
    
    # Clear agent registry
    AgentRegistry.clear()
    
    # Create agents
    agents = create_agents()
    
    # Create communication infrastructure
    comm_manager = CommunicationManager()
    
    # Create and register terminal channel
    terminal_channel = TerminalChannel("terminal_demo")
    await terminal_channel.start()
    comm_manager.register_channel(terminal_channel)
    
    # Create session
    session_id = "terminal_session_001"
    comm_manager.assign_channel_to_session(session_id, "terminal_demo")
    
    # Create Orchestra with communication
    orchestra = Orchestra(
        agent_registry=AgentRegistry,
        communication_manager=comm_manager
    )
    
    # Define topology with bidirectional User communication
    topology = {
        "agents": ["User", "Coordinator", "Researcher", "Synthesizer"],
        "flows": [
            # User can start the flow
            "User -> Coordinator",

            # Agents can return to User for questions
            "Coordinator -> User",
            "User -> Coordinator",  # User responds back

            "Coordinator -> Researcher",

            "Researcher -> User",
            "User -> Researcher",  # User responds back

            "Researcher -> Synthesizer",

            "Synthesizer -> User",
            "User -> Synthesizer",  # User responds back

            # Final flow
            "Synthesizer -> Coordinator"  # Could go back for final review
        ],
        "rules": [
            "max_steps(20)",  # Prevent infinite loops
            "timeout(600)"     # 10 minute timeout
        ]
    }
    
    # Initial task
    initial_task = "I need to analyze the market for sustainable packaging solutions"
    
    print(f"\nInitial Task: {initial_task}")
    print("\nStarting multi-agent workflow with user interaction...")
    print("(Agents will ask you clarifying questions during execution)\n")
    
    try:
        # Run the orchestration
        result = await orchestra.run(
            task=initial_task,
            topology=topology,
            context={
                "session_id": session_id,
                "communication_mode": "sync"  # Synchronous mode
            }
        )
        
        print("\n" + "="*80)
        print("EXECUTION COMPLETE")
        print("="*80)
        print(f"Success: {result.success}")
        print(f"Total Steps: {result.total_steps}")
        print(f"Final Response Preview:")
        print("-"*40)
        if isinstance(result.final_response, str):
            print(result.final_response[:500] + "..." if len(result.final_response) > 500 else result.final_response)
        else:
            print(result.final_response)
        
    except Exception as e:
        logger.error(f"Demo failed: {e}")
        raise
    finally:
        # Cleanup
        await terminal_channel.stop()
        await comm_manager.cleanup()


async def demo_async_api():
    """Demonstrate asynchronous API-based communication."""
    print("\n" + "="*80)
    print("DEMO 2: Asynchronous Web API Communication")
    print("="*80)
    
    # This would typically integrate with a web framework like FastAPI
    # For demo purposes, we'll simulate the async flow
    
    from marsys.coordination.communication.channels import WebChannel
    
    # Clear agent registry
    AgentRegistry.clear()
    
    # Create agents
    agents = create_agents()
    
    # Create communication infrastructure
    comm_manager = CommunicationManager()
    
    # Create and register web channel
    web_channel = WebChannel("web_demo")
    await web_channel.start()
    comm_manager.register_channel(web_channel)
    
    # Create session
    session_id = "web_session_001"
    comm_manager.assign_channel_to_session(session_id, "web_demo")
    
    # Create Orchestra
    orchestra = Orchestra(
        agent_registry=AgentRegistry,
        communication_manager=comm_manager
    )
    
    # Same topology as before
    topology = {
        "agents": ["User", "Coordinator", "Researcher"],
        "flows": [
            "User -> Coordinator",
            "Coordinator -> User",
            "User -> Coordinator",
            "Coordinator -> Researcher",
            "Researcher -> User",
            "User -> Researcher"
        ]
    }
    
    print("\nSimulating async web communication...")
    print("In a real web app, these would be API calls:\n")
    
    # Start the orchestration (non-blocking)
    task = asyncio.create_task(orchestra.run(
        task="Analyze electric vehicle charging market",
        topology=topology,
        context={
            "session_id": session_id,
            "communication_mode": "async_pubsub"
        }
    ))
    
    # Simulate API polling and responses
    for i in range(3):  # Simulate 3 interactions
        await asyncio.sleep(1)  # Simulate delay
        
        # Check for pending interactions (API: GET /interactions/pending)
        pending = await web_channel.get_pending_interactions(session_id)
        
        if pending:
            print(f"\n[API] GET /interactions/pending")
            print(f"Found {len(pending)} pending interactions:")
            
            for interaction in pending:
                print(f"\nInteraction ID: {interaction['id']}")
                print(f"From: {interaction['from_agent']}")
                print(f"Message: {interaction['message']}")
                
                # Simulate user response (API: POST /interactions/{id}/response)
                user_response = f"Simulated response to question {i+1}"
                print(f"\n[API] POST /interactions/{interaction['id']}/response")
                print(f"Response: {user_response}")
                
                success = await web_channel.submit_response(
                    interaction['id'],
                    user_response
                )
                print(f"Response submitted: {success}")
    
    # Wait for completion
    try:
        result = await asyncio.wait_for(task, timeout=30)
        print("\n" + "="*80)
        print("ASYNC EXECUTION COMPLETE")
        print("="*80)
        print(f"Success: {result.success}")
        
    except asyncio.TimeoutError:
        print("\nExecution timed out (this is expected in demo mode)")
    
    finally:
        # Cleanup
        await web_channel.stop()
        await comm_manager.cleanup()


async def main():
    """Run all demos."""
    print("\n" + "="*80)
    print("USER NODE COMMUNICATION SYSTEM DEMONSTRATION")
    print("="*80)
    print("\nThis demo shows how agents can communicate with users through")
    print("the User node in both synchronous and asynchronous modes.")
    
    # Run synchronous terminal demo
    await demo_sync_terminal()
    
    # Uncomment to run async demo
    # await demo_async_api()
    
    print("\n" + "="*80)
    print("ALL DEMOS COMPLETE")
    print("="*80)


if __name__ == "__main__":
    # Check for API key
    import os
    if not os.getenv("OPENAI_API_KEY"):
        print("ERROR: Please set OPENAI_API_KEY environment variable")
        sys.exit(1)
    
    asyncio.run(main())