import pytest
from prefect import flow

from prefect_census.credentials import CensusCredentials
from prefect_census.flows import run_census_sync
from prefect_census.runs import (
    CensusSyncRunCancelled,
    CensusSyncRunFailed,
    CensusSyncRunTimeout,
)
from prefect_census.syncs import (
    CensusSync,
    CensusSyncTriggerFailed,
    trigger_census_sync,
    trigger_census_sync_run_and_wait_for_completion,
)


@pytest.fixture
def census_credentials():
    return CensusCredentials(api_key="my_api_key")


@pytest.fixture
def census_sync():
    return CensusSync(
        credentials=CensusCredentials(api_key="my_api_key"),
        sync_id=5,
    )


class TestTriggerCensusSyncRun:
    async def test_trigger_sync(self, mock_trigger_sync_calls, census_credentials):
        @flow
        async def test_flow():
            return await trigger_census_sync(credentials=census_credentials, sync_id=45)

        result = await test_flow()
        assert result == 45

    async def test_trigger_nonexistent_sync(
        self, mock_sync_not_found, census_credentials
    ):
        @flow
        async def test_trigger_nonexistent_job():
            await trigger_census_sync(census_credentials, sync_id=46)

        with pytest.raises(CensusSyncTriggerFailed, match="Not found!"):
            await test_trigger_nonexistent_job()


class TestTriggerCensusSyncRunAndWaitForCompletion:
    async def test_run_success(self, mock_successful_sync_calls, census_credentials):

        result = await trigger_census_sync_run_and_wait_for_completion(
            credentials=census_credentials, sync_id=5
        )

        assert result == {"id": 5, "status": "completed", "sync_run_id": 12345}

    def test_run_in_sync_flow(self, mock_successful_sync, census_credentials):
        @flow
        def test_sync_flow():
            return trigger_census_sync_run_and_wait_for_completion(
                credentials=census_credentials, sync_id=5
            )

        result = test_sync_flow()

        assert result == {"id": 5, "status": "completed", "sync_run_id": 12345}

    async def test_run_success_with_wait(
        self, mock_successful_sync_with_wait, census_credentials
    ):

        result = await trigger_census_sync_run_and_wait_for_completion(
            credentials=census_credentials, sync_id=5, poll_frequency_seconds=2
        )

        assert result == {"id": 5, "status": "completed", "sync_run_id": 12345}

    async def test_run_failure_with_wait(
        self, mock_failed_sync_with_wait, census_credentials
    ):
        with pytest.raises(CensusSyncRunFailed):
            await trigger_census_sync_run_and_wait_for_completion(
                credentials=census_credentials, sync_id=5, poll_frequency_seconds=2
            )

    async def test_run_failure_no_run_id(
        self, mock_sync_failure_no_run_id, census_credentials
    ):
        with pytest.raises(KeyError, match="sync_run_id"):
            await trigger_census_sync_run_and_wait_for_completion(
                credentials=census_credentials, sync_id=5, poll_frequency_seconds=2
            )

    async def test_run_cancelled_with_wait(
        self, mock_cancelled_sync_with_wait, census_credentials
    ):
        with pytest.raises(CensusSyncRunCancelled):
            await trigger_census_sync_run_and_wait_for_completion(
                credentials=census_credentials, sync_id=5, poll_frequency_seconds=2
            )

    async def test_run_timed_out(self, mock_sync_timed_out, census_credentials):
        with pytest.raises(CensusSyncRunTimeout):
            await trigger_census_sync_run_and_wait_for_completion(
                credentials=census_credentials,
                sync_id=5,
                poll_frequency_seconds=1,
                max_wait_seconds=3,
            )


class TestRunCensusSync:
    async def test_run_success(self, mock_successful_sync_calls, census_sync):

        result = await run_census_sync(census_sync)

        assert result.final_status.value == "completed"
        assert result.run_data == {"id": 5, "status": "completed", "sync_run_id": 12345}

    def test_run_in_sync_flow(self, mock_successful_sync, census_sync):
        @flow
        def test_sync_flow():
            return run_census_sync(census_sync=census_sync)

        result = test_sync_flow()

        assert result.final_status.value == "completed"
        assert result.run_data == {"id": 5, "status": "completed", "sync_run_id": 12345}

    async def test_run_success_with_wait(
        self, mock_successful_sync_with_wait, census_sync
    ):

        result = await run_census_sync(census_sync)

        assert result.final_status.value == "completed"
        assert result.run_data == {"id": 5, "status": "completed", "sync_run_id": 12345}

    async def test_run_failure_with_wait(self, mock_failed_sync_with_wait, census_sync):
        with pytest.raises(CensusSyncRunFailed):
            await run_census_sync(census_sync)

    async def test_run_failure_no_run_id(
        self, mock_sync_failure_no_run_id, census_sync
    ):
        with pytest.raises(KeyError, match="sync_run_id"):
            await run_census_sync(census_sync)

    async def test_run_cancelled_with_wait(
        self, mock_cancelled_sync_with_wait, census_sync
    ):
        with pytest.raises(CensusSyncRunCancelled):
            census_sync.poll_frequency_seconds = 2
            await run_census_sync(census_sync)

    async def test_run_timed_out(self, mock_sync_timed_out, census_sync):
        with pytest.raises(CensusSyncRunTimeout):
            census_sync.poll_frequency_seconds = 2
            census_sync.max_wait_seconds = 3
            await run_census_sync(census_sync)
