src/delegation_mcp/installer/task_mapper.py

"""Task-to-agent mapping with intelligent suggestions.

This module provides functionality for mapping task categories to agents
with intelligent suggestions based on agent capabilities and routing strategies.
"""

import logging
from typing import TypedDict

from rich.console import Console
from rich.prompt import Confirm, Prompt
from rich.table import Table

from .agent_profiles import (
    get_agent_profile, 
    ROUTING_PRESETS, 
    DEFAULT_ROUTING_RULES,
    RoutingPreset
)

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


class TaskCategory(TypedDict):
    """Definition of a task category."""
    key: str
    name: str
    description: str
    pattern_examples: list[str]


# Task categories for delegation
TASK_CATEGORIES: list[TaskCategory] = [
    {
        "key": "security_audit",
        "name": "Security Audit",
        "description": "Security audits, vulnerability scans, safety checks",
        "pattern_examples": [
            "security", "vulnerability", "audit", "CVE", "harden", "secure", "protect",
            "lock down", "access control", "permissions", "rules", "firestore rules",
            "authentication", "authorization", "encrypt", "expose", "leak", "breach",
            "attack", "threat", "OWASP", "XSS", "injection", "sanitize", "exploit",
        ],
    },
    {
        "key": "code_review",
        "name": "Code Review",
        "description": "Code quality review, best practices analysis",
        "pattern_examples": [
            "review", "code quality", "best practices", "lint", "improve", "clean up",
            "tech debt", "smell", "anti-pattern", "convention", "standards",
            "maintainability", "readability", "code analysis",
        ],
    },
    {
        "key": "architecture",
        "name": "Architecture",
        "description": "System design, architecture planning, complex reasoning",
        "pattern_examples": [
            "architecture", "design", "system design", "structure", "organize", "plan",
            "approach", "strategy", "pattern", "blueprint", "diagram", "flow", "schema",
        ],
    },
    {
        "key": "refactoring",
        "name": "Refactoring",
        "description": "Code refactoring, cleanup, optimization",
        "pattern_examples": [
            "refactor", "cleanup", "optimize code", "rename", "restructure", "reorganize",
            "simplify", "DRY", "extract", "inline", "consolidate", "modularize",
        ],
    },
    {
        "key": "quick_fix",
        "name": "Quick Fixes",
        "description": "Rapid bug fixes, small code changes",
        "pattern_examples": [
            "fix", "bug", "quick change", "error", "crash", "broken", "not working",
            "issue", "problem", "patch", "hotfix",
        ],
    },
    {
        "key": "documentation",
        "name": "Documentation",
        "description": "README files, API docs, code comments",
        "pattern_examples": [
            "documentation", "docs", "README", "comments", "comment", "explain", "describe",
            "guide", "tutorial", "how-to", "API docs", "docstring", "examples",
        ],
    },
    {
        "key": "testing",
        "name": "Testing",
        "description": "Unit tests, integration tests, test coverage",
        "pattern_examples": [
            "test", "testing", "coverage", "unit test", "integration test", "e2e",
            "spec", "assertion", "mock", "stub", "test case", "test suite",
        ],
    },
    {
        "key": "performance",
        "name": "Performance",
        "description": "Performance analysis and optimization",
        "pattern_examples": [
            "performance", "optimize", "speed", "slow", "latency", "throughput",
            "bottleneck", "profiling", "benchmark", "memory", "CPU", "scalability",
        ],
    },
    {
        "key": "browser_interaction",
        "name": "Browser Interaction",
        "description": "Browser automation, web scraping, UI testing",
        "pattern_examples": ["browser", "selenium", "playwright", "chrome"],
    },
    {
        "key": "git_operations",
        "name": "Git Operations",
        "description": "Git workflows, repository management",
        "pattern_examples": [
            "git", "commit", "merge", "branch", "push", "pull", "rebase", "cherry-pick",
            "stash", "tag", "history", "checkout", "reset", "revert",
        ],
    },
    {
        "key": "shell_tasks",
        "name": "Shell/Terminal",
        "description": "Shell scripting, terminal commands",
        "pattern_examples": [
            "shell", "terminal", "bash", "script", "command", "CLI", "automation",
            "cron", "env", "environment variables", "path", "execute",
        ],
    },
    {
        "key": "exploration",
        "name": "Code Exploration",
        "description": "Code exploration, dependency tracing, implementation analysis",
        "pattern_examples": ["how does", "trace the flow", "what files implement", "understand the implementation", "map dependencies"],
    },
    {
        "key": "debugging",
        "name": "Debugging",
        "description": "Bug investigation, error analysis, root cause identification",
        "pattern_examples": ["debug", "why is failing", "investigate", "find the cause", "troubleshoot"],
    },
    {
        "key": "impact_analysis",
        "name": "Impact Analysis",
        "description": "Dependency analysis, usage finding, breaking change assessment",
        "pattern_examples": ["what would break", "find all usages", "what depends on", "impact of changing", "all references"],
    },
    {
        "key": "general",
        "name": "General Tasks",
        "description": "Default for tasks that don't fit specific categories",
        "pattern_examples": ["general", "misc", "other"],
    },
]


