import abc
from typing import Literal, overload

import polars as pl

from bayesline.api._src._utils import docstrings_from_sync
from bayesline.api._src.equity.riskdataset_settings import (
    DatasetStatus,
    RiskDatasetMetadata,
    RiskDatasetProperties,
    RiskDatasetSettings,
    RiskDatasetSettingsMenu,
    RiskDatasetUpdateResult,
)
from bayesline.api._src.registry import (
    AsyncRegistryBasedApi,
    RawSettings,
    RegistryBasedApi,
)
from bayesline.api._src.tasks import AsyncTask, Task


class DatasetError(Exception):
    """Exception raised for dataset-related errors."""

    ...


class RiskDatasetApi(abc.ABC):
    """API for managing risk datasets.

    This abstract base class defines the interface for risk dataset operations,
    including updating datasets and retrieving their properties.
    """

    @property
    @abc.abstractmethod
    def settings(self) -> RiskDatasetSettings:
        """Get the settings for this risk dataset.

        Returns
        -------
        RiskDatasetSettings
            The settings for this risk dataset.
        """
        ...

    @abc.abstractmethod
    def describe(self) -> RiskDatasetProperties:
        """Get the properties of this risk dataset.

        Returns
        -------
        RiskDatasetProperties
            The properties of this risk dataset.
        """
        ...

    @abc.abstractmethod
    def update(self, force: bool = False) -> RiskDatasetUpdateResult:
        """Check underlying datasets for updates and update this dataset to latest versions.

        Parameters
        ----------
        force : bool, default=False
            If true, the update will be forced even if the dataset is up to date.

        Raises
        ------
        DatasetError
            If an error occurs.

        Returns
        -------
        RiskDatasetUpdateResult
            The result of the update operation.
        """

    @abc.abstractmethod
    def update_as_task(self) -> Task[RiskDatasetUpdateResult]: ...  # noqa: D102


class RiskDatasetLoaderApi(
    RegistryBasedApi[RiskDatasetSettings, RiskDatasetSettingsMenu, RiskDatasetApi],
):
    """API for loading and managing risk datasets.

    This class provides functionality for loading, creating, and managing risk datasets
    through a registry-based approach.
    """

    @abc.abstractmethod
    def get_default_dataset_name(self) -> str:
        """Get the default dataset name.

        Returns
        -------
        str
            The default dataset name that will be populated into settings if no
            dataset name is provided.
        """

    @abc.abstractmethod
    def get_dataset_names(
        self,
        *,
        mode: Literal["System", "User"] | None = None,
    ) -> dict[str, DatasetStatus]:
        """Get the names of all available datasets.

        Parameters
        ----------
        mode : Literal["System", "User"] | None, default=None
            System: only system wide datasets (available to all users and provided
            by the system).
            User: only user specific datasets (available to the current user).
            None: all datasets (system wide and user specific).

        Returns
        -------
        dict[str, DatasetStatus]
            The names of all available datasets and their status.
            "ready" means the dataset is ready to use.
            "data_stale" means the dataset is ready to use but newer data is available.
            "needs_update" means local data needs to be computed/updated.
            "not_ready" means the reference dataset is not in a usable state.
            "corrupt" means the underlying settings are corrupted and need to be repaired.
        """

    @overload
    def list_riskdatasets(
        self, name: str | None = None, *, df: Literal[False] = ...
    ) -> list[RiskDatasetMetadata]: ...

    @overload
    def list_riskdatasets(
        self, name: str | None = None, *, df: Literal[True]
    ) -> pl.DataFrame: ...

    @abc.abstractmethod
    def list_riskdatasets(
        self, name: str | None = None, *, df: bool = False
    ) -> list[RiskDatasetMetadata] | pl.DataFrame:
        """List all available risk datasets.

        Parameters
        ----------
        name : str | None, default=None
            The name of the dataset to filter by.
            If not given, all datasets are returned.
        df : bool, default=False
            If True, return a polars DataFrame instead of a list of metadata objects.

        Returns
        -------
        list[RiskDatasetMetadata] | pl.DataFrame
            The metadata of all available datasets, either as a list of objects
            or as a DataFrame depending on the `df` parameter.
        """

    def create_or_replace_dataset(
        self, name: str, settings: RiskDatasetSettings
    ) -> RiskDatasetApi:
        """Create a new dataset or replace an existing one.

        Parameters
        ----------
        name : str
            The name of the dataset to create or replace.
        settings : RiskDatasetSettings
            The settings for the dataset.

        Returns
        -------
        RiskDatasetApi
            The API for the created or replaced dataset.
        """
        self.delete_dataset_if_exists(name)
        return self.create_dataset(name, settings)

    @abc.abstractmethod
    def create_dataset(
        self, name: str, settings: RiskDatasetSettings
    ) -> RiskDatasetApi:
        """Create a new dataset with the given name and settings.

        Parameters
        ----------
        name : str
            The name of the dataset to create.
        settings : RiskDatasetSettings
            The settings for the dataset to create.

        Raises
        ------
        DatasetError
            If a dataset with the given name already exists
            or if the settings are otherwise invalid.

        Returns
        -------
        RiskDatasetApi
            The API of the newly created dataset.
        """

    @abc.abstractmethod
    def create_dataset_as_task(  # noqa: D102
        self, name: str, settings: RiskDatasetSettings
    ) -> Task[RiskDatasetApi]: ...

    @abc.abstractmethod
    def delete_dataset(self, name: str) -> RawSettings:
        """Delete the given dataset.

        Parameters
        ----------
        name : str
            The name of the dataset to delete.

        Raises
        ------
        KeyError
            If the dataset does not exist.
        DatasetError
            If the dataset could not be deleted.

        Returns
        -------
        RawSettings
            The raw settings of the deleted dataset.
        """

    def delete_dataset_if_exists(self, name: str) -> RawSettings | None:
        """Delete the given dataset if it exists.

        Parameters
        ----------
        name : str
            The name of the dataset to delete.

        Raises
        ------
        DatasetError
            If the dataset could not be deleted.

        Returns
        -------
        RawSettings | None
            The raw settings of the deleted dataset if dataset existed.
        """
        if name in self.settings.names(mode="All"):
            return self.delete_dataset(name)
        return None


