"""
Integration tests for the Agent Control Loop.

Tests cover:
- Task lifecycle (assign, execute, complete)
- Multi-agent coordination
- Race condition prevention
- Escalation flows
- Health check integration
"""

import pytest
import asyncio
from datetime import datetime, timedelta
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import AsyncMock, MagicMock, patch

from agent_orchestrator.control.loop import (
    AgentControlLoop,
    AgentContext,
    LoopConfig,
    LoopState,
)
from agent_orchestrator.control.health import (
    StuckDetectionConfig,
    HealthCheckResult,
    AgentState,
    StuckReason,
)
from agent_orchestrator.control.actions import (
    ActionPolicy,
    ControlAction,
    ControlActionType,
    EscalationLevel,
)
from agent_orchestrator.persistence.database import OrchestratorDB
from agent_orchestrator.persistence.models import Task, Agent
from agent_orchestrator.adapters.base import (
    BaseAdapter,
    AgentResponse,
    AgentStatus,
    UsageStats,
    StatusPacket,
)


class MockAdapter(BaseAdapter):
    """Mock adapter for testing."""

    def __init__(self, name: str = "mock", delay: float = 0.0):
        super().__init__(agent_id=name)
        self.name = name
        self.delay = delay
        self.executed_prompts: list[str] = []
        self.execute_count = 0
        self._success = True
        self._error_message = None
        self._usage = UsageStats()

    def get_name(self) -> str:
        return self.name

    def is_healthy(self) -> bool:
        return self._status != AgentStatus.ERROR

    async def execute(self, prompt: str, context: dict = None) -> AgentResponse:
        self.execute_count += 1
        self.executed_prompts.append(prompt)

        if self.delay > 0:
            await asyncio.sleep(self.delay)

        return AgentResponse(
            content=f"Response to: {prompt}",
            success=self._success,
            error=self._error_message,
            tokens_input=100,
            tokens_output=50,
            cost=0.01,
        )

    async def stream(self, task: str, context: dict = None):
        """Stream implementation (not used in tests)."""
        yield f"Response to: {task}"

    def get_usage(self) -> UsageStats:
        """Return usage stats."""
        return self._usage

    def write_status_packet(self) -> StatusPacket:
        """Write status packet."""
        return StatusPacket(
            agent_id=self.name,
            tokens_used=self._usage.total_tokens(),
            cost_usd=self._usage.total_cost,
            tool_calls=[],
            files_modified=[],
            errors=[],
        )

    def set_error(self, error_message: str):
        """Configure adapter to return errors."""
        self._success = False
        self._error_message = error_message

    def set_success(self):
        """Configure adapter to return success."""
        self._success = True
        self._error_message = None


class TestTaskLifecycle:
    """Integration tests for task lifecycle components."""

    @pytest.fixture
    def temp_dir(self):
        """Create temporary directory for tests."""
        with TemporaryDirectory() as tmpdir:
            yield Path(tmpdir)

    @pytest.fixture
    def db(self, temp_dir):
        """Create test database."""
        return OrchestratorDB(temp_dir / "test.db")

    @pytest.fixture
    def control_loop(self, db, temp_dir):
        """Create control loop with test configuration."""
        ops_path = temp_dir / "ops"
        ops_path.mkdir(exist_ok=True)
        (ops_path / "context").mkdir(exist_ok=True)
        (ops_path / "runbooks").mkdir(exist_ok=True)

        loop = AgentControlLoop(
            db=db,
            ops_path=ops_path,
            config=LoopConfig(
                health_check_interval_seconds=1,
                task_poll_interval_seconds=1,
            ),
            stuck_config=StuckDetectionConfig(
                idle_time_threshold_seconds=60,
                repeated_error_threshold=3,
            ),
            action_policy=ActionPolicy(
                auto_prompt_max_attempts=2,
                escalate_after_auto_prompts=1,
                auto_prompt_base_delay_seconds=0.0,
            ),
        )
        return loop

    @pytest.mark.asyncio
    async def test_task_assignment_registers_context(self, control_loop, db):
        """Test that task assignment creates agent context."""
        adapter = MockAdapter("claude")
        control_loop.register_adapter("claude", adapter)

        agent = Agent(id="claude", tool="claude_code", status="idle")
        db.create_agent(agent)

        task = Task(
            id="task-001",
            description="Test task",
            task_type="test",
            status="pending",
        )
        db.create_task(task)

        # Assign task
        result = await control_loop.assign_task("task-001", "claude")

        assert result is True

        # Verify agent context was created
        assert "claude" in control_loop._agents
        context = control_loop._agents["claude"]
        assert context.current_task_id == "task-001"
        assert context.generation == 1

    @pytest.mark.asyncio
    async def test_task_assignment_rejected_for_busy_agent(self, control_loop, db):
        """Test that a second task assignment is rejected for busy agent."""
        adapter = MockAdapter("claude", delay=1.0)  # Long-running task
        control_loop.register_adapter("claude", adapter)

        agent = Agent(id="claude", tool="claude_code", status="idle")
        db.create_agent(agent)

        task1 = Task(id="task-001", description="First task", task_type="test", status="pending")
        task2 = Task(id="task-002", description="Second task", task_type="test", status="pending")
        db.create_task(task1)
        db.create_task(task2)

        # Assign first task
        result1 = await control_loop.assign_task("task-001", "claude")
        assert result1 is True

        # Try to assign second task while first is running
        result2 = await control_loop.assign_task("task-002", "claude")
        assert result2 is False

    @pytest.mark.asyncio
    async def test_task_cancellation(self, control_loop, db):
        """Test task cancellation during execution."""
        adapter = MockAdapter("claude", delay=0.5)
        control_loop.register_adapter("claude", adapter)

        agent = Agent(id="claude", tool="claude_code", status="idle")
        db.create_agent(agent)

        task = Task(id="task-003", description="Long task", task_type="test", status="pending")
        db.create_task(task)

        await control_loop.assign_task("task-003", "claude")

        # Cancel while running
        await asyncio.sleep(0.1)
        result = await control_loop.cancel_task("task-003")

        assert result is True

        # Verify task was cancelled
        updated_task = db.get_task("task-003")
        assert updated_task.status == "failed"
        assert "Cancelled" in (updated_task.error_message or "")


