"""Multi-curve environment (`fuggers_py.curves.multicurve.environment`)."""

from __future__ import annotations

from dataclasses import dataclass, field

from fuggers_py.core.traits import YieldCurve
from fuggers_py.core.types import Currency

from ..errors import CurvesError
from .index import RateIndex


class MissingCurveError(CurvesError):
    """Raised when a requested curve is missing from the environment."""


@dataclass(frozen=True, slots=True)
class MultiCurveEnvironment:
    discount_curves: dict[Currency, YieldCurve]
    projection_curves: dict[RateIndex, YieldCurve]

    def discount_curve(self, currency: Currency) -> YieldCurve:
        try:
            return self.discount_curves[currency]
        except KeyError as exc:  # pragma: no cover - simple
            raise MissingCurveError(f"Missing discount curve for currency {currency}.") from exc

    def projection_curve(self, index: RateIndex) -> YieldCurve:
        try:
            return self.projection_curves[index]
        except KeyError as exc:  # pragma: no cover - simple
            raise MissingCurveError(f"Missing projection curve for index {index}.") from exc


@dataclass(slots=True)
class MultiCurveEnvironmentBuilder:
    _discount_curves: dict[Currency, YieldCurve] = field(default_factory=dict)
    _projection_curves: dict[RateIndex, YieldCurve] = field(default_factory=dict)

    def add_discount_curve(self, currency: Currency, curve: YieldCurve) -> "MultiCurveEnvironmentBuilder":
        self._discount_curves[currency] = curve
        return self

    def add_projection_curve(self, index: RateIndex, curve: YieldCurve) -> "MultiCurveEnvironmentBuilder":
        self._projection_curves[index] = curve
        return self

    def build(self) -> MultiCurveEnvironment:
        return MultiCurveEnvironment(
            discount_curves=dict(self._discount_curves),
            projection_curves=dict(self._projection_curves),
        )


__all__ = [
    "MissingCurveError",
    "MultiCurveEnvironment",
    "MultiCurveEnvironmentBuilder",
]
