"""Module containing client for interacting with the Census API"""
from typing import Any, Dict, Optional

import prefect
from httpx import AsyncClient, Response
from prefect.utilities.asyncutils import sync_compatible


class CensusClient:
    """
    Client for interacting with the Census API.

    Attributes:
        api_key (str): API key to authenticate with the Census API.
    """

    def __init__(self, api_key: str):
        self._closed = False
        self._started = False

        self.client = AsyncClient(
            base_url="https://app.getcensus.com/api/v1",
            headers={
                "Authorization": f"Bearer {api_key}",
                "user-agent": f"prefect-{prefect.__version__}",
            },
        )

    @sync_compatible
    async def call_endpoint(
        self,
        http_method: str,
        path: str,
        params: Optional[Dict[str, Any]] = None,
        json: Optional[Dict[str, Any]] = None,
    ) -> Response:
        """
        Call an endpoint in the Census API.

        Args:
            http_method: HTTP method to call on the endpoint.
            path: The partial path for request (e.g. /syncs/42). Will be
                appended onto the base URL as determined by the client configuration.
            params: Query parameters to include in the request.
            json: JSON serializable body to send in the request.

        Returns:
            The response from the Census API.

        Example:
            ```python
            from prefect import flow
            from prefect_census import CensusCredentials
            from prefect_census.client import CensusClient

            @flow
            def my_flow(sync_id):
                creds_block = CensusCredentials(api_key="my_api_key")

                client = CensusClient(
                    api_key=creds_block.api_key.get_secret_value()
                )
                response = client.call_endpoint(
                    http_method="GET",
                    path=f"/syncs/{sync_id}"
                )
                return response

            my_flow(42)
            ```
        """

        response = await self.client.request(
            method=http_method, url=path, params=params, json=json
        )
        response.raise_for_status()
        return response

    @sync_compatible
    async def get_run_info(self, run_id: int) -> Response:
        """
        Sends a request to the [get sync id info endpoint](https://docs.getcensus.com/basics/api/syncs#get-syncs-id)

        Args:
            run_id: The ID of the sync run to get details for.

        Returns:
            The response from the Census API.
        """  # noqa
        return await self.call_endpoint(
            http_method="GET",
            path=f"/sync_runs/{run_id}",
        )

    @sync_compatible
    async def trigger_sync_run(
        self, sync_id: int, force_full_sync: bool = False
    ) -> Response:
        """
        Sends a request to the [trigger sync run endpoint](https://docs.getcensus.com/basics/api/sync-runs)
        to initiate a sync run.

        Args:
            sync_id: The ID of the sync to trigger.
            force_full_sync: If the sync should perform a full sync.

        Returns:
            The response from the Census API.
        """  # noqa
        return await self.call_endpoint(
            http_method="POST",
            path=f"/syncs/{sync_id}/trigger",
            params={"force_full_sync": force_full_sync},
        )

    async def __aenter__(self):
        """Async context manager entry method."""
        if self._closed:
            raise RuntimeError(
                "The client cannot be started again after it has been closed."
            )
        if self._started:
            raise RuntimeError("The client cannot be started more than once.")

        self._started = True
        return self

    async def __aexit__(self, *exc):
        """Async context manager exit method."""
        self._closed = True
        await self.client.__aexit__()
