"""
Tests for Adaptive Smart Budget Phase 3: Forecast-Informed Routing.

Covers:
- Per-request routing uses forecast_p_max from BudgetPacingState cache
- Blending logic: max(velocity_p_max, forecast_p_max) prevents unnecessary downgrades
- Cache hit path: locked tier uses effective_p_max directly
- Cache miss path: falls back to velocity-only calculation
- Cache miss with forecast: blends forecast signal when available
- Weekend lull scenario: forecast prevents premature downgrade
- Forecast poll writes forecast_p_max to cache (analytics path)
- Forecast_p_max propagation through apply_transition
- Edge cases: None forecast, inf forecast, expired cache
"""

import pytest
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from unittest.mock import MagicMock, patch, PropertyMock

from api.services.smart_budget_state import (
    BudgetPacingCache,
    BudgetPacingState,
    ModelTier,
    TransitionDecision,
    apply_transition,
    compute_model_tiers,
    evaluate_transition,
    find_tier_for_price,
    get_pacing_cache,
    CACHE_TTL_SECONDS,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

SAMPLE_MODELS = [
    {"model_id": "flash", "model_name": "Flash", "price": 0.50, "hle": 20.0},
    {"model_id": "haiku", "model_name": "Haiku", "price": 0.80, "hle": 28.0},
    {"model_id": "sonnet", "model_name": "Sonnet", "price": 3.00, "hle": 38.0},
    {"model_id": "pro", "model_name": "Pro", "price": 5.00, "hle": 44.7},
    {"model_id": "opus", "model_name": "Opus", "price": 12.00, "hle": 50.0},
]


def _make_tiers() -> list[ModelTier]:
    return compute_model_tiers(SAMPLE_MODELS)


def _make_state(
    sub_id=1,
    p_max=5.0,
    tier_name="balanced",
    tier_min=3.0,
    tier_max=5.0,
    forecast_p_max=None,
    velocity_ema=200_000.0,
    transition_dir=None,
    transition_started=None,
    last_transition_at=None,
) -> BudgetPacingState:
    now = datetime.now(timezone.utc)
    return BudgetPacingState(
        subscription_id=sub_id,
        effective_p_max=p_max,
        locked_model_tier=tier_name,
        locked_tier_min_price=tier_min,
        locked_tier_max_price=tier_max,
        locked_at=now,
        lock_expires_at=None,
        velocity_ema=velocity_ema,
        last_forecast_p_max=forecast_p_max,
        transition_direction=transition_dir,
        transition_started_at=transition_started,
        transition_reason=None,
        last_transition_at=last_transition_at,
    )


def _make_budget(
    *,
    budget_id=1,
    subscription_id=1,
    scope_type="global",
    limit_amount=100.0,
    period="monthly",
    tz="UTC",
    enforcement="alert",
    pacing_mode="linear",
    auto_optimize=True,
    is_active=True,
):
    b = MagicMock()
    b.id = budget_id
    b.subscription_id = subscription_id
    b.scope_type = scope_type
    b.tag_key = None
    b.tag_value = None
    b.include_untagged = False
    b.limit_amount = limit_amount
    b.period = period
    b.timezone = tz
    b.enforcement = enforcement
    b.pacing_mode = pacing_mode
    b.auto_optimize = auto_optimize
    b.is_active = is_active
    b.period_start = None
    return b


def _make_subscription(sub_id=1):
    sub = MagicMock()
    sub.id = sub_id
    sub.plan = MagicMock()
    sub.plan.price_monthly = 0
    sub.plan.tokens_per_month = 0
    return sub


# ===========================================================================
# Forecast P_max blending in per-request routing
# ===========================================================================

class TestForecastBlendingLogic:
    """
    The core Phase 3 invariant: per-request routing uses
    max(velocity_p_max, forecast_p_max) so the forecast can prevent
    unnecessary downgrades during predicted lulls.
    """

    def test_forecast_higher_than_velocity_uses_forecast(self):
        """
        When the forecast predicts lower future spend (e.g. weekend lull),
        forecast_p_max is higher than velocity_p_max.  The system should
        use the more permissive forecast value.
        """
        velocity_p_max = 3.50
        forecast_p_max = 6.20

        effective = max(velocity_p_max, forecast_p_max)

        assert effective == forecast_p_max
        assert effective > velocity_p_max

    def test_velocity_higher_than_forecast_uses_velocity(self):
        """
        When velocity is more optimistic (e.g. traffic dropped but forecast
        still projects high spend), use velocity.
        """
        velocity_p_max = 8.00
        forecast_p_max = 4.50

        effective = max(velocity_p_max, forecast_p_max)

        assert effective == velocity_p_max

    def test_both_equal_returns_same(self):
        p_max = 5.00
        assert max(p_max, p_max) == p_max


# ===========================================================================
# Cache hit path: locked tier
# ===========================================================================

class TestCacheHitPath:
    """When a locked BudgetPacingState exists in cache, routing uses it directly."""

    def test_locked_tier_uses_effective_p_max(self):
        """Cached state with a locked tier should be used as-is."""
        cache = BudgetPacingCache(ttl=60)
        state = _make_state(sub_id=42, p_max=5.45, tier_name="balanced", forecast_p_max=6.20)
        cache.put(state)

        cached = cache.get(42)
        assert cached is not None
        assert cached.locked_model_tier == "balanced"
        assert cached.effective_p_max == 5.45

    def test_locked_tier_exposes_forecast_p_max(self):
        """Cached state should carry the forecast_p_max for logging/debugging."""
        cache = BudgetPacingCache(ttl=60)
        state = _make_state(sub_id=42, forecast_p_max=7.30)
        cache.put(state)

        cached = cache.get(42)
        assert cached.last_forecast_p_max == 7.30


# ===========================================================================
# Cache miss path: forecast-informed fallback
# ===========================================================================

class TestCacheMissWithForecast:
    """
    When the tier lock has expired but a BudgetPacingState still exists
    (without a locked tier), routing should blend forecast_p_max with
    the live velocity calculation.
    """

    def test_cache_miss_no_state_uses_velocity_only(self):
        """With no cached state at all, routing falls back to pure velocity."""
        cache = BudgetPacingCache(ttl=60)
        cached = cache.get(999)
        assert cached is None

    def test_cache_state_without_lock_but_with_forecast(self):
        """
        State exists (forecast_p_max populated) but tier is not locked.
        The routing path should read last_forecast_p_max for blending.
        """
        state = _make_state(sub_id=1, tier_name=None, forecast_p_max=8.50)
        assert state.locked_model_tier is None
        assert state.last_forecast_p_max == 8.50

    def test_forecast_none_means_velocity_only(self):
        """When forecast_p_max is None, routing should use velocity_p_max alone."""
        state = _make_state(sub_id=1, tier_name=None, forecast_p_max=None)
        assert state.last_forecast_p_max is None


# ===========================================================================
# Weekend lull scenario (design doc Section 6)
# ===========================================================================

class TestWeekendLullScenario:
    """
    Simulates the weekend lull scenario from the design doc:
    - Friday evening: velocity drops (fewer requests)
    - Velocity-only P_max would drop → premature downgrade
    - Forecast predicts weekend lull → forecast_p_max stays high
    - System should NOT downgrade
    """

    def test_forecast_prevents_friday_evening_downgrade(self):
        """
        Velocity-based P_max drops to $2.50 (would downgrade from balanced).
        Forecast-based P_max is $6.00 (knows weekend is coming, spend will be low).
        Effective P_max should be $6.00 — no downgrade.
        """
        tiers = _make_tiers()
        balanced_tier = None
        for t in tiers:
            if t.min_price <= 3.0 <= t.max_price:
                balanced_tier = t
                break
        assert balanced_tier is not None, "Test requires a tier containing $3.00"

        velocity_p_max = 2.50
        forecast_p_max = 6.00

        effective_p_max = max(velocity_p_max, forecast_p_max)

        assert effective_p_max == 6.00
        assert effective_p_max >= balanced_tier.min_price, (
            "Effective P_max should keep the user in the balanced tier"
        )

    def test_weekend_lull_with_hysteresis_holds_tier(self):
        """
        Full integration: a locked balanced tier should hold through a
        weekend where velocity drops, because the forecast prevents the
        P_max from falling below the downgrade threshold.
        """
        tiers = _make_tiers()
        balanced = None
        for t in tiers:
            if "balanced" in t.name:
                balanced = t
                break
        if balanced is None:
            balanced = tiers[1] if len(tiers) > 1 else tiers[0]

        now = datetime.now(timezone.utc)
        state = _make_state(
            tier_name=balanced.name,
            tier_min=balanced.min_price,
            tier_max=balanced.max_price,
            p_max=5.00,
            forecast_p_max=6.00,
        )

        # Simulate: velocity-only P_max drops to $2.50 on Friday evening,
        # but forecast says $6.00.  The effective P_max fed to the
        # transition engine should be max(2.50, 6.00) = 6.00.
        effective_p_max = max(2.50, 6.00)

        decision = evaluate_transition(
            current_state=state,
            raw_p_max=effective_p_max,
            tiers=tiers,
            remaining_budget=60.0,
            budget_limit=100.0,
            days_remaining=20.0,
            projected_spend=40.0,
            velocity_ema=200_000.0,
            now=now,
        )

        assert decision.action == "hold", (
            f"Expected hold (no downgrade), got {decision.action}: {decision.reason}"
        )


# ===========================================================================
# Forecast P_max propagation through apply_transition
# ===========================================================================

class TestForecastPropagation:
    """Verify that forecast_p_max flows through the state machine correctly."""

    def test_apply_transition_stores_forecast_p_max(self):
        """apply_transition should persist forecast_p_max in the new state."""
        tiers = _make_tiers()
        decision = TransitionDecision("hold", tiers[0].name, "stable", 5.0)

        state = apply_transition(
            current_state=None,
            decision=decision,
            subscription_id=1,
            raw_p_max=5.0,
            velocity_ema=200_000.0,
            forecast_p_max=7.50,
            tiers=tiers,
        )

        assert state.last_forecast_p_max == 7.50

    def test_apply_transition_none_forecast(self):
        """When forecast is unavailable, last_forecast_p_max should be None."""
        tiers = _make_tiers()
        decision = TransitionDecision("hold", tiers[0].name, "stable", 5.0)

        state = apply_transition(
            current_state=None,
            decision=decision,
            subscription_id=1,
            raw_p_max=5.0,
            velocity_ema=200_000.0,
            forecast_p_max=None,
            tiers=tiers,
        )

        assert state.last_forecast_p_max is None

    def test_upgrade_preserves_forecast_p_max(self):
        """After an upgrade transition, the new state should carry forecast_p_max."""
        tiers = _make_tiers()
        target = tiers[1]
        decision = TransitionDecision("upgrade", target.name, "sustained", 8.0)

        state = apply_transition(
            current_state=None,
            decision=decision,
            subscription_id=1,
            raw_p_max=8.0,
            velocity_ema=300_000.0,
            forecast_p_max=9.25,
            tiers=tiers,
        )

        assert state.locked_model_tier == target.name
        assert state.last_forecast_p_max == 9.25

    def test_hold_updates_forecast_p_max(self):
        """On a hold decision, forecast_p_max should be updated to the latest value."""
        tiers = _make_tiers()
        old_state = _make_state(
            tier_name=tiers[0].name,
            tier_min=tiers[0].min_price,
            tier_max=tiers[0].max_price,
            forecast_p_max=4.00,
        )
        decision = TransitionDecision("hold", tiers[0].name, "stable", old_state.effective_p_max)

        new_state = apply_transition(
            current_state=old_state,
            decision=decision,
            subscription_id=1,
            raw_p_max=old_state.effective_p_max,
            velocity_ema=200_000.0,
            forecast_p_max=6.50,
            tiers=tiers,
        )

        assert new_state.last_forecast_p_max == 6.50, (
            "Forecast P_max should be updated on every poll cycle"
        )


# ===========================================================================
# Cache write from analytics poll (Phase 3 enhancement)
# ===========================================================================

class TestAnalyticsPollCacheWrite:
    """
    The analytics forecast poll (every 30s) should always write forecast_p_max
    to the BudgetPacingState cache, even when the full hysteresis engine runs.
    """

    def test_state_carries_forecast_after_transition_evaluation(self):
        """
        Simulates the analytics poll path: evaluate_transition + apply_transition
        should produce a state with last_forecast_p_max set.
        """
        tiers = _make_tiers()
        now = datetime.now(timezone.utc)
        economy = tiers[0]

        state = _make_state(
            tier_name=economy.name,
            tier_min=economy.min_price,
            tier_max=economy.max_price,
        )

        decision = evaluate_transition(
            current_state=state,
            raw_p_max=3.0,
            tiers=tiers,
            remaining_budget=70.0,
            budget_limit=100.0,
            days_remaining=20.0,
            projected_spend=30.0,
            velocity_ema=200_000.0,
            now=now,
        )

        new_state = apply_transition(
            current_state=state,
            decision=decision,
            subscription_id=1,
            raw_p_max=3.0,
            velocity_ema=200_000.0,
            forecast_p_max=5.80,
            tiers=tiers,
            now=now,
        )

        assert new_state.last_forecast_p_max == 5.80

    def test_cache_round_trip_preserves_forecast(self):
        """Write state to cache, read it back — forecast_p_max should survive."""
        cache = BudgetPacingCache(ttl=60)
        state = _make_state(sub_id=7, forecast_p_max=4.25)
        cache.put(state)

        retrieved = cache.get(7)
        assert retrieved is not None
        assert retrieved.last_forecast_p_max == 4.25

    def test_successive_polls_update_forecast(self):
        """Each poll cycle should update forecast_p_max in the cache."""
        cache = BudgetPacingCache(ttl=60)
        tiers = _make_tiers()
        now = datetime.now(timezone.utc)

        # First poll: forecast = $4.00
        state1 = _make_state(sub_id=1, forecast_p_max=4.00)
        cache.put(state1)
        assert cache.get(1).last_forecast_p_max == 4.00

        # Second poll: forecast = $6.50 (weekend prediction)
        state2 = _make_state(sub_id=1, forecast_p_max=6.50)
        cache.put(state2)
        assert cache.get(1).last_forecast_p_max == 6.50


# ===========================================================================
# End-to-end: forecast-informed routing prevents unnecessary downgrade
# ===========================================================================

class TestForecastInformedRoutingE2E:
    """
    End-to-end scenario simulating the full Phase 3 flow:
    1. Analytics poll computes forecast_p_max and writes to cache
    2. Per-request routing reads cache and uses forecast-informed P_max
    3. System avoids unnecessary downgrade during a predicted lull
    """

    def test_full_flow_prevents_downgrade_during_lull(self):
        """
        Scenario:
        - User is on balanced tier ($3-5 models)
        - Velocity spikes → velocity_p_max drops to $2.00 (would trigger downgrade)
        - Forecast knows this is temporary → forecast_p_max = $5.50
        - Per-request routing should use $5.50, keeping the user on balanced
        """
        tiers = _make_tiers()
        cache = BudgetPacingCache(ttl=60)
        now = datetime.now(timezone.utc)

        balanced = None
        for t in tiers:
            if t.min_price <= 3.0 <= t.max_price:
                balanced = t
                break
        assert balanced is not None

        # Step 1: Analytics poll writes state with forecast_p_max
        state = _make_state(
            sub_id=1,
            p_max=5.50,
            tier_name=balanced.name,
            tier_min=balanced.min_price,
            tier_max=balanced.max_price,
            forecast_p_max=5.50,
            velocity_ema=400_000.0,
        )
        cache.put(state)

        # Step 2: Per-request routing reads cache
        cached = cache.get(1)
        assert cached is not None
        assert cached.locked_model_tier == balanced.name

        # The routing path uses effective_p_max from the locked state
        assert cached.effective_p_max == 5.50

        # Step 3: Even if velocity_p_max computed live would be $2.00,
        # the locked state holds at $5.50
        velocity_p_max_live = 2.00
        effective = cached.effective_p_max  # Uses locked value, not live
        assert effective == 5.50
        assert effective >= balanced.min_price

    def test_cache_miss_blends_forecast_with_velocity(self):
        """
        When tier lock is absent but forecast_p_max exists in cache,
        routing should blend: max(velocity_p_max, forecast_p_max).
        """
        cache = BudgetPacingCache(ttl=60)

        # State exists but without a locked tier (e.g., just initialized)
        state = _make_state(
            sub_id=1,
            tier_name=None,
            forecast_p_max=7.00,
        )
        cache.put(state)

        cached = cache.get(1)
        assert cached is not None
        assert cached.locked_model_tier is None

        # Routing falls back to live calculation but can read forecast
        velocity_p_max = 4.50
        forecast_from_cache = cached.last_forecast_p_max

        effective = max(velocity_p_max, forecast_from_cache)
        assert effective == 7.00

    def test_expired_cache_falls_back_to_velocity(self):
        """When cache entry has expired, routing uses velocity-only."""
        import time
        cache = BudgetPacingCache(ttl=0.05)
        state = _make_state(sub_id=1, forecast_p_max=8.00)
        cache.put(state)

        time.sleep(0.06)
        cached = cache.get(1)
        assert cached is None


# ===========================================================================
# Forecast-routing alignment validation (design doc Section 10)
# ===========================================================================

class TestForecastRoutingAlignment:
    """
    Success criteria from Section 10: per-request routing should read
    cached forecast P_max, ensuring forecast-routing alignment.
    """

    def test_forecast_p_max_matches_what_analytics_wrote(self):
        """
        The forecast_p_max that analytics writes to the cache should be
        exactly what per-request routing reads back.
        """
        cache = BudgetPacingCache(ttl=60)
        tiers = _make_tiers()
        now = datetime.now(timezone.utc)

        # Simulate analytics poll writing state
        forecast_value = 6.78
        state = _make_state(sub_id=42, forecast_p_max=forecast_value)
        cache.put(state)

        # Simulate per-request routing reading
        cached = cache.get(42)
        assert cached.last_forecast_p_max == forecast_value

    def test_velocity_and_forecast_divergence_uses_max(self):
        """
        When velocity and forecast diverge significantly, the more
        permissive (higher) value should win.  This is the key design
        decision that prevents the forecast from being overly conservative.
        """
        scenarios = [
            (3.00, 8.00, 8.00, "forecast higher — use forecast"),
            (8.00, 3.00, 8.00, "velocity higher — use velocity"),
            (5.00, 5.00, 5.00, "equal — either works"),
            (0.50, 6.00, 6.00, "velocity very low — forecast saves the day"),
        ]

        for vel, fcast, expected, desc in scenarios:
            effective = max(vel, fcast)
            assert effective == expected, f"Failed: {desc}"


# ===========================================================================
# Edge cases
# ===========================================================================

class TestPhase3EdgeCases:
    def test_forecast_p_max_zero_uses_velocity(self):
        """forecast_p_max of 0 (budget exhausted per forecast) should not lift P_max."""
        velocity_p_max = 3.00
        forecast_p_max = 0.0
        effective = max(velocity_p_max, forecast_p_max)
        assert effective == velocity_p_max

    def test_both_zero_stays_zero(self):
        """When both signals say zero, effective is zero (budget exhausted)."""
        assert max(0.0, 0.0) == 0.0

    def test_negative_forecast_treated_as_zero(self):
        """Negative forecast_p_max (shouldn't happen) should not go below zero."""
        velocity_p_max = 2.00
        forecast_p_max = -1.0
        effective = max(velocity_p_max, max(forecast_p_max, 0.0))
        assert effective == velocity_p_max

    def test_very_large_forecast_capped_by_tier(self):
        """
        Even if forecast_p_max is very high, the hysteresis tier lock
        caps the effective model selection.
        """
        tiers = _make_tiers()
        most_expensive = tiers[-1]

        state = _make_state(
            tier_name=most_expensive.name,
            tier_min=most_expensive.min_price,
            tier_max=most_expensive.max_price,
            p_max=most_expensive.max_price,
            forecast_p_max=1000.0,
        )

        # The effective_p_max is the tier-locked value, not the raw forecast
        assert state.effective_p_max == most_expensive.max_price

    def test_state_without_forecast_is_valid(self):
        """A BudgetPacingState with no forecast data should still be usable."""
        state = _make_state(forecast_p_max=None)
        assert state.last_forecast_p_max is None
        assert state.effective_p_max > 0
        assert state.locked_model_tier is not None


# ===========================================================================
# Regression: Phase 1 + Phase 2 invariants still hold
# ===========================================================================

class TestPhase3Regressions:
    """Verify Phase 3 changes don't break Phase 1/2 invariants."""

    def test_ema_velocity_still_used_for_velocity_p_max(self):
        """
        Phase 3 adds forecast blending but the velocity path should still
        use EMA (not simple average).
        """
        from api.services.budget_service import BudgetService

        db = MagicMock()
        svc = BudgetService(db)
        budget = _make_budget(limit_amount=100.0, period="monthly")
        sub = _make_subscription()
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)

        db.query.return_value.filter.return_value.all.return_value = [budget]

        with patch.object(svc, "_compute_spend_usd", return_value=Decimal("30.0")), \
             patch.object(svc, "compute_ema_velocity", return_value=200_000.0) as mock_ema:
            svc.calculate_affordability_constraints(
                subscription=sub,
                request_tags={},
                intent_price_per_million=0.0,
                now_utc=now,
            )

        mock_ema.assert_called_once()

    def test_hysteresis_still_prevents_oscillation(self):
        """Phase 3 forecast blending should not weaken hysteresis protection."""
        tiers = _make_tiers()
        boundary_tier = tiers[1]
        boundary_price = boundary_tier.min_price

        now = datetime.now(timezone.utc)
        state = None
        locked_tiers = []

        for i in range(20):
            if i % 2 == 0:
                raw_p_max = boundary_price * 1.02
            else:
                raw_p_max = boundary_price * 0.98

            tick = now + timedelta(seconds=i * 5)

            decision = evaluate_transition(
                current_state=state,
                raw_p_max=raw_p_max,
                tiers=tiers,
                remaining_budget=50.0,
                budget_limit=100.0,
                days_remaining=15.0,
                projected_spend=50.0,
                velocity_ema=200_000.0,
                now=tick,
            )

            state = apply_transition(
                current_state=state,
                decision=decision,
                subscription_id=1,
                raw_p_max=raw_p_max,
                velocity_ema=200_000.0,
                forecast_p_max=raw_p_max * 1.1,
                tiers=tiers,
                now=tick,
            )

            locked_tiers.append(state.locked_model_tier)

        unique_tiers = set(locked_tiers)
        assert len(unique_tiers) == 1, (
            f"Oscillation detected: {unique_tiers}. Hysteresis must hold."
        )

    def test_emergency_downgrade_overrides_forecast(self):
        """
        Emergency downgrade (budget < 10%) should fire regardless of
        what the forecast says.  Safety rules are inviolable.
        """
        tiers = _make_tiers()
        premium = tiers[-1]
        now = datetime.now(timezone.utc)

        state = _make_state(
            tier_name=premium.name,
            tier_min=premium.min_price,
            tier_max=premium.max_price,
            p_max=premium.max_price,
            forecast_p_max=20.0,
        )

        decision = evaluate_transition(
            current_state=state,
            raw_p_max=20.0,
            tiers=tiers,
            remaining_budget=5.0,
            budget_limit=100.0,
            days_remaining=10.0,
            projected_spend=95.0,
            velocity_ema=200_000.0,
            now=now,
        )

        assert decision.action == "emergency_downgrade", (
            "Emergency downgrade must fire even when forecast_p_max is high"
        )
