"""
Task Types and Routing Table - Maps task categories to optimal agents.

This module defines:
1. TaskType enum for categorizing tasks
2. Agent capability profiles
3. Routing rules that match tasks to agents

The routing strategy considers:
- Task complexity and risk level
- Agent strengths (context window, tools, autonomy modes)
- Cost efficiency
- Current availability and health

Usage:
    from agent_orchestrator.routing.task_types import TaskType, ROUTING_TABLE

    task_type = TaskType.CODE_GENERATION
    agents = ROUTING_TABLE[task_type]
"""

from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Optional

from ..adapters.base import RiskLevel


class TaskType(Enum):
    """Categories of tasks that agents can perform."""

    # Code Generation
    CODE_GENERATION = auto()  # Write new code from scratch
    CODE_COMPLETION = auto()  # Complete partial code
    CODE_REFACTORING = auto()  # Improve existing code structure

    # Analysis
    CODE_REVIEW = auto()  # Review code for issues
    CODE_EXPLANATION = auto()  # Explain what code does
    ARCHITECTURE_ANALYSIS = auto()  # Analyze system architecture
    SECURITY_AUDIT = auto()  # Check for security issues
    PERFORMANCE_ANALYSIS = auto()  # Find performance bottlenecks

    # Testing
    TEST_GENERATION = auto()  # Write unit/integration tests
    TEST_EXECUTION = auto()  # Run tests and report results
    TEST_DEBUGGING = auto()  # Debug failing tests

    # Documentation
    DOCUMENTATION = auto()  # Write docs, README, comments
    API_DOCUMENTATION = auto()  # Generate API docs

    # Infrastructure
    CI_CD_TASKS = auto()  # CI/CD pipeline work
    DEPLOYMENT = auto()  # Deployment tasks
    DEVOPS = auto()  # General DevOps tasks

    # Research
    RESEARCH = auto()  # Research best practices, tools
    CODEBASE_EXPLORATION = auto()  # Explore and understand codebase

    # Debugging
    BUG_INVESTIGATION = auto()  # Investigate bugs
    BUG_FIX = auto()  # Fix identified bugs

    # Multi-file Operations
    LARGE_REFACTOR = auto()  # Refactoring across many files
    MIGRATION = auto()  # Data or code migrations

    # Interactive
    PAIR_PROGRAMMING = auto()  # Interactive coding session


class AgentTool(Enum):
    """Available agent tools/CLIs."""

    CLAUDE_CODE = "claude-code"
    CLAUDE_SDK = "claude-sdk"
    GEMINI_CLI = "gemini-cli"
    CODEX_CLI = "codex-cli"
    OPENAI_SDK = "openai-sdk"


@dataclass
class AgentCapability:
    """Describes an agent's capabilities for a task type."""

    tool: AgentTool
    priority: int = 1  # Lower is higher priority
    risk_level: RiskLevel = RiskLevel.MEDIUM
    autonomy_mode: Optional[str] = None  # For Codex: suggest/auto-edit/full-auto
    max_context_tokens: int = 100000
    supports_streaming: bool = True
    requires_api_key: bool = True
    cost_tier: str = "medium"  # low, medium, high
    strengths: list[str] = field(default_factory=list)
    limitations: list[str] = field(default_factory=list)


