"""Internal curve resolution and bump helpers for rates pricing."""

from __future__ import annotations

from dataclasses import dataclass
from decimal import Decimal

from fuggers_py.bonds.types import Tenor
from fuggers_py.core.daycounts import DayCountConvention
from fuggers_py.core.traits import YieldCurve
from fuggers_py.core.types import Currency, Date
from fuggers_py.curves.bumping import KeyRateBump, ParallelBump
from fuggers_py.curves.multicurve import MultiCurveEnvironment, RateIndex
from fuggers_py.data.pricing_specs import AnalyticsCurves


def _to_decimal(value: object) -> Decimal:
    if isinstance(value, Decimal):
        return value
    return Decimal(str(value))


def _curve_supports_discounting(curve: object | None) -> bool:
    return curve is not None and hasattr(curve, "discount_factor") and hasattr(curve, "zero_rate")


def _curve_supports_forward_projection(curve: object | None) -> bool:
    return curve is not None and hasattr(curve, "forward_rate")


def _curve_supports_projection(curve: object | None) -> bool:
    return _curve_supports_discounting(curve) or _curve_supports_forward_projection(curve)


@dataclass(frozen=True, slots=True)
class _ForwardProjectionWrapper:
    curve: YieldCurve

    def reference_date(self) -> Date:
        return self.curve.reference_date()

    def max_date(self) -> Date:
        return self.curve.max_date()

    def forward_rate(self, start_date: Date, end_date: Date) -> Decimal:
        return self.curve.forward_rate(start_date, end_date)


def _projection_keys(currency: Currency, index_name: str, index_tenor: Tenor | None) -> tuple[str, ...]:
    normalized_name = index_name.strip().upper()
    if index_tenor is None:
        return (
            normalized_name,
            f"{currency.code()}-{normalized_name}",
        )
    rate_index = RateIndex.new(normalized_name, index_tenor, currency)
    return (
        str(rate_index),
        rate_index.key(),
        normalized_name,
        f"{normalized_name}-{index_tenor}",
        f"{currency.code()}-{normalized_name}",
        f"{currency.code()}-{normalized_name}-{index_tenor}",
    )


def forward_rate_from_curve(
    curve: object,
    start_date: Date,
    end_date: Date,
    *,
    day_count_convention: DayCountConvention,
) -> Decimal:
    """Return the forward rate implied by a projection or discount curve.

    The rate is a raw decimal. When the curve only exposes discount factors,
    the forward is computed from the start and end discount factors using the
    supplied day-count convention.
    """

    if _curve_supports_forward_projection(curve) and not _curve_supports_discounting(curve):
        return _to_decimal(curve.forward_rate(start_date, end_date))
    tau = day_count_convention.to_day_count().year_fraction(start_date, end_date)
    if tau == Decimal(0):
        raise ValueError("forward_rate requires distinct start and end dates.")
    df_start = curve.discount_factor(start_date)
    df_end = curve.discount_factor(end_date)
    if df_end == Decimal(0):
        raise ValueError("forward_rate requires a non-zero end-date discount factor.")
    return (df_start / df_end - Decimal(1)) / tau


def resolve_discount_curve(curves: AnalyticsCurves, currency: Currency) -> YieldCurve:
    """Return the discount curve used for the given currency."""

    environment = curves.multicurve_environment
    if environment is not None and hasattr(environment, "discount_curve"):
        try:
            curve = environment.discount_curve(currency)
        except Exception:
            curve = None
        if _curve_supports_discounting(curve):
            return curve
    if _curve_supports_discounting(curves.discount_curve):
        return curves.discount_curve
    if _curve_supports_discounting(curves.collateral_curve):
        return curves.collateral_curve
    raise ValueError(f"Missing discount curve for currency {currency.code()}.")


def resolve_projection_curve(
    curves: AnalyticsCurves,
    *,
    currency: Currency,
    index_name: str,
    index_tenor: Tenor | None,
) -> YieldCurve:
    """Return the projection curve for the supplied rate index."""

    environment = curves.multicurve_environment
    if environment is not None and index_tenor is not None and hasattr(environment, "projection_curve"):
        try:
            curve = environment.projection_curve(RateIndex.new(index_name, index_tenor, currency))
        except Exception:
            curve = None
        if _curve_supports_projection(curve):
            return curve
    for key in _projection_keys(currency, index_name, index_tenor):
        curve = curves.projection_curves.get(key)
        if _curve_supports_projection(curve):
            return curve
    if _curve_supports_projection(curves.forward_curve):
        return curves.forward_curve
    if _curve_supports_discounting(curves.discount_curve):
        return curves.discount_curve
    raise ValueError(f"Missing projection curve for {currency.code()} {index_name}.")


def _parallel_bump_curve(curve: object | None, amount: float) -> object | None:
    if _curve_supports_discounting(curve):
        return ParallelBump(amount).apply(curve)
    inner_curve = getattr(curve, "curve", None)
    if _curve_supports_discounting(inner_curve):
        return _ForwardProjectionWrapper(ParallelBump(amount).apply(inner_curve))
    return curve


