"""
Integration Tests for Task Routing.

These tests verify that the task router:
1. Correctly classifies tasks
2. Routes to appropriate agents
3. Respects budget constraints
4. Handles approval requirements
"""

import pytest
from unittest.mock import MagicMock, AsyncMock, patch
from pathlib import Path

from agent_orchestrator.routing.task_types import (
    TaskType,
    AgentTool,
    RoutingRule,
    AgentCapability,
    ROUTING_TABLE,
    AGENT_PROFILES,
    classify_task,
    get_routing_rule,
    get_preferred_agents,
    get_agent_profile,
)
from agent_orchestrator.routing.router import (
    TaskRouter,
    RouteDecision,
    BudgetManager,
    create_router,
)
from agent_orchestrator.adapters.base import RiskLevel, AgentStatus
from agent_orchestrator.persistence.database import OrchestratorDB
from agent_orchestrator.config import Config, BudgetsConfig


class TestTaskTypeClassification:
    """Tests for task classification."""

    def test_classify_test_generation(self):
        """Test classification of test-related tasks."""
        assert classify_task("Write unit tests for auth module") == TaskType.TEST_GENERATION
        assert classify_task("Create integration tests") == TaskType.TEST_GENERATION
        assert classify_task("Add tests for the API") == TaskType.TEST_GENERATION

    def test_classify_test_execution(self):
        """Test classification of test execution tasks."""
        assert classify_task("Run tests and report results") == TaskType.TEST_EXECUTION
        assert classify_task("Execute pytest") == TaskType.TEST_EXECUTION
        assert classify_task("Run npm test") == TaskType.TEST_EXECUTION

    def test_classify_code_review(self):
        """Test classification of code review tasks."""
        assert classify_task("Review this pull request") == TaskType.CODE_REVIEW
        assert classify_task("Do a code review of the changes") == TaskType.CODE_REVIEW

    def test_classify_documentation(self):
        """Test classification of documentation tasks."""
        assert classify_task("Document this module") == TaskType.DOCUMENTATION
        assert classify_task("Write README for the project") == TaskType.DOCUMENTATION
        assert classify_task("Add docstrings to functions") == TaskType.DOCUMENTATION

    def test_classify_bug_fix(self):
        """Test classification of bug fix tasks."""
        assert classify_task("Fix the authentication bug") == TaskType.BUG_FIX
        assert classify_task("Resolve the issue with login") == TaskType.BUG_FIX

    def test_classify_refactoring(self):
        """Test classification of refactoring tasks."""
        assert classify_task("Refactor the database module") == TaskType.CODE_REFACTORING
        assert classify_task("Clean up the legacy code") == TaskType.CODE_REFACTORING

    def test_classify_research(self):
        """Test classification of research tasks."""
        assert classify_task("Research best practices for caching") == TaskType.RESEARCH
        assert classify_task("Compare different auth libraries") == TaskType.RESEARCH

    def test_classify_security(self):
        """Test classification of security tasks."""
        assert classify_task("Security audit of the API") == TaskType.SECURITY_AUDIT
        assert classify_task("Check for vulnerabilities") == TaskType.SECURITY_AUDIT

    def test_classify_cicd(self):
        """Test classification of CI/CD tasks."""
        assert classify_task("Set up GitHub Actions pipeline") == TaskType.CI_CD_TASKS
        assert classify_task("Configure GitLab CI") == TaskType.CI_CD_TASKS

    def test_classify_deployment(self):
        """Test classification of deployment tasks."""
        assert classify_task("Deploy to production") == TaskType.DEPLOYMENT
        assert classify_task("Release version 2.0") == TaskType.DEPLOYMENT

    def test_classify_default(self):
        """Test default classification for unrecognized tasks."""
        # Unrecognized tasks should default to CODE_GENERATION
        assert classify_task("Some random task") == TaskType.CODE_GENERATION
        assert classify_task("Build the thing") == TaskType.CODE_GENERATION


