"""MySQL persistence for revenue observations and forecasts."""

from __future__ import annotations

import json
from datetime import datetime
from decimal import Decimal
from typing import Any

from revenue_forecast import RevenueForecast, RevenueObservation
from revenue_speed_forecast import RevenueSpeedParameter, RevenueSpeedSample
from storage import connect


def _db_datetime(value: datetime | None) -> datetime | None:
    if value is None:
        return None
    return value.replace(tzinfo=None)


def _aware_datetime(value: datetime, timezone: Any) -> datetime:
    if value.tzinfo is None:
        return value.replace(tzinfo=timezone)
    return value.astimezone(timezone)


def load_previous_observation(
    report_time: datetime,
    source_version: str,
) -> RevenueObservation | None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT *
                FROM revenue_forecast_observation
                WHERE report_time < %s
                  AND report_time >= DATE(%s)
                  AND source_version = %s
                  AND quality_status = 'VALID'
                ORDER BY report_time DESC
                LIMIT 1
                """,
                (
                    _db_datetime(report_time),
                    _db_datetime(report_time),
                    source_version,
                ),
            )
            row = cursor.fetchone()
            return _row_to_observation(row, report_time.tzinfo) if row else None
    finally:
        connection.close()


def upsert_observation(
    observation: RevenueObservation,
    *,
    source_version: str,
    lag_seconds: int,
    quality_status: str,
    quality_message: str | None,
) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO revenue_forecast_observation
                    (source_partition, source_version, report_time, today_revenue,
                     overall_cpm, impressions, dau, fill_rate,
                     revenue_change_pct, cpm_change_pct,
                     exposure_change_pct, dau_change_pct,
                     fill_rate_change_pct,
                     source_modified_at, lag_seconds, quality_status,
                     quality_message)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                ON DUPLICATE KEY UPDATE
                    source_partition=VALUES(source_partition),
                    source_version=VALUES(source_version),
                    today_revenue=VALUES(today_revenue),
                    overall_cpm=VALUES(overall_cpm),
                    impressions=VALUES(impressions),
                    dau=VALUES(dau),
                    fill_rate=VALUES(fill_rate),
                    revenue_change_pct=VALUES(revenue_change_pct),
                    cpm_change_pct=VALUES(cpm_change_pct),
                    exposure_change_pct=VALUES(exposure_change_pct),
                    dau_change_pct=VALUES(dau_change_pct),
                    fill_rate_change_pct=VALUES(fill_rate_change_pct),
                    source_modified_at=VALUES(source_modified_at),
                    lag_seconds=VALUES(lag_seconds),
                    quality_status=VALUES(quality_status),
                    quality_message=VALUES(quality_message)
                """,
                (
                    observation.partition,
                    source_version,
                    _db_datetime(observation.report_time),
                    observation.today_revenue,
                    observation.overall_cpm,
                    observation.impressions,
                    observation.dau,
                    observation.fill_rate,
                    observation.revenue_change_pct,
                    observation.cpm_change_pct,
                    observation.exposure_change_pct,
                    observation.dau_change_pct,
                    observation.fill_rate_change_pct,
                    _db_datetime(observation.source_modified_at),
                    lag_seconds,
                    quality_status,
                    quality_message,
                ),
            )
    finally:
        connection.close()


