#!/usr/bin/env python
"""Build and publish weighted-speed parameters from completed revenue days."""

from __future__ import annotations

import argparse
import json
from datetime import date, datetime, timedelta
from zoneinfo import ZoneInfo

from odps_source import build_odps_client
from revenue_forecast_config import RevenueForecastConfig
from revenue_forecast_repository import publish_speed_parameter_release
from revenue_forecast_source import fetch_interval_revenue_series
from revenue_speed_forecast import (
    aggregate_speed_parameters,
    build_speed_samples,
)
from run_once import load_environment
from storage import initialize_schema


SHANGHAI = ZoneInfo("Asia/Shanghai")


def _parse_date(value: str) -> date:
    return datetime.strptime(value, "%Y%m%d").date()


def _date_range(start: date, end: date) -> tuple[date, ...]:
    if start > end:
        raise ValueError("start date must not be after end date")
    return tuple(
        start + timedelta(days=offset)
        for offset in range((end - start).days + 1)
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--start-date", help="Training start date in YYYYMMDD")
    parser.add_argument("--end-date", help="Training end date in YYYYMMDD")
    parser.add_argument("--parameter-version", help="Published parameter version")
    return parser.parse_args()


def main() -> int:
    load_environment()
    args = parse_args()
    config = RevenueForecastConfig.from_env()
    today = datetime.now(SHANGHAI).date()
    yesterday = today - timedelta(days=1)
    end_date = _parse_date(args.end_date) if args.end_date else yesterday
    start_date = (
        _parse_date(args.start_date)
        if args.start_date
        else end_date - timedelta(days=29)
    )
    if end_date >= today:
        raise ValueError("Training end date must be a completed historical day")
    parameter_version = (
        args.parameter_version or config.speed_parameter_version
    )
    dates = _date_range(start_date, end_date)
    series = fetch_interval_revenue_series(build_odps_client(), dates)

    samples = []
    complete_dates = []
    skipped_dates = []
    for data_date in dates:
        day_samples = build_speed_samples(
            data_date=data_date,
            intervals=series.get(data_date, []),
            parameter_version=parameter_version,
            latest_weight=config.speed_latest_weight,
            start_time=config.start_time,
            stop_time=config.stop_time,
        )
        if not day_samples:
            skipped_dates.append(data_date.strftime("%Y%m%d"))
            continue
        complete_dates.append(data_date.strftime("%Y%m%d"))
        samples.extend(day_samples)

    parameters = aggregate_speed_parameters(samples)
    if not parameters:
        raise RuntimeError("No complete historical day available for publication")
    initialize_schema()
    publish_speed_parameter_release(
        parameter_version=parameter_version,
        samples=samples,
        parameters=parameters,
    )
    print(
        json.dumps(
            {
                "status": "PUBLISHED",
                "parameter_version": parameter_version,
                "latest_weight": str(config.speed_latest_weight),
                "complete_dates": complete_dates,
                "skipped_dates": skipped_dates,
                "sample_rows": len(samples),
                "time_slots": len(parameters),
                "minimum_samples_per_slot": min(
                    item.sample_count for item in parameters
                ),
                "maximum_samples_per_slot": max(
                    item.sample_count for item in parameters
                ),
            },
            ensure_ascii=False,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