class TestRoutingTable:
    """Tests for the routing table configuration."""

    def test_all_task_types_have_rules(self):
        """Ensure all task types have routing rules."""
        for task_type in TaskType:
            assert task_type in ROUTING_TABLE, f"Missing rule for {task_type}"

    def test_routing_rules_have_valid_agents(self):
        """Ensure routing rules reference valid agents."""
        for task_type, rule in ROUTING_TABLE.items():
            for agent in rule.preferred_agents:
                assert agent in AgentTool, f"Invalid agent {agent} in {task_type}"

    def test_routing_rules_have_valid_risk_levels(self):
        """Ensure routing rules have valid risk levels."""
        for task_type, rule in ROUTING_TABLE.items():
            assert isinstance(rule.risk_level, RiskLevel)

    def test_get_routing_rule(self):
        """Test getting routing rule for a task type."""
        rule = get_routing_rule(TaskType.TEST_GENERATION)

        assert rule.task_type == TaskType.TEST_GENERATION
        assert len(rule.preferred_agents) > 0
        assert rule.risk_level == RiskLevel.LOW

    def test_get_preferred_agents(self):
        """Test getting preferred agents for a task type."""
        agents = get_preferred_agents(TaskType.ARCHITECTURE_ANALYSIS)

        assert AgentTool.GEMINI_CLI in agents  # Gemini excels at large context
        assert len(agents) > 0


class TestAgentProfiles:
    """Tests for agent capability profiles."""

    def test_all_agents_have_profiles(self):
        """Ensure all agent tools have profiles."""
        for tool in AgentTool:
            assert tool in AGENT_PROFILES, f"Missing profile for {tool}"

    def test_profile_structure(self):
        """Test that profiles have required fields."""
        for tool, profile in AGENT_PROFILES.items():
            assert isinstance(profile, AgentCapability)
            assert profile.tool == tool
            assert isinstance(profile.priority, int)
            assert isinstance(profile.risk_level, RiskLevel)
            assert isinstance(profile.max_context_tokens, int)
            assert isinstance(profile.strengths, list)
            assert isinstance(profile.limitations, list)

    def test_get_agent_profile(self):
        """Test getting agent profile."""
        profile = get_agent_profile(AgentTool.GEMINI_CLI)

        assert profile.tool == AgentTool.GEMINI_CLI
        assert profile.max_context_tokens == 1000000  # 1M tokens
        assert "context" in " ".join(profile.strengths).lower()

    def test_claude_code_profile(self):
        """Test Claude Code profile configuration."""
        profile = AGENT_PROFILES[AgentTool.CLAUDE_CODE]

        assert profile.max_context_tokens == 200000
        assert profile.cost_tier == "high"
        assert "multi-file" in " ".join(profile.strengths).lower()


class TestBudgetManager:
    """Tests for budget management."""

    @pytest.fixture
    def config(self):
        """Create test configuration."""
        config = MagicMock(spec=Config)
        config.budgets = BudgetsConfig(
            daily_budget_usd=10.0,
            task_budget_usd=2.0,
            warn_threshold_percent=80,
        )
        return config

    @pytest.fixture
    def budget_manager(self, config):
        """Create budget manager for testing."""
        return BudgetManager(config)

    def test_can_afford_under_budget(self, budget_manager):
        """Test can_afford when under budget."""
        assert budget_manager.can_afford(AgentTool.CLAUDE_CODE, 1.0)

    def test_can_afford_over_budget(self, budget_manager):
        """Test can_afford when over budget."""
        # Spend most of the budget
        budget_manager.daily_spent["claude-code"] = 9.5

        assert not budget_manager.can_afford(AgentTool.CLAUDE_CODE, 1.0)

    def test_record_spend(self, budget_manager):
        """Test recording spending."""
        budget_manager.record_spend(AgentTool.CLAUDE_CODE, "task-1", 1.5)

        assert budget_manager.daily_spent["claude-code"] == 1.5
        assert budget_manager.task_spent["task-1"] == 1.5

    def test_get_remaining_budget(self, budget_manager):
        """Test getting remaining budget."""
        budget_manager.daily_spent["claude-code"] = 3.0

        remaining = budget_manager.get_remaining_budget(AgentTool.CLAUDE_CODE)
        assert remaining == 7.0

    def test_estimate_cost(self, budget_manager):
        """Test cost estimation."""
        # High-tier agent, large context
        cost = budget_manager.estimate_cost(AgentTool.CLAUDE_CODE, "large")
        assert cost > 0

        # Low-tier agent, minimal context
        cost_low = budget_manager.estimate_cost(AgentTool.GEMINI_CLI, "minimal")
        assert cost_low < cost


