import abc
from typing import Any, Literal, Self

import polars as pl
from pydantic import Field, model_validator

from bayesline.api._src._utils import docstrings_from_sync
from bayesline.api._src.equity.calendar_settings import CalendarSettings
from bayesline.api._src.equity.portfoliohierarchy_settings import (
    PortfolioHierarchySettings,
)
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.scenarios import (
    PortfolioHoldingsScenario,
)
from bayesline.api._src.equity.report.settings import ReportSettings
from bayesline.api._src.tasks import AsyncTask, Task
from bayesline.api._src.types import DateLike, DNFFilterExpressions


class PortfolioReturnsReportSettings(ReportSettings):
    """Settings for a portfolio returns report."""

    report_type: Literal["portfolio_returns"] = Field(
        default="portfolio_returns", description="Report type discriminator"
    )
    portfolio_hierarchy_settings: PortfolioHierarchySettings = Field()
    calendar: CalendarSettings = Field(default_factory=CalendarSettings)
    currency: str | None = Field(None)
    returns_currency: str = Field("USD")
    rescale_bench: bool = Field(True)
    normalize_holdings: bool = Field(True)
    return_aggregation_type: Literal["arithmetic", "geometric"] = Field("geometric")

    @model_validator(mode="after")
    def validate_geometric_requires_normalized(self) -> Self:
        if self.return_aggregation_type == "geometric" and not self.normalize_holdings:
            raise ValueError(
                "Geometric compounding requires normalize_holdings=True. "
                "PnL values cannot be geometrically compounded."
            )
        return self


class PortfolioReturnsReportAccessor(TypedReportAccessorApi):
    """Specific accessor for a portfolio returns report."""

    scenario_types = [PortfolioHoldingsScenario]

    @abc.abstractmethod
    def get_portolio_returns(
        self,
        portfolio_names: list[str] | None = None,
        start_date: DateLike | None = None,
        end_date: DateLike | None = None,
        aggregation_type: Literal["D", "M", "Y"] | None = "D",
        value_cols: list[str] | None = None,
    ) -> pl.DataFrame:
        """
        Get the portfolio returns for the portfolio names between start and end dates.

        Parameters
        ----------
        portfolio_names: list[str] | None, default=None
            The names of the portfolios to calculate the report for. If None, all
            portfolios will be calculated.
        start_date: DateLike | None, default=None
            The start date of the report.
        end_date: DateLike | None, default=None
            The end date of the report.
        aggregation_type: Literal["D", "M", "Y"] | None, default="D"
            The aggregation type to use for the returns.
            If None, dates will be completely collapsed and the returns for the entire
            period will be returned.
        value_cols: list[str] | None, default=None
            The value columns to return. If None, all value columns will be returned.

        Returns
        -------
        pl.DataFrame
            A dataframe with the portfolio returns.
        """

    @abc.abstractmethod
    def get_asset_returns(
        self,
        portfolio_id: str,
        start_date: DateLike | None = None,
        end_date: DateLike | None = None,
        aggregation_type: Literal["D", "M", "Y"] | None = "D",
        value_cols: list[str] | None = None,
    ) -> pl.DataFrame:
        """
        Get the asset level returns for the portfolio id between start and end dates.

        Parameters
        ----------
        portfolio_id: str
            The id of the portfolio to calculate the report for.
        start_date: DateLike | None, default=None
            The start date of the report.
        end_date: DateLike | None, default=None
            The end date of the report.
        aggregation_type: Literal["D", "M", "Y"] | None, default="D"
            The aggregation type to use for the returns.
            If None, dates will be completely collapsed and the returns for the entire
            period will be returned.
        value_cols: list[str] | None, default=None
            The value columns to return. If None, all value columns will be returned.

        Returns
        -------
        pl.DataFrame
            A dataframe with the asset level returns.
        """


