from __future__ import annotations

from decimal import Decimal

from fuggers_py.core import Currency, Date, Frequency
from fuggers_py.engine import RatesPricingRouter, RoutedFraPricingResult
from fuggers_py.data import BasisSwapQuoteOutput, SwapQuoteOutput
from fuggers_py.rates import BasisSwap, FixedFloatSwap, FixedLegSpec, FloatingLegSpec, Fra, Ois, PayReceive, ScheduleDefinition

from ._rates_helpers import flat_curve, multicurve_analytics_curves, rate_index


def _swap(reference_date: Date, *, cls=FixedFloatSwap):
    return cls(
        effective_date=reference_date,
        maturity_date=reference_date.add_years(2),
        fixed_leg=FixedLegSpec(
            pay_receive=PayReceive.RECEIVE,
            notional=Decimal("1000000"),
            fixed_rate=Decimal("0.04"),
            currency=Currency.USD,
            schedule=ScheduleDefinition(frequency=Frequency.SEMI_ANNUAL),
        ),
        floating_leg=FloatingLegSpec(
            pay_receive=PayReceive.PAY,
            notional=Decimal("1000000"),
            index_name="SOFR",
            index_tenor="3M",
            currency=Currency.USD,
            schedule=ScheduleDefinition(frequency=Frequency.QUARTERLY),
        ),
    )


def test_rates_pricing_router_dispatches_swap_fra_and_basis_swap() -> None:
    reference_date = Date.from_ymd(2026, 1, 2)
    discount_curve = flat_curve(reference_date, "0.03")
    sofr_curve = flat_curve(reference_date, "0.035")
    term_curve = flat_curve(reference_date, "0.0375")
    curves = multicurve_analytics_curves(
        discount_curve=discount_curve,
        discount_currency=Currency.USD,
        forward_curve=sofr_curve,
        projection_curves={
            rate_index("SOFR", "3M"): sofr_curve,
            rate_index("TERM", "6M"): term_curve,
        },
    )
    router = RatesPricingRouter()
    swap = _swap(reference_date)
    fra = Fra(
        start_date=reference_date.add_months(6),
        end_date=reference_date.add_months(9),
        notional=Decimal("1000000"),
        fixed_rate=Decimal("0.034"),
        pay_receive=PayReceive.RECEIVE,
        currency=Currency.USD,
        index_name="SOFR",
        index_tenor="3M",
        instrument_id="FRA-1",
    )
    basis_swap = BasisSwap(
        effective_date=reference_date,
        maturity_date=reference_date.add_years(3),
        pay_leg=FloatingLegSpec(
            pay_receive=PayReceive.PAY,
            notional=Decimal("1000000"),
            index_name="SOFR",
            index_tenor="3M",
            currency=Currency.USD,
            schedule=ScheduleDefinition(frequency=Frequency.QUARTERLY),
        ),
        receive_leg=FloatingLegSpec(
            pay_receive=PayReceive.RECEIVE,
            notional=Decimal("1000000"),
            index_name="TERM",
            index_tenor="6M",
            currency=Currency.USD,
            schedule=ScheduleDefinition(frequency=Frequency.SEMI_ANNUAL),
        ),
        instrument_id="BASIS-1",
    )

    swap_result = router.price(swap, curves=curves)
    fra_result = router.price(fra, curves=curves)
    basis_result = router.price(basis_swap, curves=curves)

    assert isinstance(swap_result, SwapQuoteOutput)
    assert swap_result.pricing_path == "swap"
    assert swap_result.pv01 is not None
    assert isinstance(fra_result, RoutedFraPricingResult)
    assert fra_result.pricing_path == "fra"
    assert fra_result.instrument_id is not None
    assert isinstance(basis_result, BasisSwapQuoteOutput)
    assert basis_result.pricing_path == "basis_swap"
    assert basis_result.pv01 is not None


def test_rates_pricing_router_prices_ois_and_batch() -> None:
    reference_date = Date.from_ymd(2026, 1, 2)
    curve = flat_curve(reference_date, "0.03")
    curves = multicurve_analytics_curves(
        discount_curve=curve,
        discount_currency=Currency.USD,
        forward_curve=curve,
        projection_curves={rate_index("SOFR", "3M"): curve},
    )
    router = RatesPricingRouter()
    ois = _swap(reference_date, cls=Ois)
    fra = Fra(
        start_date=reference_date.add_months(6),
        end_date=reference_date.add_months(9),
        notional=Decimal("1000000"),
        fixed_rate=Decimal("0.03"),
        pay_receive=PayReceive.RECEIVE,
        currency=Currency.USD,
        index_name="SOFR",
        index_tenor="3M",
    )

    results = router.price_batch((ois, fra), curves=curves)

    assert isinstance(results[0], SwapQuoteOutput)
    assert results[0].pricing_path == "ois"
    assert isinstance(results[1], RoutedFraPricingResult)