def save_forecast(forecast: RevenueForecast) -> None:
    channel_costs_json = json.dumps(
        {channel: str(cost) for channel, cost in forecast.channel_costs},
        ensure_ascii=False,
        sort_keys=True,
    )
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO revenue_forecast_result
                    (report_time, forecast_version, forecast_method,
                     parameter_version, current_cumulative_revenue,
                     speed_latest_weight, latest_interval_revenue,
                     previous_interval_revenue, weighted_speed,
                     remaining_multiplier_p10, remaining_multiplier_p50,
                     remaining_multiplier_p90, parameter_sample_count,
                     forecast_revenue, forecast_lower, forecast_upper,
                     target_cost_ratio, target_daily_cost, cost_reserve_date,
                     channel_costs_json, non_miniapp_reserved_cost,
                     miniapp_target_daily_cost, forecast_status)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                ON DUPLICATE KEY UPDATE
                    forecast_method=VALUES(forecast_method),
                    parameter_version=VALUES(parameter_version),
                    current_cumulative_revenue=VALUES(current_cumulative_revenue),
                    speed_latest_weight=VALUES(speed_latest_weight),
                    latest_interval_revenue=VALUES(latest_interval_revenue),
                    previous_interval_revenue=VALUES(previous_interval_revenue),
                    weighted_speed=VALUES(weighted_speed),
                    remaining_multiplier_p10=VALUES(remaining_multiplier_p10),
                    remaining_multiplier_p50=VALUES(remaining_multiplier_p50),
                    remaining_multiplier_p90=VALUES(remaining_multiplier_p90),
                    parameter_sample_count=VALUES(parameter_sample_count),
                    forecast_revenue=VALUES(forecast_revenue),
                    forecast_lower=VALUES(forecast_lower),
                    forecast_upper=VALUES(forecast_upper),
                    target_cost_ratio=VALUES(target_cost_ratio),
                    target_daily_cost=VALUES(target_daily_cost),
                    cost_reserve_date=VALUES(cost_reserve_date),
                    channel_costs_json=VALUES(channel_costs_json),
                    non_miniapp_reserved_cost=VALUES(non_miniapp_reserved_cost),
                    miniapp_target_daily_cost=VALUES(miniapp_target_daily_cost),
                    forecast_status=VALUES(forecast_status)
                """,
                (
                    _db_datetime(forecast.report_time),
                    forecast.forecast_version,
                    forecast.forecast_method,
                    forecast.parameter_version,
                    forecast.current_cumulative_revenue,
                    forecast.speed_latest_weight,
                    forecast.latest_interval_revenue,
                    forecast.previous_interval_revenue,
                    forecast.weighted_speed,
                    forecast.remaining_multiplier_p10,
                    forecast.remaining_multiplier_p50,
                    forecast.remaining_multiplier_p90,
                    forecast.parameter_sample_count,
                    forecast.forecast_revenue,
                    forecast.forecast_lower,
                    forecast.forecast_upper,
                    forecast.target_cost_ratio,
                    forecast.target_daily_cost,
                    forecast.cost_reserve_date,
                    channel_costs_json,
                    forecast.non_miniapp_reserved_cost,
                    forecast.miniapp_target_daily_cost,
                    forecast.status,
                ),
            )
    finally:
        connection.close()


def load_speed_parameter(
    report_time: datetime,
    parameter_version: str,
) -> RevenueSpeedParameter | None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT *, TIME_FORMAT(time_slot, '%%H:%%i:%%s') AS time_slot_text
                FROM revenue_forecast_speed_parameter
                WHERE parameter_version = %s
                  AND time_slot = TIME(%s)
                LIMIT 1
                """,
                (parameter_version, _db_datetime(report_time)),
            )
            row = cursor.fetchone()
    finally:
        connection.close()
    if not row:
        return None
    return RevenueSpeedParameter(
        parameter_version=str(row["parameter_version"]),
        time_slot=datetime.strptime(row["time_slot_text"], "%H:%M:%S").time(),
        training_start=row["training_start"],
        training_end=row["training_end"],
        sample_count=int(row["sample_count"]),
        latest_weight=Decimal(str(row["latest_weight"])),
        multiplier_p10=Decimal(str(row["multiplier_p10"])),
        multiplier_p50=Decimal(str(row["multiplier_p50"])),
        multiplier_p90=Decimal(str(row["multiplier_p90"])),
        multiplier_mad=Decimal(str(row["multiplier_mad"])),
    )


