"""
Tests for the Smart Budget (Predictive Pacing) feature.

Covers:
- get_velocity() period-aware windowing (monthly, weekly, daily)
- get_velocity() cold-start fallback to previous period
- get_velocity() blending when current data is sparse
- get_velocity() month-boundary rollover
- _previous_period_window() correctness
- calculate_affordability_constraints() P_max calculation
- calculate_affordability_constraints() performance mode threshold
- End-to-end: cold-start workspace gets a finite constraint
"""

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

from api.services.budget_service import (
    BudgetService,
    COLD_START_VELOCITY_FLOOR,
    DEFAULT_EMA_ALPHA,
    VELOCITY_MIN_DATA_DAYS,
    PERFORMANCE_MODE_PREMIUM_THRESHOLD,
    compute_period_window_utc,
)


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


# ===========================================================================
# _previous_period_window
# ===========================================================================

class TestPreviousPeriodWindow:
    """Verify that _previous_period_window produces correct adjacent windows."""

    def test_monthly_previous_period(self):
        budget = _make_budget(period="monthly", tz="UTC")
        # Current period starts Mar 1
        current_start = datetime(2026, 3, 1, 0, 0, 0, tzinfo=timezone.utc)
        prev_start, prev_end = BudgetService._previous_period_window(budget, current_start)

        assert prev_end == current_start
        assert prev_start.month == 2
        assert prev_start.day == 1

    def test_monthly_january_wraps_to_december(self):
        budget = _make_budget(period="monthly", tz="UTC")
        current_start = datetime(2026, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
        prev_start, prev_end = BudgetService._previous_period_window(budget, current_start)

        assert prev_end == current_start
        assert prev_start.year == 2025
        assert prev_start.month == 12

    def test_weekly_previous_period(self):
        budget = _make_budget(period="weekly", tz="UTC")
        current_start = datetime(2026, 3, 2, 0, 0, 0, tzinfo=timezone.utc)  # Monday
        prev_start, prev_end = BudgetService._previous_period_window(budget, current_start)

        assert prev_end == current_start
        assert (current_start - prev_start).days == 7

    def test_daily_previous_period(self):
        budget = _make_budget(period="daily", tz="UTC")
        current_start = datetime(2026, 3, 4, 0, 0, 0, tzinfo=timezone.utc)
        prev_start, prev_end = BudgetService._previous_period_window(budget, current_start)

        assert prev_end == current_start
        assert (current_start - prev_start).days == 1


# ===========================================================================
# get_velocity
# ===========================================================================

class TestGetVelocity:
    """Tests for BudgetService.get_velocity with period-aware windowing."""

    def _make_service(self):
        db = MagicMock()
        return BudgetService(db), db

    def test_cold_start_no_data_returns_floor(self):
        """With zero analytics data, velocity should return the cold-start floor."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")

        with patch.object(svc, "_query_velocity_for_window", return_value=(0, None)):
            velocity = svc.get_velocity(subscription_id=1, budget=budget)

        assert velocity == COLD_START_VELOCITY_FLOOR

    def test_sufficient_current_period_data(self):
        """When current month has >= VELOCITY_MIN_DATA_DAYS of data, use it directly."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)
        earliest = now - timedelta(days=10)
        total_tokens = 1_000_000

        def mock_query(sub_id, start, end, budget=None):
            if end == now:
                return (total_tokens, earliest)
            return (0, None)

        with patch.object(svc, "_query_velocity_for_window", side_effect=mock_query):
            velocity = svc.get_velocity(subscription_id=1, budget=budget, now_utc=now)

        expected = total_tokens / 10.0
        assert abs(velocity - expected) < 1.0

    def test_current_period_window_starts_at_period_start(self):
        """The primary query window should start at the budget period start, not a fixed 7d lookback."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 20, 12, 0, 0, tzinfo=timezone.utc)
        period_start, _ = compute_period_window_utc(budget, now_utc=now)

        captured_windows = []

        def mock_query(sub_id, start, end, budget=None):
            captured_windows.append((start, end))
            return (0, None)

        with patch.object(svc, "_query_velocity_for_window", side_effect=mock_query):
            svc.get_velocity(subscription_id=1, budget=budget, now_utc=now)

        # First call should use the budget period start (Mar 1), not now - 7d
        assert captured_windows[0][0] == period_start
        assert captured_windows[0][1] == now

    def test_month_boundary_rollover(self):
        """
        On March 1st, the current period window is [Mar 1, now) which has no data.
        The fallback should look at the *previous month* (February), not a fixed 30d window.
        """
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 1, 2, 0, 0, tzinfo=timezone.utc)  # 2am on March 1st

        feb_earliest = datetime(2026, 2, 3, 0, 0, 0, tzinfo=timezone.utc)
        feb_tokens = 2_000_000

        captured_windows = []

        def mock_query(sub_id, start, end, budget=None):
            captured_windows.append((start, end))
            if len(captured_windows) == 1:
                return (0, None)  # No data yet in March
            return (feb_tokens, feb_earliest)  # February had data

        with patch.object(svc, "_query_velocity_for_window", side_effect=mock_query):
            velocity = svc.get_velocity(subscription_id=1, budget=budget, now_utc=now)

        # Second call should be the previous period (February)
        assert captured_windows[1][1] == captured_windows[0][0]  # prev_end == current_start

        # Velocity should come from February's data
        feb_end = captured_windows[1][1]
        feb_duration = (feb_end - feb_earliest).total_seconds() / 86400.0
        expected = feb_tokens / feb_duration
        assert abs(velocity - expected) < 1.0

    def test_sparse_current_data_blends_with_previous(self):
        """When current month has < VELOCITY_MIN_DATA_DAYS, blend 70/30 with previous."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 2, 12, 0, 0, tzinfo=timezone.utc)  # Day 2 of March

        cur_earliest = now - timedelta(hours=18)
        cur_tokens = 10_000
        period_start, _ = compute_period_window_utc(budget, now_utc=now)

        feb_earliest = datetime(2026, 2, 5, 0, 0, 0, tzinfo=timezone.utc)
        feb_tokens = 1_500_000

        call_count = [0]

        def mock_query(sub_id, start, end, budget=None):
            call_count[0] += 1
            if call_count[0] == 1:
                return (cur_tokens, cur_earliest)
            return (feb_tokens, feb_earliest)

        with patch.object(svc, "_query_velocity_for_window", side_effect=mock_query):
            velocity = svc.get_velocity(subscription_id=1, budget=budget, now_utc=now)

        cur_duration = (now - cur_earliest).total_seconds() / 86400.0
        cur_vel = cur_tokens / cur_duration
        prev_end = period_start  # Feb boundary
        prev_duration = (prev_end - feb_earliest).total_seconds() / 86400.0
        prev_vel = feb_tokens / prev_duration
        expected = 0.7 * cur_vel + 0.3 * prev_vel
        assert abs(velocity - expected) < 1.0

    def test_no_budget_uses_30d_lookback(self):
        """Without a budget, fall back to a 30-day rolling window."""
        svc, _ = self._make_service()
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)
        earliest = now - timedelta(days=20)

        captured_windows = []

        def mock_query(sub_id, start, end, budget=None):
            captured_windows.append((start, end))
            if len(captured_windows) == 1:
                return (500_000, earliest)
            return (0, None)

        with patch.object(svc, "_query_velocity_for_window", side_effect=mock_query):
            velocity = svc.get_velocity(subscription_id=1, now_utc=now)

        # Without a budget, primary window should be now - 30d
        assert abs((now - captured_windows[0][0]).total_seconds() - 30 * 86400) < 1.0

    def test_budget_scope_passed_to_query(self):
        """Verify budget scope is forwarded to _query_velocity_for_window."""
        svc, _ = self._make_service()
        budget = _make_budget(scope_type="tag")
        calls = []

        def mock_query(sub_id, start, end, budget=None):
            calls.append(budget)
            return (0, None)

        with patch.object(svc, "_query_velocity_for_window", side_effect=mock_query):
            svc.get_velocity(subscription_id=1, budget=budget)

        assert all(b is budget for b in calls)

    def test_daily_budget_uses_today_window(self):
        """Daily budget should use today as the primary window, yesterday as fallback."""
        svc, _ = self._make_service()
        budget = _make_budget(period="daily")
        now = datetime(2026, 3, 4, 8, 0, 0, tzinfo=timezone.utc)  # 8am
        period_start, _ = compute_period_window_utc(budget, now_utc=now)

        captured_windows = []

        def mock_query(sub_id, start, end, budget=None):
            captured_windows.append((start, end))
            return (0, None)

        with patch.object(svc, "_query_velocity_for_window", side_effect=mock_query):
            svc.get_velocity(subscription_id=1, budget=budget, now_utc=now)

        assert captured_windows[0][0] == period_start
        # Fallback should be previous day
        assert (captured_windows[0][0] - captured_windows[1][0]).total_seconds() == pytest.approx(86400, abs=1)


