from __future__ import annotations

import unittest
from datetime import date, datetime, time, timedelta
from decimal import Decimal
from zoneinfo import ZoneInfo

from revenue_forecast import RevenueObservation, RevenueTrendWindow
from revenue_speed_forecast import (
    RevenueSpeedParameter,
    SpeedForecastUnavailable,
    aggregate_speed_parameters,
    build_speed_samples,
    calculate_speed_forecast,
)


SHANGHAI = ZoneInfo("Asia/Shanghai")


def observation() -> RevenueObservation:
    return RevenueObservation(
        partition="20260803",
        report_time=datetime(2026, 8, 3, 12, 0, tzinfo=SHANGHAI),
        today_revenue=Decimal("500"),
    )


def windows() -> list[RevenueTrendWindow]:
    return [
        RevenueTrendWindow(
            report_time=datetime(2026, 8, 3, 11, 45, tzinfo=SHANGHAI),
            today_revenue=Decimal("30"),
        ),
        RevenueTrendWindow(
            report_time=datetime(2026, 8, 3, 12, 0, tzinfo=SHANGHAI),
            today_revenue=Decimal("60"),
        ),
    ]


def parameter(
    *,
    samples: int = 10,
    training_end: date | None = None,
) -> RevenueSpeedParameter:
    return RevenueSpeedParameter(
        parameter_version="params_v1",
        time_slot=time(12, 0),
        training_start=date(2026, 7, 1),
        training_end=training_end or date(2026, 8, 2),
        sample_count=samples,
        latest_weight=Decimal("0.66666667"),
        multiplier_p10=Decimal("5"),
        multiplier_p50=Decimal("10"),
        multiplier_p90=Decimal("15"),
        multiplier_mad=Decimal("2"),
    )


class RevenueForecastTest(unittest.TestCase):
    def test_forecast_uses_weighted_speed_and_p50(self) -> None:
        result = calculate_speed_forecast(
            observation(),
            windows(),
            parameter(),
            forecast_version="forecast_v6",
            parameter_version="params_v1",
            target_cost_ratio=Decimal("3.5"),
            minimum_samples=3,
        )
        expected_speed = (
            Decimal("0.66666667") * Decimal("60")
            + Decimal("0.33333333") * Decimal("30")
        )
        self.assertEqual("WEIGHTED_SPEED_P50", result.status)
        self.assertEqual(expected_speed, result.weighted_speed)
        self.assertEqual(Decimal("500") + expected_speed * 10, result.forecast_revenue)
        self.assertEqual(Decimal("500") + expected_speed * 5, result.forecast_lower)
        self.assertEqual(Decimal("500") + expected_speed * 15, result.forecast_upper)

    def test_insufficient_parameter_samples_stop_forecast(self) -> None:
        with self.assertRaisesRegex(
            SpeedForecastUnavailable,
            "INSUFFICIENT_PARAMETER_SAMPLES",
        ):
            calculate_speed_forecast(
                observation(),
                windows(),
                parameter(samples=1),
                forecast_version="forecast_v6",
                parameter_version="params_v1",
                target_cost_ratio=Decimal("3.5"),
                minimum_samples=3,
            )

    def test_current_or_future_training_data_stop_forecast(self) -> None:
        with self.assertRaisesRegex(
            SpeedForecastUnavailable,
            "PARAMETER_USES_CURRENT_OR_FUTURE_DATA",
        ):
            calculate_speed_forecast(
                observation(),
                windows(),
                parameter(training_end=date(2026, 8, 3)),
                forecast_version="forecast_v6",
                parameter_version="params_v1",
                target_cost_ratio=Decimal("3.5"),
                minimum_samples=3,
            )

    def test_samples_and_parameters_share_one_formula(self) -> None:
        intervals = [
            (
                datetime(2026, 8, 1, 0, 0, tzinfo=SHANGHAI)
                + timedelta(minutes=15 * index),
                Decimal(index + 1),
            )
            for index in range(96)
        ]
        samples = build_speed_samples(
            data_date=date(2026, 8, 1),
            intervals=intervals,
            parameter_version="params_v1",
            latest_weight=Decimal("0.66666667"),
            start_time=time(6, 15),
            stop_time=time(22, 30),
        )
        self.assertEqual(66, len(samples))
        first = samples[0]
        result = aggregate_speed_parameters([first])[0]
        self.assertEqual(first.remaining_multiplier, result.multiplier_p50)

    def test_non_contiguous_history_is_not_published(self) -> None:
        intervals = [
            (
                datetime(2026, 8, 1, 0, 0, tzinfo=SHANGHAI)
                + timedelta(minutes=15 * index),
                Decimal("1"),
            )
            for index in range(96)
        ]
        intervals[50] = (intervals[50][0] + timedelta(minutes=1), Decimal("1"))
        self.assertEqual(
            [],
            build_speed_samples(
                data_date=date(2026, 8, 1),
                intervals=intervals,
                parameter_version="params_v1",
                latest_weight=Decimal("0.66666667"),
                start_time=time(6, 15),
                stop_time=time(22, 30),
            ),
        )


if __name__ == "__main__":
    unittest.main()
