from __future__ import annotations

import json
from pathlib import Path

from backend.bootstrap.config import APP_DIR
from backend.data.contracts import DataUsage, DatasetContract, ProviderContract


class DataPolicyError(RuntimeError):
    pass


class DataSourcePolicy:
    def __init__(
        self,
        providers: dict[str, ProviderContract],
        datasets: dict[str, DatasetContract],
    ) -> None:
        self.providers = dict(providers)
        self.datasets = dict(datasets)

    @classmethod
    def load(cls, path: Path | None = None) -> "DataSourcePolicy":
        config_path = path or APP_DIR / "config" / "data-fields.config.json"
        payload = json.loads(config_path.read_text(encoding="utf-8"))
        providers = {
            provider_id: ProviderContract(
                id=provider_id,
                provider_class=str(item["class"]),
                calculation_allowed=bool(item["calculation_allowed"]),
            )
            for provider_id, item in payload["providers"].items()
        }
        datasets = {
            item["id"]: DatasetContract(
                id=str(item["id"]),
                entity=str(item["entity"]),
                frequency=str(item["frequency"]),
                primary=str(item["primary"]),
                fallbacks=tuple(str(value) for value in item.get("fallbacks", [])),
                usage=str(item["usage"]),
                fields=tuple(str(value) for value in item.get("fields", [])),
            )
            for item in payload["datasets"]
        }
        return cls(providers, datasets)

    def dataset(self, dataset_id: str) -> DatasetContract:
        try:
            return self.datasets[dataset_id]
        except KeyError as exc:
            raise DataPolicyError(f"Unregistered dataset: {dataset_id}") from exc

    def assert_allowed(
        self,
        dataset_id: str,
        provider_id: str,
        usage: DataUsage,
    ) -> DatasetContract:
        dataset = self.dataset(dataset_id)
        if dataset.usage == "blocked":
            raise DataPolicyError(f"Dataset is blocked: {dataset_id}")
        if provider_id not in dataset.providers:
            raise DataPolicyError(
                f"Provider {provider_id} is not registered for dataset {dataset_id}"
            )
        try:
            provider = self.providers[provider_id]
        except KeyError as exc:
            raise DataPolicyError(f"Unregistered provider: {provider_id}") from exc
        if usage == "calculation":
            if dataset.usage != "calculation" or not provider.calculation_allowed:
                raise DataPolicyError(
                    f"Provider {provider_id} cannot calculate dataset {dataset_id}"
                )
        return dataset