# Agent capability profiles
AGENT_PROFILES: dict[AgentTool, AgentCapability] = {
    AgentTool.CLAUDE_CODE: AgentCapability(
        tool=AgentTool.CLAUDE_CODE,
        priority=1,
        risk_level=RiskLevel.MEDIUM,
        max_context_tokens=200000,
        supports_streaming=True,
        requires_api_key=True,
        cost_tier="high",
        strengths=[
            "Complex multi-file edits",
            "Deep codebase understanding",
            "Tool use and file operations",
            "Git operations",
            "Interactive debugging",
        ],
        limitations=[
            "Per-session (no state persistence)",
            "Higher cost",
            "Requires manual approval for risky ops",
        ],
    ),
    AgentTool.CLAUDE_SDK: AgentCapability(
        tool=AgentTool.CLAUDE_SDK,
        priority=2,
        risk_level=RiskLevel.LOW,
        max_context_tokens=200000,
        supports_streaming=True,
        requires_api_key=True,
        cost_tier="high",
        strengths=[
            "Programmatic control",
            "Custom tool definitions",
            "API integration",
            "Structured output",
        ],
        limitations=[
            "No built-in file system access",
            "Requires orchestrator to execute tools",
        ],
    ),
    AgentTool.GEMINI_CLI: AgentCapability(
        tool=AgentTool.GEMINI_CLI,
        priority=2,
        risk_level=RiskLevel.MEDIUM,
        max_context_tokens=1000000,  # 1M context window!
        supports_streaming=True,
        requires_api_key=True,
        cost_tier="low",  # Free tier available
        strengths=[
            "Massive context window (1M tokens)",
            "Search grounding for current info",
            "MCP server support",
            "Free tier (60 req/min, 1000/day)",
            "Good for large codebase analysis",
        ],
        limitations=[
            "Rate limited on free tier",
            "Less precise for complex edits",
        ],
    ),
    AgentTool.CODEX_CLI: AgentCapability(
        tool=AgentTool.CODEX_CLI,
        priority=1,
        risk_level=RiskLevel.MEDIUM,
        autonomy_mode="auto-edit",
        max_context_tokens=128000,
        supports_streaming=True,
        requires_api_key=True,
        cost_tier="medium",
        strengths=[
            "Tiered autonomy modes",
            "CI/CD integration",
            "Fast code generation",
            "Good for automated pipelines",
        ],
        limitations=[
            "Smaller context than Gemini",
            "Less nuanced than Claude",
        ],
    ),
    AgentTool.OPENAI_SDK: AgentCapability(
        tool=AgentTool.OPENAI_SDK,
        priority=3,
        risk_level=RiskLevel.LOW,
        max_context_tokens=128000,
        supports_streaming=True,
        requires_api_key=True,
        cost_tier="medium",
        strengths=[
            "Programmatic control",
            "Function calling",
            "Structured outputs",
            "Wide ecosystem",
        ],
        limitations=[
            "No built-in file operations",
            "Requires tool implementation",
        ],
    ),
}


@dataclass
class RoutingRule:
    """A rule that maps a task type to preferred agents."""

    task_type: TaskType
    preferred_agents: list[AgentTool]
    risk_level: RiskLevel
    context_requirement: str = "normal"  # minimal, normal, large, massive
    autonomy_recommendation: Optional[str] = None
    notes: str = ""