# ===========================================================================
# calculate_affordability_constraints
# ===========================================================================

class TestAffordabilityConstraints:
    """Tests for BudgetService.calculate_affordability_constraints."""

    def test_basic_p_max_calculation(self):
        """P_max = remaining_budget / projected_volume_m."""
        budget = _make_budget(limit_amount=100.0, period="monthly")
        sub = _make_subscription()
        now = datetime(2026, 3, 4, 12, 0, 0, tzinfo=timezone.utc)

        db = MagicMock()
        svc = BudgetService(db)
        db.query.return_value.filter.return_value.all.return_value = [budget]

        spend = Decimal("30.0")
        velocity = 100_000.0

        with patch.object(svc, "_compute_spend_usd", return_value=spend), \
             patch.object(svc, "compute_ema_velocity", return_value=velocity):
            p_max = svc.calculate_affordability_constraints(
                subscription=sub,
                request_tags={},
                intent_price_per_million=5.0,
                now_utc=now,
            )

        remaining = 70.0
        _, end_utc = compute_period_window_utc(budget, now_utc=now)
        days_left = (end_utc - now).total_seconds() / 86400.0
        projected_vol_m = (velocity * days_left) / 1_000_000.0
        expected = remaining / projected_vol_m

        assert abs(p_max - expected) < 0.01

    def test_exhausted_budget_returns_zero(self):
        """When spend >= limit, P_max should be 0."""
        budget = _make_budget(limit_amount=50.0)
        sub = _make_subscription()
        now = datetime(2026, 3, 15, 0, 0, 0, tzinfo=timezone.utc)

        db = MagicMock()
        svc = BudgetService(db)
        db.query.return_value.filter.return_value.all.return_value = [budget]

        with patch.object(svc, "_compute_spend_usd", return_value=Decimal("55.0")), \
             patch.object(svc, "compute_ema_velocity", return_value=100_000.0):
            p_max = svc.calculate_affordability_constraints(
                subscription=sub,
                request_tags={},
                intent_price_per_million=5.0,
                now_utc=now,
            )

        assert p_max == 0.0

    def test_no_budgets_returns_infinity(self):
        """When no active auto_optimize budgets exist, return inf (no constraint)."""
        sub = _make_subscription()
        db = MagicMock()
        svc = BudgetService(db)
        db.query.return_value.filter.return_value.all.return_value = []

        p_max = svc.calculate_affordability_constraints(
            subscription=sub,
            request_tags={},
            intent_price_per_million=5.0,
        )

        assert p_max == float("inf")

    def test_performance_mode_allows_premium_when_affordable(self):
        """Performance mode returns inf when P_max >= threshold (can afford premium)."""
        budget = _make_budget(limit_amount=500.0, pacing_mode="performance")
        sub = _make_subscription()
        now = datetime(2026, 3, 4, 12, 0, 0, tzinfo=timezone.utc)

        db = MagicMock()
        svc = BudgetService(db)
        db.query.return_value.filter.return_value.all.return_value = [budget]

        with patch.object(svc, "_compute_spend_usd", return_value=Decimal("10.0")), \
             patch.object(svc, "compute_ema_velocity", return_value=50_000.0):
            p_max = svc.calculate_affordability_constraints(
                subscription=sub,
                request_tags={},
                intent_price_per_million=5.0,
                now_utc=now,
            )

        assert p_max == float("inf")

    def test_performance_mode_constrains_when_tight(self):
        """Performance mode applies ceiling when P_max < threshold."""
        budget = _make_budget(limit_amount=20.0, pacing_mode="performance")
        sub = _make_subscription()
        now = datetime(2026, 3, 4, 12, 0, 0, tzinfo=timezone.utc)

        db = MagicMock()
        svc = BudgetService(db)
        db.query.return_value.filter.return_value.all.return_value = [budget]

        with patch.object(svc, "_compute_spend_usd", return_value=Decimal("15.0")), \
             patch.object(svc, "compute_ema_velocity", return_value=500_000.0):
            p_max = svc.calculate_affordability_constraints(
                subscription=sub,
                request_tags={},
                intent_price_per_million=5.0,
                now_utc=now,
            )

        assert p_max < PERFORMANCE_MODE_PREMIUM_THRESHOLD
        assert p_max != float("inf")

    def test_cold_start_workspace_gets_finite_constraint(self):
        """
        End-to-end: a brand-new workspace with no analytics data should still
        get a finite P_max (not infinity) thanks to the velocity floor.
        """
        budget = _make_budget(limit_amount=100.0)
        sub = _make_subscription()
        now = datetime(2026, 3, 4, 12, 0, 0, tzinfo=timezone.utc)

        db = MagicMock()
        svc = BudgetService(db)
        db.query.return_value.filter.return_value.all.return_value = [budget]

        with patch.object(svc, "_compute_spend_usd", return_value=Decimal("0")), \
             patch.object(svc, "_get_daily_token_buckets", return_value=[]):
            p_max = svc.calculate_affordability_constraints(
                subscription=sub,
                request_tags={},
                intent_price_per_million=5.0,
                now_utc=now,
            )

        assert p_max != float("inf"), "Cold-start workspace must get a finite constraint"
        assert p_max > 0, "P_max should be positive (budget not exhausted)"

    def test_multiple_budgets_uses_most_restrictive(self):
        """When multiple budgets match, the lowest P_max wins."""
        generous = _make_budget(budget_id=1, limit_amount=500.0)
        tight = _make_budget(budget_id=2, limit_amount=30.0)
        sub = _make_subscription()
        now = datetime(2026, 3, 15, 0, 0, 0, tzinfo=timezone.utc)

        db = MagicMock()
        svc = BudgetService(db)
        db.query.return_value.filter.return_value.all.return_value = [generous, tight]

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

        _, end_utc = compute_period_window_utc(tight, now_utc=now)
        days_left = (end_utc - now).total_seconds() / 86400.0
        tight_p_max = 20.0 / ((200_000.0 * days_left) / 1_000_000.0)

        assert abs(p_max - tight_p_max) < 0.01