def publish_speed_parameter_release(
    *,
    parameter_version: str,
    samples: list[RevenueSpeedSample],
    parameters: list[RevenueSpeedParameter],
) -> None:
    if not samples or not parameters:
        raise ValueError("Cannot publish an empty speed parameter release")
    if any(item.parameter_version != parameter_version for item in samples):
        raise ValueError("Sample parameter version mismatch")
    if any(item.parameter_version != parameter_version for item in parameters):
        raise ValueError("Aggregate parameter version mismatch")

    connection = connect()
    try:
        connection.autocommit(False)
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT COUNT(*) AS row_count "
                "FROM revenue_forecast_speed_parameter "
                "WHERE parameter_version=%s",
                (parameter_version,),
            )
            if int(cursor.fetchone()["row_count"]) > 0:
                raise ValueError(
                    f"Parameter version already published: {parameter_version}"
                )
            cursor.executemany(
                """
                INSERT INTO revenue_forecast_speed_sample
                    (data_date, time_slot, parameter_version,
                     cumulative_revenue, final_revenue,
                     latest_interval_revenue, previous_interval_revenue,
                     latest_weight, weighted_speed, remaining_revenue,
                     remaining_multiplier)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                """,
                [
                    (
                        item.data_date,
                        item.report_time.time(),
                        item.parameter_version,
                        item.cumulative_revenue,
                        item.final_revenue,
                        item.latest_interval_revenue,
                        item.previous_interval_revenue,
                        item.latest_weight,
                        item.weighted_speed,
                        item.remaining_revenue,
                        item.remaining_multiplier,
                    )
                    for item in samples
                ],
            )
            cursor.executemany(
                """
                INSERT INTO revenue_forecast_speed_parameter
                    (parameter_version, time_slot, training_start,
                     training_end, sample_count, latest_weight,
                     multiplier_p10, multiplier_p50, multiplier_p90,
                     multiplier_mad)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                """,
                [
                    (
                        item.parameter_version,
                        item.time_slot,
                        item.training_start,
                        item.training_end,
                        item.sample_count,
                        item.latest_weight,
                        item.multiplier_p10,
                        item.multiplier_p50,
                        item.multiplier_p90,
                        item.multiplier_mad,
                    )
                    for item in parameters
                ],
            )
        connection.commit()
    except Exception:
        connection.rollback()
        raise
    finally:
        connection.autocommit(True)
        connection.close()


def _row_to_observation(row: dict[str, Any], timezone: Any) -> RevenueObservation:
    return RevenueObservation(
        partition=str(row["source_partition"]),
        report_time=_aware_datetime(row["report_time"], timezone),
        today_revenue=Decimal(str(row["today_revenue"])),
        overall_cpm=(
            Decimal(str(row["overall_cpm"]))
            if row.get("overall_cpm") is not None
            else None
        ),
        impressions=(
            int(row["impressions"]) if row.get("impressions") is not None else None
        ),
        dau=int(row["dau"]) if row.get("dau") is not None else None,
        fill_rate=(
            Decimal(str(row["fill_rate"]))
            if row.get("fill_rate") is not None
            else None
        ),
        revenue_change_pct=(
            Decimal(str(row["revenue_change_pct"]))
            if row.get("revenue_change_pct") is not None
            else None
        ),
        cpm_change_pct=(
            Decimal(str(row["cpm_change_pct"]))
            if row.get("cpm_change_pct") is not None
            else None
        ),
        exposure_change_pct=(
            Decimal(str(row["exposure_change_pct"]))
            if row.get("exposure_change_pct") is not None
            else None
        ),
        dau_change_pct=(
            Decimal(str(row["dau_change_pct"]))
            if row.get("dau_change_pct") is not None
            else None
        ),
        fill_rate_change_pct=(
            Decimal(str(row["fill_rate_change_pct"]))
            if row.get("fill_rate_change_pct") is not None
            else None
        ),
        source_modified_at=(
            _aware_datetime(row["source_modified_at"], timezone)
            if row.get("source_modified_at") is not None
            else None
        ),
    )
