"""Integration tests for deep research resilience features.

Tests cover:
- Cancellation mid-workflow
- Timeout handling with partial results
- Crash recovery from persisted state
"""

from __future__ import annotations

import asyncio
import time
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch, AsyncMock

import pytest

from foundry_mcp.config.research import ResearchConfig
from foundry_mcp.core.background_task import BackgroundTask, TaskStatus
from foundry_mcp.core.research.models.deep_research import (
    DeepResearchPhase,
    DeepResearchState,
)
from foundry_mcp.core.research.models.sources import SubQuery
from foundry_mcp.core.research.workflows.deep_research import DeepResearchWorkflow


class TestCancellationIntegration:
    """Integration tests for cancellation mid-workflow."""

    @pytest.mark.asyncio
    async def test_cancel_sets_metadata_and_persists(self):
        """Cancelling a research task sets cancelled=True in metadata and persists state."""
        config = ResearchConfig()
        workflow = DeepResearchWorkflow(config)

        # Create a state that's in progress
        state = DeepResearchState(
            id="test-cancel-integration",
            original_query="test query",
            phase=DeepResearchPhase.GATHERING,
        )
        state.sub_queries = [
            SubQuery(id="sq-1", query="sub query 1", status="pending"),
        ]

        # Create a background task and mock it as running (not done)
        bg_task = BackgroundTask(research_id=state.id)

        # Mock cancel to return True (simulates task was running)
        with patch.object(bg_task, "cancel", return_value=True):
            with patch.object(workflow, "get_background_task", return_value=bg_task):
                with patch.object(workflow.memory, "load_deep_research", return_value=state):
                    with patch.object(workflow.memory, "save_deep_research"):
                        # Execute cancel
                        result = workflow._cancel_research(state.id)

        assert result.success is True
        assert "cancelled" in result.metadata
        assert result.metadata.get("research_id") == state.id

    @pytest.mark.asyncio
    async def test_cancel_returns_partial_results(self):
        """Cancellation returns any partial results accumulated so far."""
        config = ResearchConfig()
        workflow = DeepResearchWorkflow(config)

        # Create a state with some completed work
        state = DeepResearchState(
            id="test-cancel-partial",
            original_query="test query",
            phase=DeepResearchPhase.ANALYSIS,
        )
        state.sub_queries = [
            SubQuery(id="sq-1", query="sub query 1", status="completed"),
            SubQuery(id="sq-2", query="sub query 2", status="pending"),
        ]

        bg_task = BackgroundTask(research_id=state.id)

        # Mock cancel to return True (task was running and is now cancelled)
        with patch.object(bg_task, "cancel", return_value=True):
            with patch.object(workflow, "get_background_task", return_value=bg_task):
                with patch.object(workflow.memory, "load_deep_research", return_value=state):
                    with patch.object(workflow.memory, "save_deep_research"):
                        result = workflow._cancel_research(state.id)

        assert result.success is True
        # Should include cancelled flag in metadata
        assert result.metadata.get("cancelled") is True


class TestTimeoutIntegration:
    """Integration tests for timeout handling."""

    def test_timeout_marks_state_with_abort_phase(self):
        """Timeout should mark state with abort_phase and abort_iteration."""
        config = ResearchConfig()
        workflow = DeepResearchWorkflow(config)

        # Create a timed-out background task
        bg_task = BackgroundTask(research_id="test-timeout-abort", timeout=0.01)
        time.sleep(0.02)
        bg_task.mark_timeout()

        # Create state that was in GATHERING phase
        state = DeepResearchState(
            id="test-timeout-abort",
            original_query="test query",
            phase=DeepResearchPhase.GATHERING,
            iteration=2,
        )
        state.metadata["timeout"] = True
        state.metadata["abort_phase"] = "gathering"
        state.metadata["abort_iteration"] = 2

        with patch.object(workflow, "get_background_task", return_value=bg_task):
            with patch.object(workflow.memory, "load_deep_research", return_value=state):
                result = workflow._get_status("test-timeout-abort")

        assert result.success is True
        assert result.metadata.get("is_timed_out") is True

    def test_status_includes_timeout_metadata_from_state(self):
        """Status response includes timeout metadata from persisted state."""
        config = ResearchConfig()
        workflow = DeepResearchWorkflow(config)

        # State with timeout metadata
        state = DeepResearchState(
            id="test-timeout-meta",
            original_query="test query",
            phase=DeepResearchPhase.SYNTHESIS,
        )
        state.metadata["timeout"] = True
        state.completed_at = datetime.now(timezone.utc)

        # No background task (completed/persisted state)
        with patch.object(workflow, "get_background_task", return_value=None):
            with patch.object(workflow.memory, "load_deep_research", return_value=state):
                with patch.object(workflow.memory, "save_deep_research"):
                    result = workflow._get_status("test-timeout-meta")

        assert result.success is True
        assert result.metadata.get("timed_out") is True