# ===========================================================================
# Constants sanity checks
# ===========================================================================

class TestConstants:
    def test_cold_start_floor_is_reasonable(self):
        assert COLD_START_VELOCITY_FLOOR > 0
        assert COLD_START_VELOCITY_FLOOR < 1_000_000

    def test_min_data_days_is_positive(self):
        assert VELOCITY_MIN_DATA_DAYS >= 1

    def test_performance_threshold_is_reasonable(self):
        assert PERFORMANCE_MODE_PREMIUM_THRESHOLD > 0
        assert PERFORMANCE_MODE_PREMIUM_THRESHOLD <= 20.0

    def test_default_ema_alpha_is_balanced(self):
        assert 0.05 <= DEFAULT_EMA_ALPHA <= 0.9


# ===========================================================================
# _get_daily_token_buckets
# ===========================================================================

class TestGetDailyTokenBuckets:
    """Tests for BudgetService._get_daily_token_buckets."""

    def _make_service(self):
        db = MagicMock()
        db.bind = MagicMock()
        db.bind.url = "sqlite:///test.db"
        return BudgetService(db), db

    def test_empty_result_returns_zero_filled_days(self):
        """When no analytics rows exist, return zeros for each day in the window."""
        svc, db = self._make_service()
        db.query.return_value.filter.return_value.group_by.return_value.all.return_value = []

        start = datetime(2026, 3, 1, 0, 0, 0, tzinfo=timezone.utc)
        end = datetime(2026, 3, 6, 0, 0, 0, tzinfo=timezone.utc)

        buckets = svc._get_daily_token_buckets(subscription_id=1, start=start, end=end)

        assert len(buckets) == 5  # Mar 1-5
        assert all(b == 0.0 for b in buckets)

    def test_partial_data_zero_fills_gaps(self):
        """Days with no data should be 0.0; days with data should have their totals."""
        svc, db = self._make_service()
        db.query.return_value.filter.return_value.group_by.return_value.all.return_value = [
            ("2026-03-01", 100_000, 100_000),
            ("2026-03-03", 500_000, 500_000),
        ]

        start = datetime(2026, 3, 1, 0, 0, 0, tzinfo=timezone.utc)
        end = datetime(2026, 3, 5, 0, 0, 0, tzinfo=timezone.utc)

        buckets = svc._get_daily_token_buckets(subscription_id=1, start=start, end=end)

        assert len(buckets) == 4  # Mar 1-4
        assert buckets[0] == 100_000.0  # Mar 1
        assert buckets[1] == 0.0        # Mar 2 (gap)
        assert buckets[2] == 500_000.0  # Mar 3
        assert buckets[3] == 0.0        # Mar 4 (gap)

    def test_prefers_llm_tokens_over_legacy(self):
        """When llm tokens > 0, use them; otherwise fall back to legacy."""
        svc, db = self._make_service()
        db.query.return_value.filter.return_value.group_by.return_value.all.return_value = [
            ("2026-03-01", 200_000, 150_000),  # llm > 0 → use llm
            ("2026-03-02", 0, 80_000),          # llm == 0 → use legacy
        ]

        start = datetime(2026, 3, 1, 0, 0, 0, tzinfo=timezone.utc)
        end = datetime(2026, 3, 3, 0, 0, 0, tzinfo=timezone.utc)

        buckets = svc._get_daily_token_buckets(subscription_id=1, start=start, end=end)

        assert buckets[0] == 200_000.0
        assert buckets[1] == 80_000.0

    def test_chronological_order(self):
        """Buckets must be returned oldest-first for correct EMA computation."""
        svc, db = self._make_service()
        db.query.return_value.filter.return_value.group_by.return_value.all.return_value = [
            ("2026-03-03", 300_000, 300_000),
            ("2026-03-01", 100_000, 100_000),
        ]

        start = datetime(2026, 3, 1, 0, 0, 0, tzinfo=timezone.utc)
        end = datetime(2026, 3, 4, 0, 0, 0, tzinfo=timezone.utc)

        buckets = svc._get_daily_token_buckets(subscription_id=1, start=start, end=end)

        assert buckets == [100_000.0, 0.0, 300_000.0]


