"""Module containing tasks and flows for interacting with Census syncs."""
import asyncio
from typing import Any, Dict

from httpx import HTTPStatusError
from prefect import flow, task
from prefect.blocks.abstract import JobBlock, JobRun
from prefect.logging import get_run_logger
from prefect.utilities.asyncutils import sync_compatible
from pydantic import VERSION as PYDANTIC_VERSION

if PYDANTIC_VERSION.startswith("2."):
    from pydantic.v1 import BaseModel, Field
else:
    from pydantic import BaseModel, Field

from prefect_census.credentials import CensusCredentials
from prefect_census.runs import (
    CensusSyncRunCancelled,
    CensusSyncRunFailed,
    CensusSyncRunStatus,
    CensusSyncRunTimeout,
    wait_census_sync_completion,
)
from prefect_census.utils import extract_user_message


class CensusSyncTriggerFailed(RuntimeError):
    """Used to indicate sync triggered."""


class CensusSyncResult(BaseModel):
    """The results for a Census Sync Run."""

    final_status: CensusSyncRunStatus
    run_data: Dict[str, Any]


@task(
    name="Trigger Census sync run",
    description="Triggers a Census sync run for the sync with the given sync_id.",
    retries=3,
    retry_delay_seconds=10,
)
async def trigger_census_sync(
    credentials: CensusCredentials, sync_id: int, force_full_sync: bool = False
) -> int:
    """
    A task to trigger a Census sync run.

    Args:
        credentials: Credentials for authenticating with Census.
        sync_id: The ID of the sync to trigger.
        force_full_sync: If `True`, a full sync will be triggered.

    Returns:
        The ID of the triggered sync run.

    Examples:
        Trigger a Census sync run:
        ```python
        from prefect import flow

        from prefect_census import CensusCredentials
        from prefect_census.syncs import trigger_census_sync

        @flow
        def trigger_census_sync_flow():
            credentials = CensusCredentials(api_key="my_api_key")
            trigger_census_sync(credentials=credentials, sync_id=42)

        trigger_census_sync_flow()
        ```
    """
    logger = get_run_logger()

    logger.info(f"Triggering Census sync run for sync with ID {sync_id}")
    try:
        async with credentials.get_client() as client:
            response = await client.trigger_sync_run(
                sync_id=sync_id, force_full_sync=force_full_sync
            )
    except HTTPStatusError as e:
        raise CensusSyncTriggerFailed(extract_user_message(e)) from e

    run_data = response.json()["data"]

    if "sync_run_id" in run_data:
        logger.info(
            f"Census sync run successfully triggered for sync with ID {sync_id}. "
            "You can view the status of this sync run at "
            f"https://app.getcensus.com/sync/{sync_id}/sync-history"
        )

    return run_data["sync_run_id"]


