"""Comprehensive integration tests for workflow pause functionality.

This test suite covers complete workflow pause functionality including:
- Real database interactions using containerized PostgreSQL
- Real storage operations using the test storage backend
- Complete workflow: create -> pause -> resume -> delete
- Testing with actual FileService (not mocked)
- Database transactions and rollback behavior
- Actual file upload and retrieval through storage
- Workflow status transitions in the database
- Error handling with real database constraints
- Concurrent access scenarios
- Multi-tenant isolation
- Prune functionality
- File storage integration

These tests use TestContainers to spin up real services for integration testing,
providing more reliable and realistic test scenarios than mocks.
"""

import json
import uuid
from dataclasses import dataclass
from datetime import timedelta

import pytest
from sqlalchemy import delete, select
from sqlalchemy.orm import Session, selectinload, sessionmaker

from dify_graph.entities import WorkflowExecution
from dify_graph.enums import WorkflowExecutionStatus
from extensions.ext_storage import storage
from libs.datetime_utils import naive_utc_now
from models import Account
from models import WorkflowPause as WorkflowPauseModel
from models.account import Tenant, TenantAccountJoin, TenantAccountRole
from models.model import UploadFile
from models.workflow import Workflow, WorkflowRun
from repositories.sqlalchemy_api_workflow_run_repository import (
    DifyAPISQLAlchemyWorkflowRunRepository,
    _WorkflowRunError,
)


@dataclass
class PauseWorkflowSuccessCase:
    """Test case for successful pause workflow operations."""

    name: str
    initial_status: WorkflowExecutionStatus
    description: str = ""


@dataclass
class PauseWorkflowFailureCase:
    """Test case for pause workflow failure scenarios."""

    name: str
    initial_status: WorkflowExecutionStatus
    description: str = ""


@dataclass
class ResumeWorkflowSuccessCase:
    """Test case for successful resume workflow operations."""

    name: str
    initial_status: WorkflowExecutionStatus
    description: str = ""


@dataclass
class ResumeWorkflowFailureCase:
    """Test case for resume workflow failure scenarios."""

    name: str
    initial_status: WorkflowExecutionStatus
    pause_resumed: bool
    set_running_status: bool = False
    description: str = ""


@dataclass
class PrunePausesTestCase:
    """Test case for prune pauses operations."""

    name: str
    pause_age: timedelta
    resume_age: timedelta | None
    expected_pruned_count: int
    description: str = ""


def pause_workflow_failure_cases() -> list[PauseWorkflowFailureCase]:
    """Create test cases for pause workflow failure scenarios."""
    return [
        PauseWorkflowFailureCase(
            name="pause_completed_workflow",
            initial_status=WorkflowExecutionStatus.SUCCEEDED,
            description="Should fail to pause a completed workflow",
        ),
        PauseWorkflowFailureCase(
            name="pause_failed_workflow",
            initial_status=WorkflowExecutionStatus.FAILED,
            description="Should fail to pause a failed workflow",
        ),
    ]


def resume_workflow_success_cases() -> list[ResumeWorkflowSuccessCase]:
    """Create test cases for successful resume workflow operations."""
    return [
        ResumeWorkflowSuccessCase(
            name="resume_paused_workflow",
            initial_status=WorkflowExecutionStatus.PAUSED,
            description="Should successfully resume a paused workflow",
        ),
    ]


def resume_workflow_failure_cases() -> list[ResumeWorkflowFailureCase]:
    """Create test cases for resume workflow failure scenarios."""
    return [
        ResumeWorkflowFailureCase(
            name="resume_already_resumed_workflow",
            initial_status=WorkflowExecutionStatus.PAUSED,
            pause_resumed=True,
            description="Should fail to resume an already resumed workflow",
        ),
        ResumeWorkflowFailureCase(
            name="resume_running_workflow",
            initial_status=WorkflowExecutionStatus.RUNNING,
            pause_resumed=False,
            set_running_status=True,
            description="Should fail to resume a running workflow",
        ),
    ]