# ===========================================================================
# compute_ema_velocity
# ===========================================================================

class TestComputeEmaVelocity:
    """Tests for BudgetService.compute_ema_velocity (Adaptive Smart Budget Phase 1)."""

    def _make_service(self):
        db = MagicMock()
        return BudgetService(db), db

    def test_cold_start_returns_floor(self):
        """With no data in current or previous period, return the floor."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")

        with patch.object(svc, "_get_daily_token_buckets", return_value=[]):
            velocity = svc.compute_ema_velocity(subscription_id=1, budget=budget)

        assert velocity == COLD_START_VELOCITY_FLOOR

    def test_stable_traffic_converges_to_daily_rate(self):
        """With constant daily traffic, EMA should converge close to that rate."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)

        daily_rate = 500_000.0
        buckets = [daily_rate] * 14  # 14 days of constant traffic

        with patch.object(svc, "_get_daily_token_buckets", return_value=buckets):
            velocity = svc.compute_ema_velocity(subscription_id=1, budget=budget, now_utc=now)

        # With constant input, EMA converges to the constant value.
        assert abs(velocity - daily_rate) < daily_rate * 0.01

    def test_spike_reacts_faster_than_simple_average(self):
        """
        A 5x traffic spike on the last day should move EMA significantly more
        than a simple average would.  This is the core value proposition.
        """
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)

        normal_rate = 100_000.0
        spike_rate = 500_000.0
        buckets = [normal_rate] * 13 + [spike_rate]  # spike on day 14

        with patch.object(svc, "_get_daily_token_buckets", return_value=buckets):
            ema_velocity = svc.compute_ema_velocity(
                subscription_id=1, budget=budget, now_utc=now, alpha=0.3,
            )

        simple_velocity = sum(buckets) / len(buckets)

        # EMA should be significantly higher than simple average after a spike
        assert ema_velocity > simple_velocity, (
            f"EMA ({ema_velocity:.0f}) should exceed simple avg ({simple_velocity:.0f}) after spike"
        )
        # EMA should be pulled toward the spike more aggressively
        assert ema_velocity > normal_rate * 1.3, (
            f"EMA ({ema_velocity:.0f}) should react noticeably to 5x spike"
        )

    def test_higher_alpha_reacts_more_aggressively(self):
        """alpha=0.5 should react more to a spike than alpha=0.1."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)

        buckets = [100_000.0] * 13 + [1_000_000.0]

        with patch.object(svc, "_get_daily_token_buckets", return_value=list(buckets)):
            slow = svc.compute_ema_velocity(
                subscription_id=1, budget=budget, now_utc=now, alpha=0.1,
            )

        with patch.object(svc, "_get_daily_token_buckets", return_value=list(buckets)):
            fast = svc.compute_ema_velocity(
                subscription_id=1, budget=budget, now_utc=now, alpha=0.5,
            )

        assert fast > slow, f"alpha=0.5 ({fast:.0f}) should react more than alpha=0.1 ({slow:.0f})"

    def test_weekend_lull_lowers_ema(self):
        """Two zero-traffic days (weekend) should pull EMA down."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 10, 12, 0, 0, tzinfo=timezone.utc)

        # Mon-Fri active, Sat-Sun zero, Mon active
        buckets = [200_000.0] * 5 + [0.0, 0.0] + [200_000.0] * 2

        with patch.object(svc, "_get_daily_token_buckets", return_value=buckets):
            velocity = svc.compute_ema_velocity(subscription_id=1, budget=budget, now_utc=now)

        # Should be lower than 200K due to the weekend zeros
        assert velocity < 200_000.0

    def test_falls_back_to_previous_period(self):
        """When current period is empty, EMA should use previous period data."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 1, 2, 0, 0, tzinfo=timezone.utc)

        prev_buckets = [150_000.0] * 28

        call_count = [0]

        def mock_buckets(sub_id, start, end, budget=None):
            call_count[0] += 1
            if call_count[0] == 1:
                return []  # Current period empty (just started)
            return prev_buckets

        with patch.object(svc, "_get_daily_token_buckets", side_effect=mock_buckets):
            velocity = svc.compute_ema_velocity(subscription_id=1, budget=budget, now_utc=now)

        assert velocity > COLD_START_VELOCITY_FLOOR
        assert abs(velocity - 150_000.0) < 150_000.0 * 0.01

    def test_never_returns_below_floor(self):
        """Even with tiny traffic, EMA should not go below the cold-start floor."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)

        buckets = [100.0] * 14  # Very low traffic

        with patch.object(svc, "_get_daily_token_buckets", return_value=buckets):
            velocity = svc.compute_ema_velocity(subscription_id=1, budget=budget, now_utc=now)

        assert velocity >= COLD_START_VELOCITY_FLOOR

    def test_strips_trailing_zeros(self):
        """Trailing zeros (partial today) should be stripped to avoid artificially lowering EMA."""
        svc, _ = self._make_service()
        budget = _make_budget(period="monthly")
        now = datetime(2026, 3, 10, 2, 0, 0, tzinfo=timezone.utc)

        # Day 9 had real traffic, day 10 just started (0 tokens so far)
        buckets = [200_000.0] * 9 + [0.0]

        with patch.object(svc, "_get_daily_token_buckets", return_value=buckets):
            velocity = svc.compute_ema_velocity(subscription_id=1, budget=budget, now_utc=now)

        # Should converge to ~200K, not be dragged down by the trailing zero
        assert abs(velocity - 200_000.0) < 200_000.0 * 0.01

    def test_no_budget_uses_30d_window(self):
        """Without a budget, EMA should use a 30-day rolling window."""
        svc, _ = self._make_service()
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)

        captured_args = []

        def mock_buckets(sub_id, start, end, budget=None):
            captured_args.append((start, end))
            return [100_000.0] * 14

        with patch.object(svc, "_get_daily_token_buckets", side_effect=mock_buckets):
            svc.compute_ema_velocity(subscription_id=1, now_utc=now)

        start_arg = captured_args[0][0]
        expected_start = now - timedelta(days=30)
        assert abs((start_arg - expected_start).total_seconds()) < 1.0