@flow(
    name="Trigger Census sync run and wait for completion",
    description="Triggers a Census sync run and waits for the"
    "triggered run to complete.",
)
async def trigger_census_sync_run_and_wait_for_completion(
    credentials: CensusCredentials,
    sync_id: int,
    force_full_sync: bool = False,
    max_wait_seconds: int = 900,
    poll_frequency_seconds: int = 10,
) -> Dict[str, Any]:
    """
    Flow that triggers a sync run and waits for the triggered run to complete.

    Args:
        credentials: Credentials for authenticating with Census.
        sync_id: The ID of the sync to trigger.
        force_full_sync: If `True`, a full sync will be triggered.
        max_wait_seconds: Maximum number of seconds to wait for sync to complete
        poll_frequency_seconds: Number of seconds to wait in between checks for run completion.

    Raises:
        CensusSyncRunCancelled: The triggered Census sync run was cancelled.
        CensusSyncRunFailed: The triggered Census sync run failed.
        RuntimeError: The triggered Census sync run ended in an unexpected state.

    Returns:
        The final run data returned by the Census API as dict with the following shape:
            ```
            {
                "id": 94,
                "sync_id": 52,
                "source_record_count": 1,
                "records_processed": 1,
                "records_updated": 1,
                "records_failed": 0,
                "records_invalid": 0,
                "created_at": "2021-10-20T02:51:07.546Z",
                "updated_at": "2021-10-20T02:52:29.236Z",
                "completed_at": "2021-10-20T02:52:29.234Z",
                "scheduled_execution_time": null,
                "error_code": null,
                "error_message": null,
                "error_detail": null,
                "status": "completed",
                "canceled": false,
                "full_sync": true,
                "sync_trigger_reason": {
                    "ui_tag": "Manual",
                    "ui_detail": "Manually triggered by test@getcensus.com"
                }
            }
            ```

    Examples:
        Trigger a Census sync using CensusCredentials instance and wait
        for completion as a standalone flow:
        ```python
        import asyncio

        from prefect_census import CensusCredentials
        from prefect_census.syncs import trigger_census_sync_run_and_wait_for_completion

        asyncio.run(
            trigger_census_sync_run_and_wait_for_completion(
                credentials=CensusCredentials(
                    api_key="my_api_key"
                ),
                sync_id=42
            )
        )
        ```

        Trigger a Census sync and wait for completion as a subflow:
        ```python
        from prefect import flow

        from prefect_census import CensusCredentials
        from prefect_census.syncs import trigger_census_sync_run_and_wait_for_completion

        @flow
        def my_flow():
            ...
            creds = CensusCredentials(api_key="my_api_key")
            run_result = trigger_census_sync_run_and_wait_for_completion(
                credentials=creds,
                sync_id=42
            )
            ...

        my_flow()
        ```
    """  # noqa
    logger = get_run_logger()

    triggered_run_data_future = await trigger_census_sync.submit(
        credentials=credentials, sync_id=sync_id, force_full_sync=force_full_sync
    )

    run_id = await triggered_run_data_future.result()
    if run_id is None:
        raise RuntimeError("Unable to determine run ID for triggered sync.")

    final_run_status, run_data = await wait_census_sync_completion(
        run_id=run_id,
        credentials=credentials,
        max_wait_seconds=max_wait_seconds,
        poll_frequency_seconds=poll_frequency_seconds,
    )

    if final_run_status == CensusSyncRunStatus.COMPLETED:
        logger.info(
            "Census sync run with ID %s completed successfully!",
            run_id,
        )
        return run_data

    elif final_run_status == CensusSyncRunStatus.CANCELLED:
        raise CensusSyncRunCancelled(
            f"Triggered sync run with ID {run_id} was cancelled."
        )
    elif final_run_status == CensusSyncRunStatus.FAILED:
        raise CensusSyncRunFailed(f"Triggered sync run with ID: {run_id} failed.")
    else:
        raise RuntimeError(
            f"Triggered sync run with ID: {run_id} ended with unexpected"
            f"status {final_run_status}"
        )


