import abc
from typing import Annotated, Union, overload

from pydantic import Discriminator

from bayesline.api._src._utils import docstrings_from_sync
from bayesline.api._src.equity.report.accessor import (
    AsyncTypedReportAccessorApi,
    TypedReportAccessorApi,
)
from bayesline.api._src.equity.report.api import (
    AsyncReportApi,
    ReportApi,
    S,
)
from bayesline.api._src.equity.report.covariance import (
    AsyncFactorCovarianceReportApi,
    FactorCovarianceReportApi,
    FactorCovarianceReportSettings,
)
from bayesline.api._src.equity.report.exposure import (
    AsyncExposureReportApi,
    ExposureReportApi,
    ExposureReportSettings,
)
from bayesline.api._src.equity.report.factor_stats import (
    AsyncFactorStatsReportApi,
    FactorStatsReportApi,
    FactorStatsReportSettings,
)
from bayesline.api._src.equity.report.idio import (
    AsyncIdioReportApi,
    IdioReportApi,
    IdioReportSettings,
)
from bayesline.api._src.equity.report.mcap_exposure import (
    AsyncMcapExposureReportApi,
    McapExposureReportApi,
    McapExposureReportSettings,
)
from bayesline.api._src.equity.report.portfolio_holdings import (
    AsyncPortfolioHoldingsReportApi,
    PortfolioHoldingsReportApi,
    PortfolioHoldingsReportSettings,
)
from bayesline.api._src.equity.report.portfolio_returns import (
    AsyncPortfolioReturnsReportApi,
    PortfolioReturnsReportApi,
    PortfolioReturnsReportSettings,
)
from bayesline.api._src.equity.report.return_attribution import (
    AsyncReturnAttributionReportApi,
    ReturnAttributionReportApi,
    ReturnAttributionReportSettings,
)
from bayesline.api._src.equity.report.universe import (
    AsyncUniverseCountReportApi,
    UniverseCountReportApi,
    UniverseCountReportSettings,
)
from bayesline.api._src.equity.report.weights import (
    AsyncWeightsReportApi,
    WeightsReportApi,
    WeightsReportSettings,
)
from bayesline.api._src.equity.report.xsr import (
    AsyncXSRReportApi,
    XSRReportApi,
    XSRReportSettings,
)

# Discriminated union of all report settings for polymorphic deserialization
AnyReportSettings = Annotated[
    Union[
        UniverseCountReportSettings,
        ExposureReportSettings,
        FactorStatsReportSettings,
        FactorCovarianceReportSettings,
        PortfolioHoldingsReportSettings,
        PortfolioReturnsReportSettings,
        ReturnAttributionReportSettings,
        XSRReportSettings,
        IdioReportSettings,
        McapExposureReportSettings,
        WeightsReportSettings,
    ],
    Discriminator("report_type"),
]


class ReportLoaderApi(abc.ABC):
    """The main interface for loading different types of reports."""

    @overload
    def load(self, settings: UniverseCountReportSettings) -> UniverseCountReportApi: ...

    @overload
    def load(self, settings: ExposureReportSettings) -> ExposureReportApi: ...

    @overload
    def load(self, settings: FactorStatsReportSettings) -> FactorStatsReportApi: ...

    @overload
    def load(
        self, settings: FactorCovarianceReportSettings
    ) -> FactorCovarianceReportApi: ...

    @overload
    def load(
        self, settings: PortfolioHoldingsReportSettings
    ) -> PortfolioHoldingsReportApi: ...

    @overload
    def load(
        self, settings: PortfolioReturnsReportSettings
    ) -> PortfolioReturnsReportApi: ...

    @overload
    def load(
        self, settings: ReturnAttributionReportSettings
    ) -> ReturnAttributionReportApi: ...

    @overload
    def load(self, settings: IdioReportSettings) -> IdioReportApi: ...

    @overload
    def load(self, settings: XSRReportSettings) -> XSRReportApi: ...

    @overload
    def load(self, settings: McapExposureReportSettings) -> McapExposureReportApi: ...

    @overload
    def load(self, settings: WeightsReportSettings) -> WeightsReportApi: ...

    @abc.abstractmethod
    def load(self, settings: S) -> ReportApi[..., TypedReportAccessorApi, S]:
        """
        Load a report using the specified settings.

        Parameters
        ----------
        settings: S
            The settings to use to load the report.

        Returns
        -------
        ReportApi[..., TypedReportAccessorApi, S]
            The loaded report API.
        """

    @abc.abstractmethod
    def load_accessor(self, identifier: str) -> TypedReportAccessorApi:
        """
        Load a report accessor from an identifier.

        This accessor is obtained from an existing accessor after a report has been
        calculated. The identifier is ephemeral, and should only be used for short-
        lived sessions, and not persisted over time.

        Parameters
        ----------
        identifier: str
            The short-lived ephemeral identifier to load the report accessor from.

        Returns
        -------
        TypedReportAccessorApi
            The loaded report accessor.

        Raises
        ------
        KeyError
            If the identifier is not found. This can happen if the identifier has
            expired, belongs to a different session, or was never created.
        """


@docstrings_from_sync
class AsyncReportLoaderApi(abc.ABC):

    @overload
    async def load(  # noqa: D102
        self, settings: UniverseCountReportSettings
    ) -> AsyncUniverseCountReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: ExposureReportSettings
    ) -> AsyncExposureReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: FactorStatsReportSettings
    ) -> AsyncFactorStatsReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: FactorCovarianceReportSettings
    ) -> AsyncFactorCovarianceReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: PortfolioHoldingsReportSettings
    ) -> AsyncPortfolioHoldingsReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: PortfolioReturnsReportSettings
    ) -> AsyncPortfolioReturnsReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: ReturnAttributionReportSettings
    ) -> AsyncReturnAttributionReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: IdioReportSettings
    ) -> AsyncIdioReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: XSRReportSettings
    ) -> AsyncXSRReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: McapExposureReportSettings
    ) -> AsyncMcapExposureReportApi: ...

    @overload
    async def load(  # noqa: D102
        self, settings: WeightsReportSettings
    ) -> AsyncWeightsReportApi: ...

    @abc.abstractmethod
    async def load(  # noqa: D102
        self, settings: S
    ) -> AsyncReportApi[..., AsyncTypedReportAccessorApi, S]: ...

    @abc.abstractmethod
    async def load_accessor(  # noqa: D102
        self, identifier: str
    ) -> AsyncTypedReportAccessorApi: ...