# ===========================================================================
# calculate_affordability_constraints with use_ema
# ===========================================================================

class TestAffordabilityConstraintsEma:
    """Tests for calculate_affordability_constraints EMA velocity (default)."""

    def test_default_uses_ema_velocity(self):
        """By default, should call compute_ema_velocity (not get_velocity)."""
        budget = _make_budget(limit_amount=100.0, period="monthly")
        sub = _make_subscription()
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)

        db = MagicMock()
        svc = BudgetService(db)
        db.query.return_value.filter.return_value.all.return_value = [budget]

        ema_velocity = 200_000.0

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

        mock_ema.assert_called_once()
        mock_simple.assert_not_called()
        assert p_max > 0

    def test_use_ema_false_falls_back_to_simple(self):
        """Passing use_ema=False should use the legacy simple average."""
        budget = _make_budget(limit_amount=100.0, period="monthly")
        sub = _make_subscription()
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)

        db = MagicMock()
        svc = BudgetService(db)
        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, "get_velocity", return_value=100_000.0) as mock_simple, \
             patch.object(svc, "compute_ema_velocity") as mock_ema:
            p_max = svc.calculate_affordability_constraints(
                subscription=sub,
                request_tags={},
                intent_price_per_million=0.0,
                now_utc=now,
                use_ema=False,
            )

        mock_simple.assert_called_once()
        mock_ema.assert_not_called()
        assert p_max > 0

    def test_ema_spike_produces_lower_p_max(self):
        """
        After a traffic spike, EMA velocity is higher → projected volume is higher
        → P_max is lower (more conservative). This is the desired behavior.
        """
        budget = _make_budget(limit_amount=100.0, period="monthly")
        sub = _make_subscription()
        now = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)

        db = MagicMock()
        svc = BudgetService(db)
        db.query.return_value.filter.return_value.all.return_value = [budget]

        # Simple velocity: low (hasn't caught up to spike)
        simple_velocity = 100_000.0
        # EMA velocity: high (reacted to spike)
        ema_velocity = 400_000.0

        with patch.object(svc, "_compute_spend_usd", return_value=Decimal("30.0")), \
             patch.object(svc, "get_velocity", return_value=simple_velocity):
            simple_p_max = svc.calculate_affordability_constraints(
                subscription=sub, request_tags={}, intent_price_per_million=0.0,
                now_utc=now, use_ema=False,
            )

        with patch.object(svc, "_compute_spend_usd", return_value=Decimal("30.0")), \
             patch.object(svc, "compute_ema_velocity", return_value=ema_velocity):
            ema_p_max = svc.calculate_affordability_constraints(
                subscription=sub, request_tags={}, intent_price_per_million=0.0,
                now_utc=now,
            )

        assert ema_p_max < simple_p_max, (
            f"EMA P_max (${ema_p_max:.2f}) should be lower than simple (${simple_p_max:.2f}) "
            "after a spike because higher velocity → more projected usage → tighter budget"
        )