class CensusSync(JobBlock):
    """
    A Job Block for triggering a Census sync run and waiting for completion.

    Attributes:
        credentials: Credentials for authenticating with Census.
        sync_id: The ID of the sync to trigger.
        force_full_sync: If `True`, a full sync will be triggered.
        max_wait_seconds: Maximum number of seconds to wait for sync to complete.
        poll_frequency_seconds: Number of seconds to wait in between checks
            for run completion.

    Example:
        Trigger a Census sync and wait for completion as a subflow:
        ```python
        from prefect import flow
        from prefect_census import CensusSync, run_census_sync

        @flow
        def my_census_flow():
            census_sync = CensusSync.load("BLOCK_NAME")

            # do some setup

            run_census_sync(census_sync)

            # do some cleanup
        ```
    """

    credentials: CensusCredentials = Field(
        default=...,
        description="Credentials for authenticating with Census.",
    )

    sync_id: int = Field(
        default=...,
        description="The ID of the sync to trigger.",
    )

    force_full_sync: bool = Field(
        default=False,
        description="If `True`, a full sync will be triggered.",
    )

    max_wait_seconds: int = Field(
        default=3600,
        description="The maximum number of seconds to wait for the sync to complete.",
    )

    poll_frequency_seconds: int = Field(
        default=5,
        description="Number of seconds to wait between sync status checks.",
    )

    _block_type_name = "Census Sync"
    _description = "Runs a Census sync"
    _documentation_url = "https://prefecthq.github.io/prefect-census/syncs/"
    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/b2f805555c778f37f861b67c4a861908f66a6e35-700x700.png"  # noqa

    @sync_compatible
    async def trigger(self) -> "CensusSyncRun":
        """
        Trigger a Census sync run.

        Returns:
            A CensusSyncRun instance representing the triggered sync run.
        """
        self.logger.info(f"Triggering Census sync run for sync with ID {self.sync_id}")
        try:
            async with self.credentials.get_client() as client:
                response = await client.trigger_sync_run(
                    sync_id=self.sync_id, force_full_sync=self.force_full_sync
                )
        except HTTPStatusError as e:
            raise CensusSyncTriggerFailed(extract_user_message(e)) from e

        run_data = response.json()["data"]

        if "sync_run_id" in run_data:
            self.logger.info(
                f"Census sync with ID: {self.sync_id} successfully triggered. "
                "You can view the status of this sync run at "
                f"https://app.getcensus.com/sync/{self.sync_id}/sync-history"
            )

        run_id = run_data["sync_run_id"]
        if run_id is None:
            raise RuntimeError("Unable to determine run ID for triggered sync.")

        return CensusSyncRun(census_sync=self, run_id=run_id)


class CensusSyncRun(JobRun):
    """
    A Job Run representing a Census sync run.

    Attributes:
        sync: The CensusSync instance that triggered this run.
        run_id: The ID of the Census sync run.
        run_data: The data returned by the Census API for this run.
        status: The status of the Census sync run.
    """

    def __init__(self, census_sync, run_id) -> None:
        self.sync = census_sync
        self.run_id = run_id
        self.run_data = None
        self.status = None

    @sync_compatible
    async def wait_for_completion(self):
        """Wait for the Census sync run to complete."""
        seconds_waited_for_run_completion = 0
        last_status = None
        while seconds_waited_for_run_completion <= self.sync.max_wait_seconds:
            try:
                async with self.sync.credentials.get_client() as client:
                    response = await client.get_run_info(self.run_id)
            except HTTPStatusError as e:
                raise RuntimeError(extract_user_message(e)) from e

            run_data = response.json()["data"]
            self.status = run_data.get("status")

            if CensusSyncRunStatus.is_terminal_status_code(self.status):
                self.run_data = run_data
                if self.status == CensusSyncRunStatus.COMPLETED.value:
                    self.logger.info(
                        "Census sync run with ID %s completed successfully!",
                        self.run_id,
                    )
                    return

                elif self.status == CensusSyncRunStatus.CANCELLED.value:
                    raise CensusSyncRunCancelled(
                        f"Triggered sync run with ID {self.run_id} was cancelled."
                    )
                elif self.status == CensusSyncRunStatus.FAILED.value:
                    raise CensusSyncRunFailed(
                        f"Triggered sync run with ID {self.run_id} has failed."
                    )
                else:
                    raise RuntimeError(
                        f"Sync run with ID: {self.run_id} ended with unexpected "
                        f"status {self.status}"
                    )
            if self.status != last_status:
                self.logger.info(
                    "Census sync run with ID %i has status %s.",
                    self.run_id,
                    CensusSyncRunStatus(self.status).name,
                )
                last_status = self.status

            await asyncio.sleep(self.sync.poll_frequency_seconds)
            seconds_waited_for_run_completion += self.sync.poll_frequency_seconds

        raise CensusSyncRunTimeout(
            f"Timeout of {self.sync.max_wait_seconds} seconds exceeded while "
            f"waiting for sync run with ID {self.run_id} to complete."
        )

    @sync_compatible
    async def fetch_result(self) -> CensusSyncResult:
        """Fetch the result of the Census sync run.

        Returns:
            A CensusSyncResult instance representing the result of the sync run.
        """
        return CensusSyncResult(
            final_status=CensusSyncRunStatus(self.status),
            run_data=self.run_data,
        )
