import abc
from typing import Any, Literal

import polars as pl
from pydantic import Field

from bayesline.api._src._utils import docstrings_from_sync
from bayesline.api._src.equity.portfolioreport import (
    AsyncReportAccessorApi,
    ReportAccessorApi,
)
from bayesline.api._src.equity.report.accessor import (
    AsyncTypedReportAccessorApi,
    TypedReportAccessorApi,
)
from bayesline.api._src.equity.report.api import AsyncReportApi, ReportApi
from bayesline.api._src.equity.report.settings import ReportSettings
from bayesline.api._src.equity.universe_settings import UniverseSettings
from bayesline.api._src.tasks import AsyncTask, Task
from bayesline.api._src.types import DateLike


class UniverseCountReportSettings(ReportSettings):
    """Settings for a universe count report."""

    report_type: Literal["universe_count"] = Field(
        default="universe_count", description="Report type discriminator"
    )
    universe_settings: UniverseSettings = Field()
    date_aggregation: Literal["last", "mean", "sum"] = Field(
        default="last", description="The date aggregation to use"
    )


class UniverseCountReportAccessor(TypedReportAccessorApi):
    """Specific accessor for a universe count report."""

    @abc.abstractmethod
    def get_counts(  # noqa: D102
        self,
        *,
        dates: bool = True,
        categorical_hierarchy_levels: dict[str, int] | None = None,
    ) -> pl.DataFrame:
        """
        Get the counts for the universe.

        Parameters
        ----------
        dates : bool, default=True
            If True, groups by dates.
        categorical_hierarchy_levels : dict[str, int] | None, default=None
            The level of categorical aggregation to group by.
            The key is the hierarchy name, the value is the level.
            If None, no categorical aggregation is done. A level of -1 means
            to use all levels.


        Returns
        -------
        pl.DataFrame
            A dataframe with `[date, count]` columns.
        """


@docstrings_from_sync
class AsyncUniverseCountReportAccessor(AsyncTypedReportAccessorApi):

    @abc.abstractmethod
    async def get_counts(  # noqa: D102
        self,
        *,
        dates: bool = True,
        categorical_hierarchy_levels: dict[str, int] | None = None,
    ) -> pl.DataFrame: ...


class UniverseCountReportApi(
    ReportApi[
        [DateLike | None, DateLike | None],
        UniverseCountReportAccessor,
        UniverseCountReportSettings,
    ]
):
    """API for a universe count report."""

    @abc.abstractmethod
    def get_maybe_cached(
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> UniverseCountReportAccessor | None:
        """
        Get the cached universe count report accessor API if it is available.

        Parameters
        ----------
        start_date: DateLike | None
            The start date of the report.
        end_date: DateLike | None
            The end date of the report.
        **kwargs: Any
            Additional keyword arguments.

        Returns
        -------
        UniverseCountReportAccessor | None
            The cached universe count report accessor API, or None if the report is
            not cached.
        """

    @abc.abstractmethod
    def calculate(
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> UniverseCountReportAccessor:
        """
        Calculate the universe count report.

        Parameters
        ----------
        start_date: DateLike | None
            The start date of the report.
        end_date: DateLike | None
            The end date of the report.
        **kwargs: Any
            Additional keyword arguments.

        Returns
        -------
        UniverseCountReportAccessor
            The universe count report accessor.
        """

    @abc.abstractmethod
    def calculate_as_task(  # noqa: D102
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> Task[UniverseCountReportAccessor]:
        """
        Calculate the universe count report as a task.

        Parameters
        ----------
        start_date: DateLike | None
            The start date of the report.
        end_date: DateLike | None
            The end date of the report.
        **kwargs: Any
            Additional keyword arguments.

        Returns
        -------
        Task[UniverseCountReportAccessor]
            The task for the universe count report.
        """


@docstrings_from_sync
class AsyncUniverseCountReportApi(
    AsyncReportApi[
        [DateLike | None, DateLike | None],
        AsyncUniverseCountReportAccessor,
        UniverseCountReportSettings,
    ]
):

    @abc.abstractmethod
    async def get_maybe_cached(  # noqa: D102
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> AsyncUniverseCountReportAccessor | None: ...

    @abc.abstractmethod
    async def calculate(  # noqa: D102
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> AsyncUniverseCountReportAccessor: ...

    @abc.abstractmethod
    async def calculate_as_task(  # noqa: D102
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> AsyncTask[AsyncUniverseCountReportAccessor]: ...


class UniverseCountReportAccessorImpl(UniverseCountReportAccessor):  # noqa: D101

    def __init__(self, accessor: ReportAccessorApi):
        self._accessor = accessor

    @property
    def accessor(self) -> ReportAccessorApi:  # noqa: D102
        return self._accessor

    def get_counts(  # noqa: D102
        self,
        *,
        dates: bool = True,
        categorical_hierarchy_levels: dict[str, int] | None = None,
    ) -> pl.DataFrame:
        expand = ["date"] if dates else []
        if categorical_hierarchy_levels:
            for name, max_level in categorical_hierarchy_levels.items():
                for name_level in self.accessor.axes["asset"]:
                    if not name_level.startswith(f"{name}_level_"):
                        continue
                    _, level = name_level.split("_level_")
                    if max_level == -1 or int(level) <= max_level:
                        expand.append(name_level)
        df = self.accessor.get_data([], expand=tuple(expand), value_cols=("count",))
        return df.filter(pl.col("count") != 0)


class AsyncUniverseCountReportAccessorImpl(  # noqa: D101
    AsyncUniverseCountReportAccessor
):
    def __init__(self, accessor: AsyncReportAccessorApi):
        self._accessor = accessor

    @property
    def accessor(self) -> AsyncReportAccessorApi:  # noqa: D102
        return self._accessor

    async def get_counts(  # noqa: D102
        self,
        *,
        dates: bool = True,
        categorical_hierarchy_levels: dict[str, int] | None = None,
    ) -> pl.DataFrame:
        expand = ["date"] if dates else []
        if categorical_hierarchy_levels:
            for name, max_level in categorical_hierarchy_levels.items():
                for name_level in self.accessor.axes["asset"]:
                    if not name_level.startswith(f"{name}_level_"):
                        continue
                    _, level = name_level.split("_level_")
                    if max_level == -1 or int(level) <= max_level:
                        expand.append(name_level)
        df = await self.accessor.get_data(
            [], expand=tuple(expand), value_cols=("count",)
        )
        return df.filter(pl.col("count") != 0)
