ai-orchestrator

#!/usr/bin/env python3
"""
AI Orchestrator - Main CLI Entry Point

Coordinates multiple AI coding assistants to collaborate on software development tasks.
"""

import sys
import os
import logging
from pathlib import Path
import click
from rich.console import Console
from rich.logging import RichHandler
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.table import Table
import yaml

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

from orchestrator import Orchestrator, InteractiveShell
from adapters import AgentCapability


console = Console()


def setup_logging(log_level: str, log_file: str):
    """Setup logging configuration."""
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format="%(message)s",
        handlers=[
            RichHandler(rich_tracebacks=True, console=console),
            logging.FileHandler(log_file)
        ]
    )


@click.group()
@click.version_option(version='1.0.0')
def cli():
    """
    AI Orchestrator - Coordinate multiple AI coding assistants.

    This tool enables Codex, Gemini, Claude, and Copilot to collaborate
    on software development tasks with automated implementation, review,
    and refinement cycles.
    """
    pass


@cli.command()
@click.argument('task')
@click.option('--workflow', '-w', default='default',
              help='Workflow to use (default, quick, thorough, etc.)')
@click.option('--config', '-c', type=click.Path(exists=True),
              help='Path to configuration file')
@click.option('--max-iterations', '-m', type=int,
              help='Maximum number of refinement iterations')
@click.option('--output', '-o', type=click.Path(),
              help='Output directory for generated code')
@click.option('--verbose', '-v', is_flag=True,
              help='Enable verbose logging')
@click.option('--dry-run', is_flag=True,
              help='Show execution plan without running')
@click.option('--agents', help='Comma-separated list of agents to use')
def run(task, workflow, config, max_iterations, output, verbose, dry_run, agents):
    """
    Execute a task using AI agent collaboration.

    Example:
        ai-orchestrator run "Create a REST API with user authentication"
    """
    # Load config
    config_path = config or Path(__file__).parent / "config" / "agents.yaml"

    # Setup logging
    log_level = 'DEBUG' if verbose else 'INFO'
    setup_logging(log_level, 'ai-orchestrator.log')

    # Create orchestrator
    try:
        orchestrator = Orchestrator(config_path)
    except Exception as e:
        console.print(f"[red]Error initializing orchestrator: {e}[/red]")
        sys.exit(1)

    # Show header
    console.print(Panel.fit(
        f"[bold cyan]AI Orchestrator[/bold cyan]\n"
        f"Task: {task}\n"
        f"Workflow: {workflow}",
        border_style="cyan"
    ))

    # Show available agents
    available = orchestrator.get_available_agents()
    console.print(f"\n[green]Available agents:[/green] {', '.join(available)}")

    if not available:
        console.print("[red]No agents available! Please install and configure AI CLI tools.[/red]")
        sys.exit(1)

    # Dry run - show plan
    if dry_run:
        console.print("\n[yellow]Dry run - execution plan:[/yellow]")
        workflows = orchestrator.config.get('workflows', {})
        workflow_config = workflows.get(workflow, [])

        table = Table(title=f"Workflow: {workflow}")
        table.add_column("Step", style="cyan")
        table.add_column("Agent", style="green")
        table.add_column("Task", style="yellow")

        for i, step in enumerate(workflow_config, 1):
            table.add_row(
                str(i),
                step.get('agent', 'unknown'),
                step.get('task', 'unknown')
            )

        console.print(table)
        return

    # Execute task
    console.print("\n[bold]Starting execution...[/bold]\n")

    try:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task_progress = progress.add_task("[cyan]Orchestrating agents...", total=None)

            results = orchestrator.execute_task(
                task=task,
                workflow_name=workflow,
                max_iterations=max_iterations
            )

            progress.update(task_progress, completed=True)

        # Display results
        console.print("\n[bold green]Execution complete![/bold green]\n")

        # Show iterations
        for i, iteration in enumerate(results.get('iterations', []), 1):
            console.print(f"[bold]Iteration {i}:[/bold]")

            for step in iteration.get('steps', []):
                status = "✓" if step.get('success') else "✗"
                color = "green" if step.get('success') else "red"
                agent = step.get('agent')
                task_type = step.get('task')

                console.print(f"  [{color}]{status}[/{color}] {agent} - {task_type}")

                if step.get('suggestions'):
                    console.print(f"     Suggestions: {len(step['suggestions'])}")

        # Show final output
        if results.get('final_output'):
            console.print("\n[bold]Final Output:[/bold]")
            console.print(Panel(results['final_output'][:500] + "..." if len(results['final_output']) > 500 else results['final_output']))

        # Show success status
        if results.get('success'):
            console.print("\n[bold green]Task completed successfully! ✓[/bold green]")
        else:
            console.print("\n[bold yellow]Task completed with some issues[/bold yellow]")

    except Exception as e:
        console.print(f"\n[bold red]Error during execution: {e}[/bold red]")
        if verbose:
            console.print_exception()
        sys.exit(1)


@cli.command()
@click.option('--config', '-c', type=click.Path(exists=True),
              help='Path to configuration file')
