src/delegation_mcp/installer/config_generator.py

"""Configuration file generator for orchestrators and delegation rules.

This module generates YAML configuration files based on user selections
during installation.
"""

import logging
import shutil
from datetime import datetime
from pathlib import Path
from typing import Any

import yaml
from rich.console import Console
from rich.prompt import Confirm

from ..agent_discovery import AgentMetadata
from .agent_profiles import get_agent_profile
from .task_mapper import TASK_CATEGORIES

logger = logging.getLogger(__name__)
console = Console()


class ConfigGenerator:
    """Generates configuration files based on user selections."""

    def __init__(self):
        """Initialize the config generator."""
        self.orchestrators_config: dict[str, Any] = {}
        self.delegation_config: dict[str, Any] = {}

    def check_existing_configs(self, project_dir: Path) -> bool:
        """
        Check if configuration files already exist.
        
        Args:
            project_dir: Project directory path
            
        Returns:
            True if configs exist, False otherwise
        """
        config_dir = project_dir / "config"
        orchestrators_file = config_dir / "orchestrators.yaml"
        delegation_file = config_dir / "delegation_rules.yaml"

        return orchestrators_file.exists() or delegation_file.exists()

    def prompt_existing_configs(self, project_dir: Path) -> str:
        """
        Prompt user for how to handle existing configs.
        
        Args:
            project_dir: Project directory path
            
        Returns:
            User choice: "overwrite", "backup", or "skip"
        """
        console.print("\n[yellow]⚠ Existing configuration detected[/yellow]")
        console.print("\nConfiguration files already exist in this project.")
        console.print("You can:")
        console.print("  [cyan]O[/cyan]verwrite - Replace existing files with new configuration")
        console.print("  [cyan]B[/cyan]ackup   - Backup existing files and create new ones")
        console.print("  [cyan]S[/cyan]kip     - Keep existing files unchanged\n")

        choice = ""
        while choice not in ["o", "b", "s"]:
            choice = input("Choose an option (O/B/S): ").lower().strip()

        choice_map = {"o": "overwrite", "b": "backup", "s": "skip"}
        return choice_map[choice]

    def backup_existing_configs(self, project_dir: Path) -> None:
        """
        Backup existing configuration files.
        
        Args:
            project_dir: Project directory path
        """
        config_dir = project_dir / "config"
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        files_to_backup = [
            "orchestrators.yaml",
            "delegation_rules.yaml",
        ]

        for filename in files_to_backup:
            source = config_dir / filename
            if source.exists():
                backup_name = f"{source.stem}.backup_{timestamp}{source.suffix}"
                backup_path = config_dir / backup_name
                shutil.copy2(source, backup_path)
                console.print(f"[green]✓[/green] Backed up {filename} to {backup_name}")

    def generate_orchestrators_yaml(
        self,
        selected_agents: list[str],
        agent_metadata: dict[str, AgentMetadata],
    ) -> dict[str, Any]:
        """
        Generate orchestrators.yaml configuration.
        
        Args:
            selected_agents: List of selected agent names
            agent_metadata: Dictionary of agent name to metadata
            
        Returns:
            Orchestrators configuration dictionary
        """
        config: dict[str, Any] = {"orchestrators": {}}

        for agent_name in selected_agents:
            metadata = agent_metadata.get(agent_name)
            if not metadata:
                continue

            # Get command and args from metadata
            # command can be a string or list[str]
            if isinstance(metadata.command, list):
                command = metadata.command[0] if metadata.command else agent_name
                args = metadata.command[1:] if len(metadata.command) > 1 else []
            else:
                command = metadata.command or agent_name
                args = []

            # Build orchestrator config
            orch_config: dict[str, Any] = {
                "name": agent_name,
                "command": command,
                "args": args,
                "enabled": True,
                "env": {},
                "timeout": 300,
                "max_retries": 3,
            }

            config["orchestrators"][agent_name] = orch_config

        self.orchestrators_config = config
        return config

    def generate_delegation_rules_yaml(
        self,
        selected_agents: list[str],
        task_mappings: dict[str, str],
        agent_metadata: dict[str, AgentMetadata],
        primary_orchestrator: str,
    ) -> dict[str, Any]:
        """
        Generate delegation_rules.yaml configuration.
        
        Args:
            selected_agents: List of selected agent names
            task_mappings: Dictionary of task_key -> agent_name
            agent_metadata: Dictionary of agent name to metadata
            primary_orchestrator: Name of primary orchestrator
            
        Returns:
            Delegation rules configuration dictionary
        """
        config: dict[str, Any] = {
            "orchestrator": primary_orchestrator,
            "routing_strategy": "hybrid",
            "orchestrators": {},
            "rules": [],
        }

        # Build orchestrators section with capabilities
        for agent_name in selected_agents:
            metadata = agent_metadata.get(agent_name)
            if not metadata:
                continue

            profile = get_agent_profile(agent_name)
            capabilities = profile["capabilities"]

            # Get command and args from metadata
            # command can be a string or list[str]
            if isinstance(metadata.command, list):
                command = metadata.command[0] if metadata.command else agent_name
                args = metadata.command[1:] if len(metadata.command) > 1 else []
            else:
                command = metadata.command or agent_name
                args = []

            orch_config: dict[str, Any] = {
                "name": agent_name,
                "command": command,
                "args": args,
                "enabled": True,
                "env": {},
                "timeout": 300,
                "max_retries": 3,
                "cost_per_1k_tokens": 0.001,
                "capabilities": dict(capabilities),
            }

            config["orchestrators"][agent_name] = orch_config

        # Build delegation rules from task mappings
        # Create task key to category mapping
        category_map = {cat["key"]: cat for cat in TASK_CATEGORIES}

        priority = 10  # Start with high priority
        for task_key, agent_name in task_mappings.items():
            category = category_map.get(task_key)
            if not category:
                continue

            # Build pattern from examples
            pattern_parts = category["pattern_examples"]
            pattern = "|".join(pattern_parts)

            rule: dict[str, Any] = {
                "delegate_to": agent_name,
                "description": category["description"],
                "pattern": pattern,
                "priority": max(1, priority),
                "requires_approval": False,
            }

            config["rules"].append(rule)
            priority -= 1  # Decrease priority for next rule

        # Add fallback rule (general tasks)
        if "general" not in task_mappings:
            # Use primary orchestrator as fallback
            fallback_agent = primary_orchestrator
        else:
            fallback_agent = task_mappings["general"]

        fallback_rule: dict[str, Any] = {
            "delegate_to": fallback_agent,
            "description": "General queries and fallback",
            "pattern": ".*",
            "priority": 1,
            "requires_approval": False,
        }
        config["rules"].append(fallback_rule)

        self.delegation_config = config
        return config

    def save_configs(self, project_dir: Path) -> None:
        """
        Save configuration files to disk.
        
        Args:
            project_dir: Project directory path
        """
        config_dir = project_dir / "config"
        config_dir.mkdir(parents=True, exist_ok=True)

        # Save orchestrators.yaml
        orchestrators_file = config_dir / "orchestrators.yaml"
        with open(orchestrators_file, "w", encoding="utf-8") as f:
            yaml.dump(
                self.orchestrators_config,
                f,
                default_flow_style=False,
                sort_keys=False,
            )
        console.print(f"[green]✓[/green] Generated {orchestrators_file}")

        # Save delegation_rules.yaml
        delegation_file = config_dir / "delegation_rules.yaml"
        with open(delegation_file, "w", encoding="utf-8") as f:
            # Add header comment
            f.write("# Delegation MCP Configuration\n")
            f.write("# Auto-generated based on user selections\n\n")
            yaml.dump(
                self.delegation_config,
                f,
                default_flow_style=False,
                sort_keys=False,
            )
        console.print(f"[green]✓[/green] Generated {delegation_file}")

    def generate_configs(
        self,
        selected_agents: list[str],
        task_mappings: dict[str, str],
        agent_metadata: dict[str, AgentMetadata],
        primary_orchestrator: str,
        project_dir: Path,
        scope: str = "local",
    ) -> None:
        """
        Complete configuration generation flow.
        
        Args:
            selected_agents: List of selected agent names
            task_mappings: Dictionary of task_key -> agent_name
            agent_metadata: Dictionary of agent name to metadata
            primary_orchestrator: Name of primary orchestrator
            project_dir: Project directory path
            scope: Installation scope - "local" for project-level, "user" for user-level
        """
        # Determine config directory based on scope
        if scope == "user":
            config_base_dir = Path.home() / ".delegation-mcp"
        else:
            config_base_dir = project_dir
        
        # Check for existing configs
        if self.check_existing_configs(config_base_dir):
            choice = self.prompt_existing_configs(config_base_dir)

            if choice == "skip":
                console.print("\n[yellow]⚠[/yellow] Keeping existing configuration files\n")
                return
            elif choice == "backup":
                self.backup_existing_configs(config_base_dir)

        # Generate configurations
        self.generate_orchestrators_yaml(selected_agents, agent_metadata)
        self.generate_delegation_rules_yaml(
            selected_agents, task_mappings, agent_metadata, primary_orchestrator
        )

        # Save to disk
        self.save_configs(config_base_dir)
        console.print("\n[green]✓[/green] Configuration files generated successfully\n")