class TestAgentContextGeneration:
    """Tests for generation-based race condition prevention."""

    def test_generation_increment(self):
        """Test that generation is incremented correctly."""
        context = AgentContext(
            agent_id="test",
            adapter=MagicMock(),
            generation=1,
        )

        assert context.generation == 1

        new_gen = context.increment_generation()
        assert new_gen == 2
        assert context.generation == 2

    def test_stale_detection(self):
        """Test stale operation detection."""
        context = AgentContext(
            agent_id="test",
            adapter=MagicMock(),
            generation=1,
        )

        # Same generation is not stale
        assert context.is_stale(1) is False

        # Different generation is stale
        assert context.is_stale(0) is True
        assert context.is_stale(2) is True

    def test_terminated_context_is_stale(self):
        """Test that terminated context is always stale."""
        context = AgentContext(
            agent_id="test",
            adapter=MagicMock(),
            generation=1,
        )

        assert context.is_stale(1) is False

        context.is_terminated = True
        assert context.is_stale(1) is True


class TestMultiAgentCoordination:
    """Integration tests for multi-agent scenarios."""

    @pytest.fixture
    def temp_dir(self):
        with TemporaryDirectory() as tmpdir:
            yield Path(tmpdir)

    @pytest.fixture
    def db(self, temp_dir):
        return OrchestratorDB(temp_dir / "test.db")

    @pytest.fixture
    def control_loop(self, db, temp_dir):
        ops_path = temp_dir / "ops"
        ops_path.mkdir(exist_ok=True)
        (ops_path / "context").mkdir(exist_ok=True)
        (ops_path / "runbooks").mkdir(exist_ok=True)

        loop = AgentControlLoop(
            db=db,
            ops_path=ops_path,
            config=LoopConfig(
                health_check_interval_seconds=1,
                max_concurrent_agents=3,
            ),
            action_policy=ActionPolicy(auto_prompt_base_delay_seconds=0.0),
        )
        return loop

    @pytest.mark.asyncio
    async def test_multiple_agents_can_be_registered(self, control_loop, db):
        """Test that multiple agents can be registered and tracked."""
        # Register multiple adapters
        for name in ["claude", "gemini", "codex"]:
            adapter = MockAdapter(name, delay=0.1)
            control_loop.register_adapter(name, adapter)

            agent = Agent(id=name, tool="claude_code", status="idle")
            db.create_agent(agent)

        # Verify all registered
        status = control_loop.get_loop_status()
        assert status["registered_adapters"] == 3

    @pytest.mark.asyncio
    async def test_multiple_agents_can_have_concurrent_contexts(self, control_loop, db):
        """Test that multiple agents can have active contexts simultaneously."""
        # Register multiple adapters with delay so tasks stay active
        for name in ["claude", "gemini"]:
            adapter = MockAdapter(name, delay=1.0)
            control_loop.register_adapter(name, adapter)

            agent = Agent(id=name, tool="claude_code", status="idle")
            db.create_agent(agent)

        # Create tasks
        task1 = Task(id="task-1", description="Task 1", task_type="test", status="pending")
        task2 = Task(id="task-2", description="Task 2", task_type="test", status="pending")
        db.create_task(task1)
        db.create_task(task2)

        # Assign tasks to different agents
        result1 = await control_loop.assign_task("task-1", "claude")
        result2 = await control_loop.assign_task("task-2", "gemini")

        assert result1 is True
        assert result2 is True

        # Both agents should have active contexts
        assert "claude" in control_loop._agents
        assert "gemini" in control_loop._agents
        assert control_loop._agents["claude"].current_task_id == "task-1"
        assert control_loop._agents["gemini"].current_task_id == "task-2"

    @pytest.mark.asyncio
    async def test_agent_cannot_run_multiple_tasks(self, control_loop, db):
        """Test that one agent cannot run multiple tasks simultaneously."""
        adapter = MockAdapter("claude", delay=0.5)
        control_loop.register_adapter("claude", adapter)

        agent = Agent(id="claude", tool="claude_code", status="idle")
        db.create_agent(agent)

        # Create two tasks
        task1 = Task(id="task-1", description="First task", task_type="test", status="pending")
        task2 = Task(id="task-2", description="Second task", task_type="test", status="pending")
        db.create_task(task1)
        db.create_task(task2)

        # Assign first task
        result1 = await control_loop.assign_task("task-1", "claude")
        assert result1 is True

        # Try to assign second task while first is running
        await asyncio.sleep(0.1)
        result2 = await control_loop.assign_task("task-2", "claude")

        # Should be rejected
        assert result2 is False

    @pytest.mark.asyncio
    async def test_get_agent_status(self, control_loop, db):
        """Test agent status reporting."""
        adapter = MockAdapter("claude")
        control_loop.register_adapter("claude", adapter)

        agent = Agent(id="claude", tool="claude_code", status="idle")
        db.create_agent(agent)

        # Check status before task
        status = control_loop.get_agent_status("claude")
        assert status["is_registered"] is True
        assert status["is_active"] is False

        # Assign task
        task = Task(id="task-1", description="Test", task_type="test", status="pending")
        db.create_task(task)
        await control_loop.assign_task("task-1", "claude")

        # Check status with active task
        status = control_loop.get_agent_status("claude")
        assert status["is_active"] is True
        assert status["current_task"] == "task-1"

    @pytest.mark.asyncio
    async def test_loop_status_reporting(self, control_loop, db):
        """Test overall loop status reporting."""
        # Register adapters
        for name in ["claude", "gemini"]:
            adapter = MockAdapter(name)
            control_loop.register_adapter(name, adapter)

        status = control_loop.get_loop_status()

        assert status["state"] == "stopped"
        assert status["registered_adapters"] == 2
        assert status["active_agents"] == 0


