from __future__ import annotations

from dataclasses import dataclass

from backend.data.contracts import DataSource, DataUsage


class DataPolicyError(RuntimeError):
    pass


@dataclass(frozen=True, slots=True)
class DataSourcePolicy:
    calculation_sources: frozenset[DataSource] = frozenset(
        {DataSource.TUSHARE, DataSource.IFIND, DataSource.LOCAL}
    )
    display_sources: frozenset[DataSource] = frozenset(DataSource)

    def assert_allowed(self, source: DataSource, usage: DataUsage) -> None:
        allowed = (
            self.calculation_sources if usage is DataUsage.CALCULATION else self.display_sources
        )
        if source not in allowed:
            raise DataPolicyError(f"{source.value} cannot be used for {usage.value}")

    def candidates(self, dataset: str, usage: DataUsage) -> tuple[DataSource, ...]:
        routes = {
            ("calendar", DataUsage.CALCULATION): (DataSource.TUSHARE,),
            ("entities", DataUsage.CALCULATION): (DataSource.TUSHARE,),
            ("daily_chart", DataUsage.DISPLAY): (DataSource.IFIND, DataSource.TUSHARE),
            ("minute_chart", DataUsage.DISPLAY): (DataSource.IFIND, DataSource.EASTMONEY),
            ("realtime_quote", DataUsage.CALCULATION): (DataSource.IFIND, DataSource.TUSHARE),
            ("market_insight", DataUsage.CALCULATION): (DataSource.TUSHARE,),
            ("dynamic_auction", DataUsage.CALCULATION): (DataSource.IFIND,),
            ("screener_factors", DataUsage.CALCULATION): (DataSource.TUSHARE,),
        }
        return routes.get((dataset, usage), ())