def prune_pauses_test_cases() -> list[PrunePausesTestCase]:
    """Create test cases for prune pauses operations."""
    return [
        PrunePausesTestCase(
            name="prune_old_active_pauses",
            pause_age=timedelta(days=7),
            resume_age=None,
            expected_pruned_count=1,
            description="Should prune old active pauses",
        ),
        PrunePausesTestCase(
            name="prune_old_resumed_pauses",
            pause_age=timedelta(hours=12),  # Created 12 hours ago (recent)
            resume_age=timedelta(days=7),
            expected_pruned_count=1,
            description="Should prune old resumed pauses",
        ),
        PrunePausesTestCase(
            name="keep_recent_active_pauses",
            pause_age=timedelta(hours=1),
            resume_age=None,
            expected_pruned_count=0,
            description="Should keep recent active pauses",
        ),
        PrunePausesTestCase(
            name="keep_recent_resumed_pauses",
            pause_age=timedelta(days=1),
            resume_age=timedelta(hours=1),
            expected_pruned_count=0,
            description="Should keep recent resumed pauses",
        ),
    ]


class TestWorkflowPauseIntegration:
    """Comprehensive integration tests for workflow pause functionality."""

    @pytest.fixture(autouse=True)
    def setup_test_data(self, db_session_with_containers):
        """Set up test data for each test method using TestContainers."""
        # Create test tenant and account

        tenant = Tenant(
            name="Test Tenant",
            status="normal",
        )
        db_session_with_containers.add(tenant)
        db_session_with_containers.commit()

        account = Account(
            email="test@example.com",
            name="Test User",
            interface_language="en-US",
            status="active",
        )
        db_session_with_containers.add(account)
        db_session_with_containers.commit()

        # Create tenant-account join
        tenant_join = TenantAccountJoin(
            tenant_id=tenant.id,
            account_id=account.id,
            role=TenantAccountRole.OWNER,
            current=True,
        )
        db_session_with_containers.add(tenant_join)
        db_session_with_containers.commit()

        # Set test data
        self.test_tenant_id = tenant.id
        self.test_user_id = account.id
        self.test_app_id = str(uuid.uuid4())
        self.test_workflow_id = str(uuid.uuid4())

        # Create test workflow
        self.test_workflow = Workflow(
            id=self.test_workflow_id,
            tenant_id=self.test_tenant_id,
            app_id=self.test_app_id,
            type="workflow",
            version="draft",
            graph='{"nodes": [], "edges": []}',
            features='{"file_upload": {"enabled": false}}',
            created_by=self.test_user_id,
            created_at=naive_utc_now(),
        )

        # Store session instance
        self.session = db_session_with_containers

        # Save test data to database
        self.session.add(self.test_workflow)
        self.session.commit()

        yield

        # Cleanup
        self._cleanup_test_data()

    def _cleanup_test_data(self):
        """Clean up test data after each test method."""
        # Clean up workflow pauses
        self.session.execute(delete(WorkflowPauseModel))
        # Clean up upload files
        self.session.execute(
            delete(UploadFile).where(
                UploadFile.tenant_id == self.test_tenant_id,
            )
        )
        # Clean up workflow runs
        self.session.execute(
            delete(WorkflowRun).where(
                WorkflowRun.tenant_id == self.test_tenant_id,
                WorkflowRun.app_id == self.test_app_id,
            )
        )
        # Clean up workflows
        self.session.execute(
            delete(Workflow).where(
                Workflow.tenant_id == self.test_tenant_id,
                Workflow.app_id == self.test_app_id,
            )
        )
        self.session.commit()

    def _create_test_workflow_run(
        self, status: WorkflowExecutionStatus = WorkflowExecutionStatus.RUNNING
    ) -> WorkflowRun:
        """Create a test workflow run with specified status."""
        workflow_run = WorkflowRun(
            id=str(uuid.uuid4()),
            tenant_id=self.test_tenant_id,
            app_id=self.test_app_id,
            workflow_id=self.test_workflow_id,
            type="workflow",
            triggered_from="debugging",
            version="draft",
            status=status,
            created_by=self.test_user_id,
            created_by_role="account",
            created_at=naive_utc_now(),
        )
        self.session.add(workflow_run)
        self.session.commit()
        return workflow_run

    def _create_test_state(self) -> str:
        """Create a test state string."""
        return json.dumps(
            {
                "node_id": "test-node",
                "node_type": "llm",
                "status": "paused",
                "data": {"key": "value"},
                "timestamp": naive_utc_now().isoformat(),
            }
        )

    def _get_workflow_run_repository(self):
        """Get workflow run repository instance for testing."""
        # Create session factory from the test session
        engine = self.session.get_bind()
        session_factory = sessionmaker(bind=engine, expire_on_commit=False)

        # Create a test-specific repository that implements the missing save method
        class TestWorkflowRunRepository(DifyAPISQLAlchemyWorkflowRunRepository):
            """Test-specific repository that implements the missing save method."""

            def save(self, execution: WorkflowExecution):
                """Implement the missing save method for testing."""
                # For testing purposes, we don't need to implement this method
                # as it's not used in the pause functionality tests
                pass

        # Create and return repository instance
        repository = TestWorkflowRunRepository(session_maker=session_factory)
        return repository

    # ==================== Complete Pause Workflow Tests ====================

    def test_complete_pause_resume_workflow(self):
        """Test complete workflow: create -> pause -> resume -> delete."""
        # Arrange
        workflow_run = self._create_test_workflow_run()
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        # Act - Create pause state
        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )

        # Assert - Pause state created
        assert pause_entity is not None
        assert pause_entity.id is not None
        assert pause_entity.workflow_execution_id == workflow_run.id
        assert list(pause_entity.get_pause_reasons()) == []
        # Convert both to strings for comparison
        retrieved_state = pause_entity.get_state()
        if isinstance(retrieved_state, bytes):
            retrieved_state = retrieved_state.decode()
        assert retrieved_state == test_state

        # Verify database state
        query = select(WorkflowPauseModel).where(WorkflowPauseModel.workflow_run_id == workflow_run.id)
        pause_model = self.session.scalars(query).first()
        assert pause_model is not None
        assert pause_model.resumed_at is None
        assert pause_model.id == pause_entity.id

        self.session.refresh(workflow_run)
        assert workflow_run.status == WorkflowExecutionStatus.PAUSED

        # Act - Get pause state
        retrieved_entity = repository.get_workflow_pause(workflow_run.id)

        # Assert - Pause state retrieved
        assert retrieved_entity is not None
        assert retrieved_entity.id == pause_entity.id
        retrieved_state = retrieved_entity.get_state()
        if isinstance(retrieved_state, bytes):
            retrieved_state = retrieved_state.decode()
        assert retrieved_state == test_state
        assert list(retrieved_entity.get_pause_reasons()) == []

        # Act - Resume workflow
        resumed_entity = repository.resume_workflow_pause(
            workflow_run_id=workflow_run.id,
            pause_entity=pause_entity,
        )

        # Assert - Workflow resumed
        assert resumed_entity is not None
        assert resumed_entity.id == pause_entity.id
        assert resumed_entity.resumed_at is not None

        # Verify database state
        self.session.refresh(workflow_run)
        assert workflow_run.status == WorkflowExecutionStatus.RUNNING
        self.session.refresh(pause_model)
        assert pause_model.resumed_at is not None

        # Act - Delete pause state
        repository.delete_workflow_pause(pause_entity)

        # Assert - Pause state deleted
        with Session(bind=self.session.get_bind()) as session:
            deleted_pause = session.get(WorkflowPauseModel, pause_entity.id)
            assert deleted_pause is None

    def test_pause_workflow_success(self):
        """Test successful pause workflow scenarios."""
        workflow_run = self._create_test_workflow_run(status=WorkflowExecutionStatus.RUNNING)
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )

        assert pause_entity is not None
        assert pause_entity.workflow_execution_id == workflow_run.id

        retrieved_state = pause_entity.get_state()
        if isinstance(retrieved_state, bytes):
            retrieved_state = retrieved_state.decode()
        assert retrieved_state == test_state

        self.session.refresh(workflow_run)
        assert workflow_run.status == WorkflowExecutionStatus.PAUSED
        pause_query = select(WorkflowPauseModel).where(WorkflowPauseModel.workflow_run_id == workflow_run.id)
        pause_model = self.session.scalars(pause_query).first()
        assert pause_model is not None
        assert pause_model.id == pause_entity.id
        assert pause_model.resumed_at is None

    @pytest.mark.parametrize("test_case", pause_workflow_failure_cases(), ids=lambda tc: tc.name)
    def test_pause_workflow_failure(self, test_case: PauseWorkflowFailureCase):
        """Test pause workflow failure scenarios."""
        workflow_run = self._create_test_workflow_run(status=test_case.initial_status)
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        with pytest.raises(_WorkflowRunError):
            repository.create_workflow_pause(
                workflow_run_id=workflow_run.id,
                state_owner_user_id=self.test_user_id,
                state=test_state,
                pause_reasons=[],
            )

    @pytest.mark.parametrize("test_case", resume_workflow_success_cases(), ids=lambda tc: tc.name)
    def test_resume_workflow_success(self, test_case: ResumeWorkflowSuccessCase):
        """Test successful resume workflow scenarios."""
        workflow_run = self._create_test_workflow_run(status=test_case.initial_status)
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        if workflow_run.status != WorkflowExecutionStatus.RUNNING:
            workflow_run.status = WorkflowExecutionStatus.RUNNING
            self.session.commit()

        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )

        self.session.refresh(workflow_run)
        assert workflow_run.status == WorkflowExecutionStatus.PAUSED

        resumed_entity = repository.resume_workflow_pause(
            workflow_run_id=workflow_run.id,
            pause_entity=pause_entity,
        )
        assert resumed_entity is not None
        assert resumed_entity.id == pause_entity.id
        assert resumed_entity.resumed_at is not None

        self.session.refresh(workflow_run)
        assert workflow_run.status == WorkflowExecutionStatus.RUNNING
        pause_query = select(WorkflowPauseModel).where(WorkflowPauseModel.workflow_run_id == workflow_run.id)
        pause_model = self.session.scalars(pause_query).first()
        assert pause_model is not None
        assert pause_model.id == pause_entity.id
        assert pause_model.resumed_at is not None

    def test_resume_running_workflow(self):
        """Test resume workflow failure scenarios."""
        workflow_run = self._create_test_workflow_run(status=WorkflowExecutionStatus.RUNNING)
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )

        self.session.refresh(workflow_run)
        workflow_run.status = WorkflowExecutionStatus.RUNNING
        self.session.add(workflow_run)
        self.session.commit()

        with pytest.raises(_WorkflowRunError):
            repository.resume_workflow_pause(
                workflow_run_id=workflow_run.id,
                pause_entity=pause_entity,
            )

    def test_resume_resumed_pause(self):
        """Test resume workflow failure scenarios."""
        workflow_run = self._create_test_workflow_run(status=WorkflowExecutionStatus.RUNNING)
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )
        pause_model = self.session.get(WorkflowPauseModel, pause_entity.id)
        pause_model.resumed_at = naive_utc_now()
        self.session.add(pause_model)
        self.session.commit()

        with pytest.raises(_WorkflowRunError):
            repository.resume_workflow_pause(
                workflow_run_id=workflow_run.id,
                pause_entity=pause_entity,
            )

    # ==================== Error Scenario Tests ====================

    def test_pause_nonexistent_workflow_run(self):
        """Test pausing a non-existent workflow run."""
        # Arrange
        nonexistent_id = str(uuid.uuid4())
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        # Act & Assert
        with pytest.raises(ValueError, match="WorkflowRun not found"):
            repository.create_workflow_pause(
                workflow_run_id=nonexistent_id,
                state_owner_user_id=self.test_user_id,
                state=test_state,
                pause_reasons=[],
            )

    def test_resume_nonexistent_workflow_run(self):
        """Test resuming a non-existent workflow run."""
        # Arrange
        workflow_run = self._create_test_workflow_run()
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )

        nonexistent_id = str(uuid.uuid4())

        # Act & Assert
        with pytest.raises(ValueError, match="WorkflowRun not found"):
            repository.resume_workflow_pause(
                workflow_run_id=nonexistent_id,
                pause_entity=pause_entity,
            )

    # ==================== Prune Functionality Tests ====================

    @pytest.mark.parametrize("test_case", prune_pauses_test_cases(), ids=lambda tc: tc.name)
    def test_prune_pauses_scenarios(self, test_case: PrunePausesTestCase):
        """Test various prune pauses scenarios."""
        now = naive_utc_now()

        # Create pause state
        workflow_run = self._create_test_workflow_run()
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )

        # Manually adjust timestamps for testing
        pause_model = self.session.get(WorkflowPauseModel, pause_entity.id)
        pause_model.created_at = now - test_case.pause_age

        if test_case.resume_age is not None:
            # Resume pause and adjust resume time
            repository.resume_workflow_pause(
                workflow_run_id=workflow_run.id,
                pause_entity=pause_entity,
            )
            # Need to refresh to get the updated model
            self.session.refresh(pause_model)
            # Manually set the resumed_at to an older time for testing
            pause_model.resumed_at = now - test_case.resume_age
            self.session.commit()  # Commit the resumed_at change
            # Refresh again to ensure the change is persisted
            self.session.refresh(pause_model)

        self.session.commit()

        # Act - Prune pauses
        expiration_time = now - timedelta(days=1, seconds=1)  # Expire pauses older than 1 day (plus 1 second)
        resumption_time = now - timedelta(
            days=7, seconds=1
        )  # Clean up pauses resumed more than 7 days ago (plus 1 second)

        # Debug: Check pause state before pruning
        self.session.refresh(pause_model)
        print(f"Pause created_at: {pause_model.created_at}")
        print(f"Pause resumed_at: {pause_model.resumed_at}")
        print(f"Expiration time: {expiration_time}")
        print(f"Resumption time: {resumption_time}")

        # Force commit to ensure timestamps are saved
        self.session.commit()

        # Determine if the pause should be pruned based on timestamps
        should_be_pruned = False
        if test_case.resume_age is not None:
            # If resumed, check if resumed_at is older than resumption_time
            should_be_pruned = pause_model.resumed_at < resumption_time
        else:
            # If not resumed, check if created_at is older than expiration_time
            should_be_pruned = pause_model.created_at < expiration_time

        # Act - Prune pauses
        pruned_ids = repository.prune_pauses(
            expiration=expiration_time,
            resumption_expiration=resumption_time,
        )

        # Assert - Check pruning results
        if should_be_pruned:
            assert len(pruned_ids) == test_case.expected_pruned_count
            # Verify pause was actually deleted
            # The pause should be in the pruned_ids list if it was pruned
            assert pause_entity.id in pruned_ids
        else:
            assert len(pruned_ids) == 0

    def test_prune_pauses_with_limit(self):
        """Test prune pauses with limit parameter."""
        now = naive_utc_now()

        # Create multiple pause states
        pause_entities = []
        repository = self._get_workflow_run_repository()

        for i in range(5):
            workflow_run = self._create_test_workflow_run()
            test_state = self._create_test_state()

            pause_entity = repository.create_workflow_pause(
                workflow_run_id=workflow_run.id,
                state_owner_user_id=self.test_user_id,
                state=test_state,
                pause_reasons=[],
            )
            pause_entities.append(pause_entity)

            # Make all pauses old enough to be pruned
            pause_model = self.session.get(WorkflowPauseModel, pause_entity.id)
            pause_model.created_at = now - timedelta(days=7)

        self.session.commit()

        # Act - Prune with limit
        expiration_time = now - timedelta(days=1)
        resumption_time = now - timedelta(days=7)

        pruned_ids = repository.prune_pauses(
            expiration=expiration_time,
            resumption_expiration=resumption_time,
            limit=3,
        )

        # Assert
        assert len(pruned_ids) == 3

        # Verify only 3 were deleted
        remaining_count = (
            self.session.query(WorkflowPauseModel)
            .filter(WorkflowPauseModel.id.in_([pe.id for pe in pause_entities]))
            .count()
        )
        assert remaining_count == 2

    # ==================== Multi-tenant Isolation Tests ====================

    def test_multi_tenant_pause_isolation(self):
        """Test that pause states are properly isolated by tenant."""
        # Arrange - Create second tenant

        tenant2 = Tenant(
            name="Test Tenant 2",
            status="normal",
        )
        self.session.add(tenant2)
        self.session.commit()

        account2 = Account(
            email="test2@example.com",
            name="Test User 2",
            interface_language="en-US",
            status="active",
        )
        self.session.add(account2)
        self.session.commit()

        tenant2_join = TenantAccountJoin(
            tenant_id=tenant2.id,
            account_id=account2.id,
            role=TenantAccountRole.OWNER,
            current=True,
        )
        self.session.add(tenant2_join)
        self.session.commit()

        # Create workflow for tenant 2
        workflow2 = Workflow(
            id=str(uuid.uuid4()),
            tenant_id=tenant2.id,
            app_id=str(uuid.uuid4()),
            type="workflow",
            version="draft",
            graph='{"nodes": [], "edges": []}',
            features='{"file_upload": {"enabled": false}}',
            created_by=account2.id,
            created_at=naive_utc_now(),
        )
        self.session.add(workflow2)
        self.session.commit()

        # Create workflow runs for both tenants
        workflow_run1 = self._create_test_workflow_run()
        workflow_run2 = WorkflowRun(
            id=str(uuid.uuid4()),
            tenant_id=tenant2.id,
            app_id=workflow2.app_id,
            workflow_id=workflow2.id,
            type="workflow",
            triggered_from="debugging",
            version="draft",
            status=WorkflowExecutionStatus.RUNNING,
            created_by=account2.id,
            created_by_role="account",
            created_at=naive_utc_now(),
        )
        self.session.add(workflow_run2)
        self.session.commit()

        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        # Act - Create pause for tenant 1
        pause_entity1 = repository.create_workflow_pause(
            workflow_run_id=workflow_run1.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )

        # Try to access pause from tenant 2 using tenant 1's repository
        # This should work because we're using the same repository
        pause_entity2 = repository.get_workflow_pause(workflow_run2.id)
        assert pause_entity2 is None  # No pause for tenant 2 yet

        # Create pause for tenant 2
        pause_entity2 = repository.create_workflow_pause(
            workflow_run_id=workflow_run2.id,
            state_owner_user_id=account2.id,
            state=test_state,
            pause_reasons=[],
        )

        # Assert - Both pauses should exist and be separate
        assert pause_entity1 is not None
        assert pause_entity2 is not None
        assert pause_entity1.id != pause_entity2.id
        assert pause_entity1.workflow_execution_id != pause_entity2.workflow_execution_id

    def test_cross_tenant_access_restriction(self):
        """Test that cross-tenant access is properly restricted."""
        # This test would require tenant-specific repositories
        # For now, we test that pause entities are properly scoped by tenant_id
        workflow_run = self._create_test_workflow_run()
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )

        # Verify pause is properly scoped
        pause_model = self.session.get(WorkflowPauseModel, pause_entity.id)
        assert pause_model.workflow_id == self.test_workflow_id

    # ==================== File Storage Integration Tests ====================

    def test_file_storage_integration(self):
        """Test that state files are properly stored and retrieved."""
        # Arrange
        workflow_run = self._create_test_workflow_run()
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        # Act - Create pause state
        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=test_state,
            pause_reasons=[],
        )

        # Assert - Verify file was uploaded to storage
        pause_model = self.session.get(WorkflowPauseModel, pause_entity.id)
        assert pause_model.state_object_key != ""

        # Verify file content in storage

        file_key = pause_model.state_object_key
        storage_content = storage.load(file_key).decode()
        assert storage_content == test_state

        # Verify retrieval through entity
        retrieved_state = pause_entity.get_state()
        if isinstance(retrieved_state, bytes):
            retrieved_state = retrieved_state.decode()
        assert retrieved_state == test_state

    def test_file_cleanup_on_pause_deletion(self):
        """Test that files are properly handled on pause deletion."""
        # Arrange
        workflow_run = self._create_test_workflow_run()
        test_state = self._create_test_state()
        repository = self._get_workflow_run_repository()

        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id, state_owner_user_id=self.test_user_id, state=test_state, pause_reasons=[]
        )

        # Get file info before deletion
        pause_model = self.session.get(WorkflowPauseModel, pause_entity.id)
        file_key = pause_model.state_object_key

        # Act - Delete pause state
        repository.delete_workflow_pause(pause_entity)

        # Assert - Pause record should be deleted
        self.session.expire_all()  # Clear session to ensure fresh query
        deleted_pause = self.session.get(WorkflowPauseModel, pause_entity.id)
        assert deleted_pause is None

        try:
            content = storage.load(file_key).decode()
            pytest.fail("File should be deleted from storage after pause deletion")
        except FileNotFoundError:
            # This is expected - file should be deleted from storage
            pass
        except Exception as e:
            pytest.fail(f"Unexpected error when checking file deletion: {e}")

    def test_large_state_file_handling(self):
        """Test handling of large state files."""
        # Arrange - Create a large state (1MB)
        large_state = "x" * (1024 * 1024)  # 1MB of data
        large_state_json = json.dumps({"large_data": large_state})

        workflow_run = self._create_test_workflow_run()
        repository = self._get_workflow_run_repository()

        # Act
        pause_entity = repository.create_workflow_pause(
            workflow_run_id=workflow_run.id,
            state_owner_user_id=self.test_user_id,
            state=large_state_json,
            pause_reasons=[],
        )

        # Assert
        assert pause_entity is not None
        retrieved_state = pause_entity.get_state()
        if isinstance(retrieved_state, bytes):
            retrieved_state = retrieved_state.decode()
        assert retrieved_state == large_state_json

        # Verify file size in database
        pause_model = self.session.get(WorkflowPauseModel, pause_entity.id)
        assert pause_model.state_object_key != ""
        loaded_state = storage.load(pause_model.state_object_key)
        assert loaded_state.decode() == large_state_json

    def test_multiple_pause_resume_cycles(self):
        """Test multiple pause/resume cycles on the same workflow run."""
        # Arrange
        workflow_run = self._create_test_workflow_run()
        repository = self._get_workflow_run_repository()

        # Act & Assert - Multiple cycles
        for i in range(3):
            state = json.dumps({"cycle": i, "data": f"state_{i}"})

            # Reset workflow run status to RUNNING before each pause (after first cycle)
            if i > 0:
                self.session.refresh(workflow_run)  # Refresh to get latest state from session
                workflow_run.status = WorkflowExecutionStatus.RUNNING
                self.session.commit()
                self.session.refresh(workflow_run)  # Refresh again after commit

            # Pause
            pause_entity = repository.create_workflow_pause(
                workflow_run_id=workflow_run.id, state_owner_user_id=self.test_user_id, state=state, pause_reasons=[]
            )
            assert pause_entity is not None

            # Verify pause
            self.session.expire_all()  # Clear session to ensure fresh query
            self.session.refresh(workflow_run)

            # Use the test session directly to verify the pause
            stmt = select(WorkflowRun).options(selectinload(WorkflowRun.pause)).where(WorkflowRun.id == workflow_run.id)
            workflow_run_with_pause = self.session.scalar(stmt)
            pause_model = workflow_run_with_pause.pause

            # Verify pause using test session directly
            assert pause_model is not None
            assert pause_model.id == pause_entity.id
            assert pause_model.state_object_key != ""

            # Load file content using storage directly
            file_content = storage.load(pause_model.state_object_key)
            if isinstance(file_content, bytes):
                file_content = file_content.decode()
            assert file_content == state

            # Resume
            resumed_entity = repository.resume_workflow_pause(
                workflow_run_id=workflow_run.id,
                pause_entity=pause_entity,
            )
            assert resumed_entity is not None
            assert resumed_entity.resumed_at is not None

            # Verify resume - check that pause is marked as resumed
            self.session.expire_all()  # Clear session to ensure fresh query
            stmt = select(WorkflowPauseModel).where(WorkflowPauseModel.id == pause_entity.id)
            resumed_pause_model = self.session.scalar(stmt)
            assert resumed_pause_model is not None
            assert resumed_pause_model.resumed_at is not None

            # Verify workflow run status
            self.session.refresh(workflow_run)
            assert workflow_run.status == WorkflowExecutionStatus.RUNNING