class TestPendingActionDeduplication:
    """Tests for preventing duplicate actions."""

    @pytest.fixture
    def temp_dir(self):
        with TemporaryDirectory() as tmpdir:
            yield Path(tmpdir)

    @pytest.fixture
    def db(self, temp_dir):
        return OrchestratorDB(temp_dir / "test.db")

    @pytest.fixture
    def control_loop(self, db, temp_dir):
        ops_path = temp_dir / "ops"
        ops_path.mkdir(exist_ok=True)
        (ops_path / "context").mkdir(exist_ok=True)
        (ops_path / "runbooks").mkdir(exist_ok=True)

        loop = AgentControlLoop(
            db=db,
            ops_path=ops_path,
            action_policy=ActionPolicy(auto_prompt_base_delay_seconds=0.0),
        )
        return loop

    def test_pending_actions_tracking(self, control_loop):
        """Test that pending actions are tracked."""
        assert len(control_loop._pending_actions) == 0

        control_loop._pending_actions.add("agent-1")
        assert "agent-1" in control_loop._pending_actions

        control_loop._pending_actions.discard("agent-1")
        assert "agent-1" not in control_loop._pending_actions


class TestActionPolicyIntegration:
    """Integration tests for action policy with control loop."""

    def test_backoff_prevents_rapid_auto_prompts(self):
        """Test that backoff prevents flooding agents with prompts."""
        policy = ActionPolicy(
            auto_prompt_base_delay_seconds=30.0,
            auto_prompt_max_attempts=5,
        )

        health_result = HealthCheckResult(
            agent_id="agent-1",
            timestamp=datetime.now(),
            state=AgentState.RUNNING,
            is_healthy=False,
            is_stuck=True,
            stuck_reason=StuckReason.IDLE_BURNING_TOKENS,
            stuck_details="No progress",
        )

        # First auto-prompt allowed
        action1 = policy.decide_action(health_result)
        assert action1.action_type == ControlActionType.AUTO_PROMPT
        policy.record_auto_prompt("agent-1")

        # Second blocked by backoff
        action2 = policy.decide_action(health_result)
        assert action2.action_type == ControlActionType.CONTINUE
        assert "backoff" in action2.reason.lower()

    def test_escalation_cooldown_prevents_flooding(self):
        """Test that escalation cooldown prevents rapid escalations."""
        policy = ActionPolicy()
        policy._escalation_cooldown_seconds = 60.0

        # First escalation allowed
        assert policy.can_escalate("agent-1") is True
        policy.record_escalation("agent-1", EscalationLevel.WARN)

        # Second blocked by cooldown
        assert policy.can_escalate("agent-1") is False

        # Verify state tracking
        state = policy.get_escalation_state("agent-1")
        assert state["count"] == 1
        assert state["last_escalation_level"] == "warn"