class TestTaskRouter:
    """Tests for the task router."""

    @pytest.fixture
    def db(self, tmp_path):
        """Create test database."""
        db_path = tmp_path / "test.db"
        return OrchestratorDB(str(db_path))

    @pytest.fixture
    def mock_config(self):
        """Create mock configuration."""
        config = MagicMock(spec=Config)
        config.budgets = BudgetsConfig(
            daily_budget_usd=10.0,
            task_budget_usd=2.0,
        )
        config.control_loop = MagicMock()
        config.control_loop.auto_approve_low_risk = True
        # Add cli_auth mock
        config.cli_auth = MagicMock()
        config.cli_auth.check_enabled = False
        # Add routing mock
        config.routing = MagicMock()
        config.routing.cli_rebalance_enabled = False
        return config

    @pytest.fixture
    def mock_adapter(self):
        """Create mock adapter."""
        adapter = MagicMock()
        adapter.agent_id = "mock-agent"
        adapter.is_healthy.return_value = True
        adapter._status = AgentStatus.IDLE
        adapter.execute = AsyncMock(return_value=MagicMock(cost=0.5))
        return adapter

    @pytest.fixture
    def router(self, db, mock_config):
        """Create router for testing."""
        return TaskRouter(db=db, config=mock_config)

    @pytest.mark.asyncio
    async def test_route_with_registered_adapter(self, router, mock_adapter):
        """Test routing when adapter is registered."""
        router.register_adapter(AgentTool.CLAUDE_CODE, mock_adapter)

        decision = await router.route("Write unit tests")

        assert decision.task_type == TaskType.TEST_GENERATION
        assert decision.adapter is not None
        assert decision.approved  # Low risk, auto-approved

    @pytest.mark.asyncio
    async def test_route_without_registered_adapter(self, router):
        """Test routing when no adapter is registered."""
        decision = await router.route("Write unit tests")

        assert not decision.approved
        assert "No suitable agent" in decision.rejection_reason

    @pytest.mark.asyncio
    async def test_route_with_force_agent(self, router, mock_adapter):
        """Test forcing specific agent."""
        router.register_adapter(AgentTool.GEMINI_CLI, mock_adapter)

        decision = await router.route(
            "Write unit tests",
            force_agent=AgentTool.GEMINI_CLI,
        )

        assert decision.selected_agent == AgentTool.GEMINI_CLI
        assert "Forced to" in decision.routing_notes

    @pytest.mark.asyncio
    async def test_route_respects_budget(self, router, mock_adapter, mock_config):
        """Test routing respects budget constraints."""
        router.register_adapter(AgentTool.CLAUDE_CODE, mock_adapter)

        # Exhaust budget
        router.budget_manager.daily_spent["claude-code"] = 10.0

        decision = await router.route("Write unit tests")

        # Should fail due to budget
        assert not decision.approved

    @pytest.mark.asyncio
    async def test_route_high_risk_requires_approval(self, router, mock_adapter, mock_config):
        """Test high-risk tasks require approval."""
        router.register_adapter(AgentTool.CLAUDE_CODE, mock_adapter)
        mock_config.control_loop.auto_approve_low_risk = False

        decision = await router.route("Deploy to production")

        assert decision.task_type == TaskType.DEPLOYMENT
        assert decision.risk_level == RiskLevel.CRITICAL
        assert decision.requires_approval
        assert not decision.approved

    @pytest.mark.asyncio
    async def test_execute_routed_task(self, router, mock_adapter):
        """Test executing a routed task."""
        router.register_adapter(AgentTool.CLAUDE_CODE, mock_adapter)

        decision = await router.route("Write unit tests")
        result = await router.execute_routed_task(decision, {})

        assert result is not None
        mock_adapter.execute.assert_called_once()

    @pytest.mark.asyncio
    async def test_execute_unapproved_task(self, router, mock_adapter, mock_config):
        """Test executing unapproved task returns None."""
        router.register_adapter(AgentTool.CLAUDE_CODE, mock_adapter)
        mock_config.control_loop.auto_approve_low_risk = False

        decision = await router.route("Deploy to production")
        assert not decision.approved

        result = await router.execute_routed_task(decision, {})
        assert result is None

    def test_get_available_agents(self, router, mock_adapter):
        """Test getting list of available agents."""
        router.register_adapter(AgentTool.CLAUDE_CODE, mock_adapter)
        router.register_adapter(AgentTool.GEMINI_CLI, mock_adapter)

        available = router.get_available_agents()

        assert AgentTool.CLAUDE_CODE in available
        assert AgentTool.GEMINI_CLI in available

    def test_get_agent_status(self, router, mock_adapter):
        """Test getting agent status."""
        router.register_adapter(AgentTool.CLAUDE_CODE, mock_adapter)

        status = router.get_agent_status()

        assert "claude-code" in status
        assert status["claude-code"]["healthy"] is True


