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.riskmodels_settings import FactorRiskModelSettings
from bayesline.api._src.tasks import AsyncTask, Task
from bayesline.api._src.types import DateLike


class McapExposureReportSettings(ReportSettings):
    """Settings for a market-cap-weighted exposure report."""

    report_type: Literal["mcap_exposure"] = Field(
        default="mcap_exposure", description="Report type discriminator"
    )
    factor_model_settings: FactorRiskModelSettings = Field()


class McapExposureReportAccessor(TypedReportAccessorApi):
    """Specific accessor for a market-cap-weighted exposure report."""

    @abc.abstractmethod
    def get_mcap_weighted_exposures(self) -> pl.DataFrame:
        """
        Get the time-series of market-cap-weighted exposures for the factors.

        For each date and factor, this computes the sum of exposure_i * mcap_i
        across all assets in the estimation universe.

        Returns
        -------
        pl.DataFrame
            A dataframe with a `date` column and columns for each factor with their
            respective market-cap-weighted exposures, in factor_group^factor format.
        """


@docstrings_from_sync
class AsyncMcapExposureReportAccessor(AsyncTypedReportAccessorApi):

    @abc.abstractmethod
    async def get_mcap_weighted_exposures(self) -> pl.DataFrame: ...  # noqa: D102


class McapExposureReportApi(
    ReportApi[
        [DateLike, DateLike], McapExposureReportAccessor, McapExposureReportSettings
    ]
):
    """API for a market-cap-weighted exposure report."""

    @abc.abstractmethod
    def get_maybe_cached(
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> McapExposureReportAccessor | None:
        """
        Get the cached mcap exposure 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
        -------
        McapExposureReportAccessor | None
            The cached 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,
    ) -> McapExposureReportAccessor:
        """
        Calculate the market-cap-weighted exposure 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
        -------
        McapExposureReportAccessor
            The report accessor.
        """

    @abc.abstractmethod
    def calculate_as_task(
        self,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> Task[McapExposureReportAccessor]:
        """
        Calculate the market-cap-weighted exposure 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[McapExposureReportAccessor]
            The task for the report accessor.
        """


@docstrings_from_sync
class AsyncMcapExposureReportApi(
    AsyncReportApi[
        [DateLike | None, DateLike | None],
        AsyncMcapExposureReportAccessor,
        McapExposureReportSettings,
    ]
):

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

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

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


class McapExposureReportAccessorImpl(McapExposureReportAccessor):  # noqa: D101

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

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

    def get_mcap_weighted_exposures(self) -> pl.DataFrame:  # noqa: D102
        df = self.accessor.get_data(
            [],
            expand=("date",),
            value_cols=("mcap_exposure",),
            pivot_cols=("factor_group", "factor"),
        )
        return df.rename(lambda x: x.removesuffix("^mcap_exposure"))


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

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

    async def get_mcap_weighted_exposures(self) -> pl.DataFrame:  # noqa: D102
        df = await self.accessor.get_data(
            [],
            expand=("date",),
            value_cols=("mcap_exposure",),
            pivot_cols=("factor_group", "factor"),
        )
        return df.rename(lambda x: x.removesuffix("^mcap_exposure"))