def agents(config):
    """List available agents and their status."""
    config_path = config or Path(__file__).parent / "config" / "agents.yaml"

    try:
        orchestrator = Orchestrator(config_path)
    except Exception as e:
        console.print(f"[red]Error: {e}[/red]")
        sys.exit(1)

    table = Table(title="AI Agents")
    table.add_column("Agent", style="cyan")
    table.add_column("Status", style="green")
    table.add_column("Command", style="yellow")
    table.add_column("Role", style="magenta")

    agents_config = orchestrator.config.get('agents', {})

    for name, adapter in orchestrator.adapters.items():
        agent_config = agents_config.get(name, {})
        table.add_row(
            name,
            "✓ Available",
            adapter.command,
            agent_config.get('role', 'N/A')
        )

    # Show unavailable agents
    for name, agent_config in agents_config.items():
        if name not in orchestrator.adapters:
            status = "✗ Not available"
            if not agent_config.get('enabled'):
                status = "○ Disabled"

            table.add_row(
                name,
                status,
                agent_config.get('command', 'N/A'),
                agent_config.get('role', 'N/A')
            )

    console.print(table)


@cli.command()
@click.option('--config', '-c', type=click.Path(exists=True),
              help='Path to configuration file')
def workflows(config):
    """List available workflows."""
    config_path = config or Path(__file__).parent / "config" / "agents.yaml"

    with open(config_path, 'r') as f:
        config_data = yaml.safe_load(f)

    workflows_config = config_data.get('workflows', {})

    table = Table(title="Available Workflows")
    table.add_column("Workflow", style="cyan")
    table.add_column("Steps", style="green")
    table.add_column("Description", style="yellow")

    for name, steps in workflows_config.items():
        step_summary = " → ".join([step.get('agent', '?') for step in steps])
        description = steps[0].get('description', 'N/A') if steps else 'N/A'

        table.add_row(name, step_summary, description)

    console.print(table)


@cli.command()
@click.option('--config', '-c', type=click.Path(exists=True),
              help='Path to configuration file')
def validate(config):
    """Validate configuration and check agent availability."""
    config_path = config or Path(__file__).parent / "config" / "agents.yaml"

    console.print("[bold]Validating configuration...[/bold]\n")

    # Check config file exists
    if not Path(config_path).exists():
        console.print(f"[red]✗ Config file not found: {config_path}[/red]")
        sys.exit(1)

    console.print(f"[green]✓ Config file found: {config_path}[/green]")

    # Try to load config
    try:
        with open(config_path, 'r') as f:
            config_data = yaml.safe_load(f)
        console.print("[green]✓ Config file is valid YAML[/green]")
    except Exception as e:
        console.print(f"[red]✗ Invalid YAML: {e}[/red]")
        sys.exit(1)

    # Check required sections
    required_sections = ['agents', 'workflows', 'settings']
    for section in required_sections:
        if section in config_data:
            console.print(f"[green]✓ Section '{section}' present[/green]")
        else:
            console.print(f"[yellow]! Section '{section}' missing[/yellow]")

    # Initialize orchestrator
    try:
        orchestrator = Orchestrator(config_path)
        console.print(f"\n[green]✓ Orchestrator initialized successfully[/green]")
        console.print(f"[green]✓ {len(orchestrator.adapters)} agents available[/green]")
    except Exception as e:
        console.print(f"\n[red]✗ Orchestrator initialization failed: {e}[/red]")
        sys.exit(1)

    console.print("\n[bold green]Configuration is valid! ✓[/bold green]")


@cli.command()
def version():
    """Show version information."""
    console.print("[bold cyan]AI Orchestrator[/bold cyan] version 1.0.0")
    console.print("A collaborative AI coding assistant orchestration system")


@cli.command()
@click.option('--config', '-c', type=click.Path(exists=True),
              help='Path to configuration file')
@click.option('--workflow', '-w', default='default',
              help='Default workflow to use')
def shell(config, workflow):
    """
    Start an interactive shell for multi-round conversations.

    The interactive shell provides a REPL-style interface similar to
    Claude Code and Codex CLIs, allowing multi-round conversations,
    context preservation, and iterative development.

    Examples:
        ai-orchestrator shell
        ai-orchestrator shell --workflow thorough
    """
    config_path = config or Path(__file__).parent / "config" / "agents.yaml"

    # Setup logging for shell
    setup_logging('INFO', 'ai-orchestrator.log')

    try:
        # Create and start interactive shell
        interactive_shell = InteractiveShell(config_path)
        interactive_shell.history.workflow = workflow
        interactive_shell.start()
    except KeyboardInterrupt:
        console.print("\n[yellow]Exiting...[/yellow]")
    except Exception as e:
        console.print(f"[red]Error starting shell: {e}[/red]")
        sys.exit(1)


@cli.command()
@click.option('--config', '-c', type=click.Path(exists=True),
              help='Path to configuration file')
@click.option('--workflow', '-w', default='default',
              help='Default workflow to use')
def interactive(config, workflow):
    """
    Alias for 'shell' command - start interactive mode.

    Same as running 'ai-orchestrator shell'.
    """
    # Just call shell with same parameters
    ctx = click.get_current_context()
    ctx.invoke(shell, config=config, workflow=workflow)


if __name__ == '__main__':
    cli()