class TaskMapper:
    """Manages task-to-agent mapping during installation."""

    def __init__(self):
        """Initialize the task mapper."""
        self.task_mappings: dict[str, str] = {}
        self.selected_strategy: str = "balanced"

    def select_strategy(self) -> str:
        """
        Prompt user to select a routing strategy.
        
        Returns:
            Selected strategy key
        """
        console.print("\n[bold]Delegation Strategy[/bold]")
        console.print("Choose how tasks should be distributed among agents:\n")
        
        table = Table(show_header=True, header_style="bold magenta")
        table.add_column("Option", style="cyan", justify="center")
        table.add_column("Strategy", style="green")
        table.add_column("Description", style="white")
        table.add_column("Priorities", style="yellow")
        
        strategies = list(ROUTING_PRESETS.items())
        
        for i, (key, preset) in enumerate(strategies, 1):
            priorities = f"Cost: {preset['cost_priority']}, Quality: {preset['quality_priority']}"
            table.add_row(str(i), preset["name"], preset["description"], priorities)
            
        console.print(table)
        console.print("\n")
        
        choices = [str(i) for i in range(1, len(strategies) + 1)]
        default_idx = [k for k, _ in strategies].index("balanced") + 1
        
        selection = Prompt.ask(
            "Select strategy",
            choices=choices,
            default=str(default_idx)
        )
        
        selected_key = strategies[int(selection) - 1][0]
        self.selected_strategy = selected_key
        
        console.print(f"\n[green]✓[/green] Selected: {ROUTING_PRESETS[selected_key]['name']}\n")
        return selected_key

    def suggest_mappings(self, agent_names: list[str], strategy_key: str) -> dict[str, tuple[str, str]]:
        """
        Generate intelligent mapping suggestions based on strategy and agent capabilities.
        
        Args:
            agent_names: List of available agent names
            strategy_key: Key of the selected routing strategy
            
        Returns:
            Dictionary of task_key -> (suggested_agent, reasoning)
        """
        suggestions: dict[str, tuple[str, str]] = {}
        preset = ROUTING_PRESETS[strategy_key]
        
        # Helper to find best agent from a list of preferred ones
        def find_best_available(preferred: list[str], fallback_reason: str) -> tuple[str, str]:
            for agent in preferred:
                if agent in agent_names:
                    # Find specific reason from rules if available
                    return agent, fallback_reason
            
            # Fallback logic based on strategy
            if preset["cost_priority"] == "high":
                # Prefer free/local agents
                for agent in agent_names:
                    profile = get_agent_profile(agent)
                    if profile["cost_tier"] == "free":
                        return agent, "Selected for cost efficiency"
            
            if preset["quality_priority"] == "high":
                # Prefer Claude/Gemini
                for agent in ["claude", "gemini"]:
                    if agent in agent_names:
                        return agent, "Selected for high quality"
                        
            # Default to first available
            return agent_names[0], "Best available option"

        for category in TASK_CATEGORIES:
            task_key = category["key"]
            
            # Get default rule
            rule = DEFAULT_ROUTING_RULES.get(task_key)
            if not rule:
                suggestions[task_key] = (agent_names[0], "Default assignment")
                continue
                
            # Apply strategy overrides
            preferred = rule["preferred"]
            reason = rule["reason"]
            
            if strategy_key == "cost_optimized":
                # Prioritize free agents
                free_agents = [a for a in agent_names if get_agent_profile(a)["cost_tier"] == "free"]
                if free_agents:
                    preferred = free_agents + preferred
                    reason = "Cost optimized choice"
                    
            elif strategy_key == "speed_first":
                # Prioritize fast agents
                fast_agents = [a for a in agent_names if get_agent_profile(a)["response_speed"] == "fast"]
                if fast_agents:
                    preferred = fast_agents + preferred
                    reason = "Optimized for speed"
            
            elif strategy_key == "token_saver":
                # Prioritize large context or concise agents
                # (Simplified logic: prefer Gemini for context, Aider for conciseness)
                if task_key in ["architecture", "exploration"]:
                    preferred = ["gemini"] + preferred
                    reason = "Large context window"
                else:
                    preferred = ["aider"] + preferred
                    reason = "Concise responses"

            # Find best agent
            agent, final_reason = find_best_available(preferred, reason)
            suggestions[task_key] = (agent, final_reason)

        return suggestions

    def display_suggestions(
        self, 
        suggestions: dict[str, tuple[str, str]],
        agent_names: list[str]
    ) -> None:
        """
        Display mapping suggestions in a formatted table.
        
        Args:
            suggestions: Dictionary of task_key -> (agent, reasoning)
            agent_names: List of available agent names for context
        """
        table = Table(
            title=f"Suggested Mappings ({ROUTING_PRESETS[self.selected_strategy]['name']})",
            show_header=True,
            header_style="bold magenta"
        )
        table.add_column("Task Category", style="cyan", no_wrap=True)
        table.add_column("Suggested Agent", style="green")
        table.add_column("Reasoning", style="yellow")

        # Create task key to category mapping for lookup
        category_map = {cat["key"]: cat for cat in TASK_CATEGORIES}

        for task_key, (agent, reasoning) in suggestions.items():
            category = category_map.get(task_key)
            if category:
                task_name = category["name"]
                table.add_row(task_name, agent, reasoning)

        console.print("\n")
        console.print(table)
        console.print("\n")

    def prompt_task_assignments(
        self, 
        agent_names: list[str],
        suggestions: dict[str, tuple[str, str]]
    ) -> dict[str, str]:
        """
        Interactive prompt for task-to-agent assignment.
        
        Args:
            agent_names: List of available agent names
            suggestions: Pre-computed suggestions
            
        Returns:
            Dictionary of task_key -> agent_name
        """
        self.display_suggestions(suggestions, agent_names)

        console.print("[bold]Task Assignment Configuration[/bold]")
        console.print("You can accept all suggestions or customize individual mappings.\n")

        # Ask if user wants to use all suggestions
        accept_all = Confirm.ask(
            "Accept all suggested mappings?",
            default=True
        )

        if accept_all:
            self.task_mappings = {
                task_key: agent 
                for task_key, (agent, _) in suggestions.items()
            }
            console.print("\n[green]✓[/green] Using all suggested mappings\n")
            return self.task_mappings

        # Custom assignment
        console.print("\nCustomize task assignments:\n")
        self.task_mappings = {}

        # Create task key to category mapping
        category_map = {cat["key"]: cat for cat in TASK_CATEGORIES}

        for task_key, (suggested_agent, reasoning) in suggestions.items():
            category = category_map.get(task_key)
            if not category:
                continue

            task_name = category["name"]
            description = category["description"]

            console.print(f"\n[cyan]{task_name}[/cyan]: {description}")
            console.print(f"  Suggested: [green]{suggested_agent}[/green] ({reasoning})")

            # Ask if user wants to change
            use_suggestion = Confirm.ask(
                f"  Use {suggested_agent} for {task_name}?",
                default=True
            )

            if use_suggestion:
                self.task_mappings[task_key] = suggested_agent
                console.print(f"  [green]✓[/green] Assigned to {suggested_agent}")
            else:
                # Let user pick an agent
                console.print(f"  Available agents: {', '.join(agent_names)}")
                
                while True:
                    chosen_agent = Prompt.ask(
                        f"  Select agent for {task_name}",
                        choices=agent_names,
                        default=suggested_agent
                    )
                    
                    if chosen_agent in agent_names:
                        self.task_mappings[task_key] = chosen_agent
                        console.print(f"  [green]✓[/green] Assigned to {chosen_agent}")
                        break
                    else:
                        console.print(f"  [red]✗[/red] Invalid agent. Please choose from: {', '.join(agent_names)}")

        console.print(f"\n[green]✓[/green] Task assignment configuration complete\n")
        return self.task_mappings

    def get_task_mappings(self) -> dict[str, str]:
        """
        Get the task-to-agent mappings.
        
        Returns:
            Dictionary of task_key -> agent_name
        """
        return self.task_mappings

    def map_tasks(self, agent_names: list[str]) -> dict[str, str]:
        """
        Complete task mapping flow with suggestions and user input.
        
        Args:
            agent_names: List of available agent names
            
        Returns:
            Dictionary of task_key -> agent_name
        """
        if len(agent_names) < 2:
            logger.warning("Need at least 2 agents for task mapping")
            return {}

        # Select strategy
        strategy_key = self.select_strategy()

        # Generate suggestions based on strategy
        suggestions = self.suggest_mappings(agent_names, strategy_key)

        # Get user assignments
        return self.prompt_task_assignments(agent_names, suggestions)