class TestCrashRecoveryIntegration:
    """Integration tests for crash recovery from persisted state."""

    def test_continue_loads_persisted_state(self):
        """Continue action loads state from persistence."""
        config = ResearchConfig()
        workflow = DeepResearchWorkflow(config)

        # Create persisted state from "previous session"
        state = DeepResearchState(
            id="test-recovery-state",
            original_query="test query",
            phase=DeepResearchPhase.GATHERING,
            iteration=1,
        )
        state.sub_queries = [
            SubQuery(id="sq-1", query="sub query 1", status="completed"),
            SubQuery(id="sq-2", query="sub query 2", status="pending"),
        ]

        with patch.object(workflow.memory, "load_deep_research", return_value=state):
            # Verify state can be loaded
            loaded = workflow.memory.load_deep_research("test-recovery-state")
            assert loaded is not None
            assert loaded.id == "test-recovery-state"
            assert loaded.phase == DeepResearchPhase.GATHERING
            assert len(loaded.sub_queries) == 2

    def test_status_returns_partial_progress_after_crash(self):
        """Status shows partial progress from persisted state after crash."""
        config = ResearchConfig()
        workflow = DeepResearchWorkflow(config)

        # State representing crash mid-workflow
        state = DeepResearchState(
            id="test-crash-progress",
            original_query="test query",
            phase=DeepResearchPhase.ANALYSIS,
        )
        state.sub_queries = [
            SubQuery(id="sq-1", query="q1", status="completed"),
            SubQuery(id="sq-2", query="q2", status="completed"),
            SubQuery(id="sq-3", query="q3", status="failed", error="crash"),
        ]
        state.metadata["failed"] = True
        state.metadata["failure_error"] = "Unexpected error during analysis"

        # No background task (crashed)
        with patch.object(workflow, "get_background_task", return_value=None):
            with patch.object(workflow.memory, "load_deep_research", return_value=state):
                with patch.object(workflow.memory, "save_deep_research"):
                    result = workflow._get_status("test-crash-progress")

        assert result.success is True
        assert result.metadata.get("is_failed") is True
        assert result.metadata.get("sub_queries_completed") == 2
        assert "failure_error" in result.metadata


class TestHeartbeatVisibility:
    """Integration tests for heartbeat visibility during execution."""

    def test_status_shows_last_heartbeat_during_execution(self):
        """Status response includes last_heartbeat_at for progress visibility."""
        config = ResearchConfig()
        workflow = DeepResearchWorkflow(config)

        # Active task with recent heartbeat
        bg_task = BackgroundTask(research_id="test-heartbeat-visible")

        state = DeepResearchState(
            id="test-heartbeat-visible",
            original_query="test query",
            phase=DeepResearchPhase.GATHERING,
        )
        heartbeat = datetime(2026, 1, 26, 12, 0, 0, tzinfo=timezone.utc)
        state.last_heartbeat_at = heartbeat

        with patch.object(workflow, "get_background_task", return_value=bg_task):
            with patch.object(workflow.memory, "load_deep_research", return_value=state):
                result = workflow._get_status("test-heartbeat-visible")

        assert result.success is True
        assert result.metadata.get("last_heartbeat_at") == heartbeat.isoformat()
        assert result.metadata.get("phase") == "gathering"