class TestRouteDecision:
    """Tests for RouteDecision dataclass."""

    def test_to_dict(self):
        """Test converting decision to dictionary."""
        decision = RouteDecision(
            task="Write tests",
            task_type=TaskType.TEST_GENERATION,
            selected_agent=AgentTool.CLAUDE_CODE,
            adapter=None,
            risk_level=RiskLevel.LOW,
            autonomy_mode="auto-edit",
            approved=True,
            requires_approval=False,
            estimated_cost=0.05,
            alternatives=[AgentTool.CODEX_CLI],
        )

        result = decision.to_dict()

        assert result["task"] == "Write tests"
        assert result["task_type"] == "TEST_GENERATION"
        assert result["selected_agent"] == "claude-code"
        assert result["risk_level"] == "low"
        assert result["approved"] is True
        assert "codex-cli" in result["alternatives"]


class TestRoutingStrategies:
    """Tests for specific routing strategies based on task types."""

    def test_large_context_tasks_prefer_gemini(self):
        """Test that large context tasks prefer Gemini."""
        rule = get_routing_rule(TaskType.ARCHITECTURE_ANALYSIS)

        # Gemini should be high priority due to 1M context
        assert AgentTool.GEMINI_CLI in rule.preferred_agents[:2]
        assert rule.context_requirement == "massive"

    def test_ci_tasks_prefer_codex(self):
        """Test that CI/CD tasks prefer Codex."""
        rule = get_routing_rule(TaskType.CI_CD_TASKS)

        assert AgentTool.CODEX_CLI in rule.preferred_agents[:2]
        assert rule.autonomy_recommendation == "suggest"  # Safety for CI

    def test_interactive_tasks_prefer_claude(self):
        """Test that interactive tasks prefer Claude Code."""
        rule = get_routing_rule(TaskType.PAIR_PROGRAMMING)

        assert rule.preferred_agents[0] == AgentTool.CLAUDE_CODE

    def test_test_generation_settings(self):
        """Test that test generation has appropriate settings."""
        rule = get_routing_rule(TaskType.TEST_GENERATION)

        assert rule.risk_level == RiskLevel.LOW  # Tests are low risk
        assert rule.autonomy_recommendation == "auto-edit"