# Routing table: Maps task types to optimal agent configurations
ROUTING_TABLE: dict[TaskType, RoutingRule] = {
    # Code Generation - Claude Code excels, Codex good for quick gen
    TaskType.CODE_GENERATION: RoutingRule(
        task_type=TaskType.CODE_GENERATION,
        preferred_agents=[AgentTool.CLAUDE_CODE, AgentTool.CODEX_CLI, AgentTool.CLAUDE_SDK],
        risk_level=RiskLevel.MEDIUM,
        context_requirement="normal",
        autonomy_recommendation="auto-edit",
        notes="Claude Code for complex, Codex for quick iterations",
    ),
    TaskType.CODE_COMPLETION: RoutingRule(
        task_type=TaskType.CODE_COMPLETION,
        preferred_agents=[AgentTool.CODEX_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.LOW,
        context_requirement="minimal",
        autonomy_recommendation="auto-edit",
        notes="Codex is fast for completions",
    ),
    TaskType.CODE_REFACTORING: RoutingRule(
        task_type=TaskType.CODE_REFACTORING,
        preferred_agents=[AgentTool.CLAUDE_CODE, AgentTool.CODEX_CLI],
        risk_level=RiskLevel.MEDIUM,
        context_requirement="normal",
        autonomy_recommendation="suggest",
        notes="Suggest mode for safety on refactoring",
    ),
    # Analysis - Gemini's large context is valuable
    TaskType.CODE_REVIEW: RoutingRule(
        task_type=TaskType.CODE_REVIEW,
        preferred_agents=[AgentTool.CLAUDE_CODE, AgentTool.GEMINI_CLI, AgentTool.CLAUDE_SDK],
        risk_level=RiskLevel.LOW,
        context_requirement="large",
        notes="Read-only, low risk",
    ),
    TaskType.CODE_EXPLANATION: RoutingRule(
        task_type=TaskType.CODE_EXPLANATION,
        preferred_agents=[AgentTool.GEMINI_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.LOW,
        context_requirement="large",
        notes="Gemini's context window helps for large files",
    ),
    TaskType.ARCHITECTURE_ANALYSIS: RoutingRule(
        task_type=TaskType.ARCHITECTURE_ANALYSIS,
        preferred_agents=[AgentTool.GEMINI_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.LOW,
        context_requirement="massive",
        notes="Gemini's 1M context ideal for full codebase analysis",
    ),
    TaskType.SECURITY_AUDIT: RoutingRule(
        task_type=TaskType.SECURITY_AUDIT,
        preferred_agents=[AgentTool.CLAUDE_CODE, AgentTool.CLAUDE_SDK],
        risk_level=RiskLevel.HIGH,
        context_requirement="large",
        notes="Security tasks need careful review",
    ),
    TaskType.PERFORMANCE_ANALYSIS: RoutingRule(
        task_type=TaskType.PERFORMANCE_ANALYSIS,
        preferred_agents=[AgentTool.GEMINI_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.LOW,
        context_requirement="large",
        notes="Read-only analysis",
    ),
    # Testing - Codex good for test generation
    TaskType.TEST_GENERATION: RoutingRule(
        task_type=TaskType.TEST_GENERATION,
        preferred_agents=[AgentTool.CODEX_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.LOW,
        context_requirement="normal",
        autonomy_recommendation="auto-edit",
        notes="Test files are low risk to generate",
    ),
    TaskType.TEST_EXECUTION: RoutingRule(
        task_type=TaskType.TEST_EXECUTION,
        preferred_agents=[AgentTool.CLAUDE_CODE, AgentTool.CODEX_CLI],
        risk_level=RiskLevel.LOW,
        context_requirement="minimal",
        autonomy_recommendation="full-auto",
        notes="Running tests is safe",
    ),
    TaskType.TEST_DEBUGGING: RoutingRule(
        task_type=TaskType.TEST_DEBUGGING,
        preferred_agents=[AgentTool.CLAUDE_CODE, AgentTool.CODEX_CLI],
        risk_level=RiskLevel.MEDIUM,
        context_requirement="normal",
        autonomy_recommendation="suggest",
        notes="May need to modify test or code",
    ),
    # Documentation - Gemini with search grounding
    TaskType.DOCUMENTATION: RoutingRule(
        task_type=TaskType.DOCUMENTATION,
        preferred_agents=[AgentTool.GEMINI_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.LOW,
        context_requirement="large",
        notes="Gemini's search grounding helps with docs",
    ),
    TaskType.API_DOCUMENTATION: RoutingRule(
        task_type=TaskType.API_DOCUMENTATION,
        preferred_agents=[AgentTool.GEMINI_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.LOW,
        context_requirement="large",
        notes="Needs to understand full API surface",
    ),
    # Infrastructure - Codex with appropriate autonomy
    TaskType.CI_CD_TASKS: RoutingRule(
        task_type=TaskType.CI_CD_TASKS,
        preferred_agents=[AgentTool.CODEX_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.HIGH,
        context_requirement="normal",
        autonomy_recommendation="suggest",
        notes="CI/CD changes need review",
    ),
    TaskType.DEPLOYMENT: RoutingRule(
        task_type=TaskType.DEPLOYMENT,
        preferred_agents=[AgentTool.CODEX_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.CRITICAL,
        context_requirement="normal",
        autonomy_recommendation="suggest",
        notes="Deployment is critical - always review",
    ),
    TaskType.DEVOPS: RoutingRule(
        task_type=TaskType.DEVOPS,
        preferred_agents=[AgentTool.CODEX_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.HIGH,
        context_requirement="normal",
        autonomy_recommendation="suggest",
        notes="Infrastructure changes need review",
    ),
    # Research - Gemini with search
    TaskType.RESEARCH: RoutingRule(
        task_type=TaskType.RESEARCH,
        preferred_agents=[AgentTool.GEMINI_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.LOW,
        context_requirement="normal",
        notes="Gemini's search grounding is ideal",
    ),
    TaskType.CODEBASE_EXPLORATION: RoutingRule(
        task_type=TaskType.CODEBASE_EXPLORATION,
        preferred_agents=[AgentTool.GEMINI_CLI, AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.LOW,
        context_requirement="massive",
        notes="Gemini's 1M context for large codebases",
    ),
    # Debugging
    TaskType.BUG_INVESTIGATION: RoutingRule(
        task_type=TaskType.BUG_INVESTIGATION,
        preferred_agents=[AgentTool.CLAUDE_CODE, AgentTool.GEMINI_CLI],
        risk_level=RiskLevel.LOW,
        context_requirement="large",
        notes="Investigation is read-only",
    ),
    TaskType.BUG_FIX: RoutingRule(
        task_type=TaskType.BUG_FIX,
        preferred_agents=[AgentTool.CLAUDE_CODE, AgentTool.CODEX_CLI],
        risk_level=RiskLevel.MEDIUM,
        context_requirement="normal",
        autonomy_recommendation="auto-edit",
        notes="Bug fixes need testing after",
    ),
    # Multi-file Operations
    TaskType.LARGE_REFACTOR: RoutingRule(
        task_type=TaskType.LARGE_REFACTOR,
        preferred_agents=[AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.HIGH,
        context_requirement="large",
        autonomy_recommendation="suggest",
        notes="Large refactors need careful review",
    ),
    TaskType.MIGRATION: RoutingRule(
        task_type=TaskType.MIGRATION,
        preferred_agents=[AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.HIGH,
        context_requirement="large",
        autonomy_recommendation="suggest",
        notes="Migrations are risky - always review",
    ),
    # Interactive
    TaskType.PAIR_PROGRAMMING: RoutingRule(
        task_type=TaskType.PAIR_PROGRAMMING,
        preferred_agents=[AgentTool.CLAUDE_CODE],
        risk_level=RiskLevel.MEDIUM,
        context_requirement="normal",
        notes="Claude Code is best for interactive sessions",
    ),
}


def get_routing_rule(task_type: TaskType) -> RoutingRule:
    """Get the routing rule for a task type."""
    return ROUTING_TABLE[task_type]


def get_preferred_agents(task_type: TaskType) -> list[AgentTool]:
    """Get the list of preferred agents for a task type."""
    return ROUTING_TABLE[task_type].preferred_agents


def get_agent_profile(tool: AgentTool) -> AgentCapability:
    """Get the capability profile for an agent tool."""
    return AGENT_PROFILES[tool]


def classify_task(task_description: str) -> TaskType:
    """
    Classify a task description into a TaskType.

    This is a simple keyword-based classifier.
    In production, consider using an LLM for better classification.
    """
    task_lower = task_description.lower()

    # Check for bug fix pattern: "fix" + "bug" or "resolve" + "issue" anywhere in text
    if ("fix" in task_lower and "bug" in task_lower) or \
       ("resolve" in task_lower and "issue" in task_lower):
        return TaskType.BUG_FIX

    # Check for security-related patterns
    if any(term in task_lower for term in ["vulnerability", "vulnerabilities", "security audit", "security check"]):
        return TaskType.SECURITY_AUDIT

    # Keywords to task type mapping
    keywords = {
        TaskType.TEST_GENERATION: ["write test", "create test", "add test", "unit test", "integration test"],
        TaskType.TEST_EXECUTION: ["run test", "execute test", "pytest", "jest", "npm test"],
        TaskType.TEST_DEBUGGING: ["debug test", "fix test", "failing test"],
        TaskType.CODE_REVIEW: ["review", "code review", "pr review", "check code"],
        TaskType.CODE_EXPLANATION: ["explain", "what does", "how does", "understand"],
        TaskType.DOCUMENTATION: ["document", "readme", "docstring", "jsdoc", "write docs"],
        TaskType.API_DOCUMENTATION: ["api doc", "swagger", "openapi", "api reference"],
        TaskType.BUG_FIX: ["fix bug", "fix issue", "resolve bug", "resolve issue", "patch"],
        TaskType.BUG_INVESTIGATION: ["investigate", "debug", "trace", "find cause"],
        TaskType.CODE_REFACTORING: ["refactor", "clean up", "restructure", "improve code"],
        TaskType.LARGE_REFACTOR: ["major refactor", "rewrite", "overhaul"],
        TaskType.MIGRATION: ["migrate", "migration", "upgrade", "convert"],
        TaskType.SECURITY_AUDIT: ["security", "vulnerability", "audit", "penetration"],
        TaskType.PERFORMANCE_ANALYSIS: ["performance", "optimize", "profil", "benchmark"],
        TaskType.CI_CD_TASKS: ["ci/cd", "pipeline", "github action", "gitlab ci", "jenkins"],
        TaskType.DEPLOYMENT: ["deploy", "release", "publish", "ship"],
        TaskType.DEVOPS: ["devops", "infrastructure", "terraform", "kubernetes", "docker"],
        TaskType.RESEARCH: ["research", "best practice", "compare", "evaluate"],
        TaskType.CODEBASE_EXPLORATION: ["explore", "find", "search", "locate", "where is"],
        TaskType.ARCHITECTURE_ANALYSIS: ["architecture", "design", "structure", "system design"],
        TaskType.CODE_COMPLETION: ["complete", "finish", "continue"],
        TaskType.PAIR_PROGRAMMING: ["pair", "collaborate", "work together", "help me with"],
    }

    for task_type, kw_list in keywords.items():
        for kw in kw_list:
            if kw in task_lower:
                return task_type

    # Default to code generation for unclassified tasks
    return TaskType.CODE_GENERATION