def _key_rate_bump_curve(curve: object | None, key_rate_bump: KeyRateBump) -> object | None:
    if _curve_supports_discounting(curve):
        return key_rate_bump.apply(curve)
    inner_curve = getattr(curve, "curve", None)
    if _curve_supports_discounting(inner_curve):
        return _ForwardProjectionWrapper(key_rate_bump.apply(inner_curve))
    return curve


def analytics_curves_with_parallel_bump(
    curves: AnalyticsCurves,
    *,
    currency: Currency,
    projection_indices: tuple[RateIndex, ...],
    bump: Decimal,
) -> AnalyticsCurves:
    """Return a copy of ``curves`` with a parallel rate bump applied.

    The bump is applied to the discount curve for ``currency`` and to any
    matching projection curves. The bump is expressed as a raw decimal shift.
    """

    amount = float(bump)
    environment = curves.multicurve_environment
    bumped_environment = environment
    if environment is not None and hasattr(environment, "discount_curves") and hasattr(environment, "projection_curves"):
        discount_curves = dict(environment.discount_curves)
        if currency in discount_curves:
            discount_curves[currency] = _parallel_bump_curve(discount_curves[currency], amount)
        projection_curves = dict(environment.projection_curves)
        for index in projection_indices:
            if index in projection_curves:
                projection_curves[index] = _parallel_bump_curve(projection_curves[index], amount)
        bumped_environment = MultiCurveEnvironment(
            discount_curves=discount_curves,
            projection_curves=projection_curves,
        )
    bumped_projection_curves = {
        key: _parallel_bump_curve(value, amount)
        for key, value in curves.projection_curves.items()
    }
    return AnalyticsCurves(
        discount_curve=_parallel_bump_curve(curves.discount_curve, amount),
        forward_curve=_parallel_bump_curve(curves.forward_curve, amount),
        government_curve=curves.government_curve,
        benchmark_curve=curves.benchmark_curve,
        credit_curve=curves.credit_curve,
        repo_curve=curves.repo_curve,
        collateral_curve=_parallel_bump_curve(curves.collateral_curve, amount),
        fx_forward_curve=curves.fx_forward_curve,
        multicurve_environment=bumped_environment,
        projection_curves=bumped_projection_curves,
        vol_surface=curves.vol_surface,
    )


def analytics_curves_with_key_rate_bump(
    curves: AnalyticsCurves,
    *,
    currency: Currency,
    projection_indices: tuple[RateIndex, ...],
    tenor_grid: tuple[Tenor, ...],
    key_tenor: Tenor,
    bump: Decimal,
) -> AnalyticsCurves:
    """Return a copy of ``curves`` with a key-rate bump applied.

    The bump is applied at ``key_tenor`` using the supplied tenor grid. The
    amount is a raw decimal shift.
    """

    key_rate_bump = KeyRateBump(list(tenor_grid), key_tenor=key_tenor, bump=float(bump))
    environment = curves.multicurve_environment
    bumped_environment = environment
    if environment is not None and hasattr(environment, "discount_curves") and hasattr(environment, "projection_curves"):
        discount_curves = dict(environment.discount_curves)
        if currency in discount_curves:
            discount_curves[currency] = _key_rate_bump_curve(discount_curves[currency], key_rate_bump)
        projection_curves = dict(environment.projection_curves)
        for index in projection_indices:
            if index in projection_curves:
                projection_curves[index] = _key_rate_bump_curve(projection_curves[index], key_rate_bump)
        bumped_environment = MultiCurveEnvironment(
            discount_curves=discount_curves,
            projection_curves=projection_curves,
        )
    bumped_projection_curves = {
        key: _key_rate_bump_curve(value, key_rate_bump)
        for key, value in curves.projection_curves.items()
    }
    return AnalyticsCurves(
        discount_curve=_key_rate_bump_curve(curves.discount_curve, key_rate_bump),
        forward_curve=_key_rate_bump_curve(curves.forward_curve, key_rate_bump),
        government_curve=curves.government_curve,
        benchmark_curve=curves.benchmark_curve,
        credit_curve=curves.credit_curve,
        repo_curve=curves.repo_curve,
        collateral_curve=_key_rate_bump_curve(curves.collateral_curve, key_rate_bump),
        fx_forward_curve=curves.fx_forward_curve,
        multicurve_environment=bumped_environment,
        projection_curves=bumped_projection_curves,
        vol_surface=curves.vol_surface,
    )


def curve_zero_rate(curve: YieldCurve, date: Date) -> Decimal:
    """Return the curve zero rate as a raw decimal."""

    resolved = curve.zero_rate(date)
    if hasattr(resolved, "value"):
        return _to_decimal(resolved.value())
    return _to_decimal(resolved)


__all__ = [
    "analytics_curves_with_key_rate_bump",
    "analytics_curves_with_parallel_bump",
    "curve_zero_rate",
    "forward_rate_from_curve",
    "resolve_discount_curve",
    "resolve_projection_curve",
]