@docstrings_from_sync
class AsyncPortfolioReturnsReportAccessor(AsyncTypedReportAccessorApi):
    scenario_types = [PortfolioHoldingsScenario]

    @abc.abstractmethod
    async def get_portolio_returns(  # noqa: D102
        self,
        portfolio_names: list[str] | None = None,
        start_date: DateLike | None = None,
        end_date: DateLike | None = None,
        aggregation_type: Literal["D", "M", "Y"] | None = "D",
        value_cols: list[str] | None = None,
    ) -> pl.DataFrame: ...

    @abc.abstractmethod
    async def get_asset_returns(  # noqa: D102
        self,
        portfolio_id: str,
        start_date: DateLike | None = None,
        end_date: DateLike | None = None,
        aggregation_type: Literal["D", "M", "Y"] | None = "D",
        value_cols: list[str] | None = None,
    ) -> pl.DataFrame: ...


class PortfolioReturnsReportApi(
    ReportApi[
        [list[str] | None, DateLike | None, DateLike | None],
        PortfolioReturnsReportAccessor,
        PortfolioReturnsReportSettings,
    ]
):
    """API for a portfolio returns report."""

    @abc.abstractmethod
    def get_maybe_cached(
        self,
        portfolio_names: list[str] | None,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> PortfolioReturnsReportAccessor | None:
        """
        Get the cached portfolio returns report accessor API if it is available.

        Parameters
        ----------
        portfolio_names: list[str] | None
            The names of the portfolios to calculate the report for. If None, all
            portfolios will be calculated.
        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
        -------
        PortfolioReturnsReportAccessor | None
            The cached portfolio returns report accessor API, or None if the report is
            not cached.
        """

    @abc.abstractmethod
    def calculate(
        self,
        portfolio_names: list[str] | None,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> PortfolioReturnsReportAccessor:
        """
        Calculate the portfolio returns report.

        Parameters
        ----------
        portfolio_names: list[str] | None
            The names of the portfolios to calculate the report for. If None, all
            portfolios will be calculated.
        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
        -------
        PortfolioReturnsReportAccessor
            The portfolio returns report accessor.
        """

    @abc.abstractmethod
    def calculate_as_task(
        self,
        portfolio_names: list[str] | None,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> Task[PortfolioReturnsReportAccessor]:
        """
        Calculate the portfolio returns report as a task.

        Parameters
        ----------
        portfolio_names: list[str] | None
            The names of the portfolios to calculate the report for. If None, all
            portfolios will be calculated.
        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[PortfolioReturnsReportAccessor]
            The task for the portfolio returns report accessor.
        """


@docstrings_from_sync
class AsyncPortfolioReturnsReportApi(
    AsyncReportApi[
        [list[str] | None, DateLike | None, DateLike | None],
        AsyncPortfolioReturnsReportAccessor,
        PortfolioReturnsReportSettings,
    ]
):

    @abc.abstractmethod
    async def get_maybe_cached(  # noqa: D102
        self,
        portfolio_names: list[str] | None,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> AsyncPortfolioReturnsReportAccessor | None: ...

    @abc.abstractmethod
    async def calculate(  # noqa: D102
        self,
        portfolio_names: list[str] | None,
        start_date: DateLike | None,
        end_date: DateLike | None,
        **kwargs: Any,
    ) -> AsyncPortfolioReturnsReportAccessor: ...

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


def _get_portfolio_returns_filters(
    start_date: DateLike | None,
    end_date: DateLike | None,
    portfolio_names: list[str] | None,
) -> DNFFilterExpressions:
    filters: DNFFilterExpressions = []
    if start_date is not None:
        filters.append(("date", ">=", start_date))
    if end_date is not None:
        filters.append(("date", "<=", end_date))
    if portfolio_names is not None:
        filters.append(("portfolio_id", "in", portfolio_names))
    return filters


def _get_portfolio_returns_expand(
    aggregation_type: Literal["D", "M", "Y"] | None = "D",
) -> tuple[str, ...]:
    expand = ["portfolio_id"]
    if aggregation_type == "D":
        expand.extend(["year", "month", "date"])
    elif aggregation_type == "M":
        expand.extend(["year", "month"])
    elif aggregation_type == "Y":
        expand.extend(["year"])
    return tuple(expand)


class PortfolioReturnsReportAccessorImpl(PortfolioReturnsReportAccessor):  # noqa: D101

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

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

    def get_portolio_returns(  # noqa: D102
        self,
        portfolio_names: list[str] | None = None,
        start_date: DateLike | None = None,
        end_date: DateLike | None = None,
        aggregation_type: Literal["D", "M", "Y"] | None = "D",
        value_cols: list[str] | None = None,
    ) -> pl.DataFrame:
        available_value_cols = self._accessor.metric_cols
        if value_cols is not None:
            if set(value_cols) > set(available_value_cols):
                raise ValueError(
                    f"Unknown value_cols: "
                    f"{set(value_cols) - set(available_value_cols)}. "
                    f"Available: {available_value_cols}."
                )
        else:
            value_cols = available_value_cols
        return self._accessor.get_data(
            path=list(),
            expand=_get_portfolio_returns_expand(aggregation_type),
            filters=_get_portfolio_returns_filters(
                start_date, end_date, portfolio_names
            ),
            value_cols=tuple(value_cols),
        )

    def get_asset_returns(  # noqa: D102
        self,
        portfolio_id: str,
        start_date: DateLike | None = None,
        end_date: DateLike | None = None,
        aggregation_type: Literal["D", "M", "Y"] | None = "D",
        value_cols: list[str] | None = None,
    ) -> pl.DataFrame:
        available_value_cols = self._accessor.metric_cols
        if value_cols is not None:
            if set(value_cols) > set(available_value_cols):
                raise ValueError(
                    f"Unknown value_cols: "
                    f"{set(value_cols) - set(available_value_cols)}. "
                    f"Available: {available_value_cols}."
                )
        else:
            value_cols = available_value_cols
        return self._accessor.get_data(
            path=[("portfolio_id", portfolio_id)],
            expand=tuple(
                list(
                    [
                        p
                        for p in _get_portfolio_returns_expand(aggregation_type)
                        if p != "portfolio_id"
                    ]
                    + [
                        "input_asset_id",
                        "input_asset_id_type",
                        "asset_id",
                        "asset_name",
                    ]
                )
            ),
            filters=_get_portfolio_returns_filters(
                start_date, end_date, [portfolio_id]
            ),
            value_cols=tuple(value_cols),
        )


class AsyncPortfolioReturnsReportAccessorImpl(  # noqa: D101
    AsyncPortfolioReturnsReportAccessor
):

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

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

    async def get_portolio_returns(  # noqa: D102
        self,
        portfolio_names: list[str] | None = None,
        start_date: DateLike | None = None,
        end_date: DateLike | None = None,
        aggregation_type: Literal["D", "M", "Y"] | None = "D",
        value_cols: list[str] | None = None,
    ) -> pl.DataFrame:
        available_value_cols = self._accessor.metric_cols
        if value_cols is not None:
            if set(value_cols) > set(available_value_cols):
                raise ValueError(
                    f"Unknown value_cols: "
                    f"{set(value_cols) - set(available_value_cols)}. "
                    f"Available: {available_value_cols}."
                )
        else:
            value_cols = available_value_cols
        return await self._accessor.get_data(
            path=[],
            expand=_get_portfolio_returns_expand(aggregation_type),
            filters=_get_portfolio_returns_filters(
                start_date, end_date, portfolio_names
            ),
            value_cols=tuple(value_cols) if value_cols is not None else (),
        )

    async def get_asset_returns(  # noqa: D102
        self,
        portfolio_id: str,
        start_date: DateLike | None = None,
        end_date: DateLike | None = None,
        aggregation_type: Literal["D", "M", "Y"] | None = "D",
        value_cols: list[str] | None = None,
    ) -> pl.DataFrame:
        available_value_cols = self._accessor.metric_cols
        if value_cols is not None:
            if set(value_cols) > set(available_value_cols):
                raise ValueError(
                    f"Unknown value_cols: "
                    f"{set(value_cols) - set(available_value_cols)}. "
                    f"Available: {available_value_cols}."
                )
        else:
            value_cols = available_value_cols
        return await self._accessor.get_data(
            path=[("portfolio_id", portfolio_id)],
            expand=tuple(
                list(
                    [
                        p
                        for p in _get_portfolio_returns_expand(aggregation_type)
                        if p != "portfolio_id"
                    ]
                    + [
                        "input_asset_id",
                        "input_asset_id_type",
                        "asset_id",
                        "asset_name",
                    ]
                )
            ),
            filters=_get_portfolio_returns_filters(
                start_date, end_date, [portfolio_id]
            ),
            value_cols=tuple(value_cols),
        )