@docstrings_from_sync
class AsyncRiskDatasetApi(abc.ABC):

    @property
    @abc.abstractmethod
    def settings(self) -> RiskDatasetSettings: ...  # noqa: D102

    @abc.abstractmethod
    async def describe(self) -> RiskDatasetProperties: ...  # noqa: D102

    @abc.abstractmethod
    async def update(  # noqa: D102
        self, force: bool = False
    ) -> RiskDatasetUpdateResult: ...

    @abc.abstractmethod
    async def update_as_task(  # noqa: D102
        self, force: bool = False
    ) -> AsyncTask[RiskDatasetUpdateResult]: ...


@docstrings_from_sync
class AsyncRiskDatasetLoaderApi(
    AsyncRegistryBasedApi[
        RiskDatasetSettings, RiskDatasetSettingsMenu, AsyncRiskDatasetApi
    ],
):

    @abc.abstractmethod
    async def get_default_dataset_name(self) -> str: ...  # noqa: D102

    @abc.abstractmethod
    async def get_dataset_names(  # noqa: D102
        self,
        *,
        mode: Literal["System", "User"] | None = None,
    ) -> dict[str, DatasetStatus]: ...

    @overload
    async def list_riskdatasets(
        self, name: str | None = None, *, df: Literal[False] = ...
    ) -> list[RiskDatasetMetadata]: ...

    @overload
    async def list_riskdatasets(
        self, name: str | None = None, *, df: Literal[True]
    ) -> pl.DataFrame: ...

    @abc.abstractmethod
    async def list_riskdatasets(  # noqa: D102
        self, name: str | None = None, *, df: bool = False
    ) -> list[RiskDatasetMetadata] | pl.DataFrame: ...

    async def create_or_replace_dataset(  # noqa: D102
        self, name: str, settings: RiskDatasetSettings
    ) -> AsyncRiskDatasetApi:
        await self.delete_dataset_if_exists(name)
        return await self.create_dataset(name, settings)

    @abc.abstractmethod
    async def create_dataset(  # noqa: D102
        self, name: str, settings: RiskDatasetSettings
    ) -> AsyncRiskDatasetApi: ...

    @abc.abstractmethod
    async def create_dataset_as_task(  # noqa: D102
        self, name: str, settings: RiskDatasetSettings
    ) -> AsyncTask[AsyncRiskDatasetApi]: ...

    @abc.abstractmethod
    async def delete_dataset(self, name: str) -> RawSettings: ...  # noqa: D102

    async def delete_dataset_if_exists(  # noqa: D102
        self, name: str
    ) -> RawSettings | None:
        if name in await self.settings.names(mode="All"):
            return await self.delete_dataset(name)
        return None
