from __future__ import annotations

"""Вторая миграция: ограничивает значения заказов 1 в день по каждой категории."""

import sqlite3
from pathlib import Path

MIGRATION_ID = "0002_cap_orders_to_one"
DESCRIPTION = "Cap daily orders per category at 1 for each client/day (values >1 -> 1)"

ORDER_COLS = [
    "orders_amt_ent",
    "orders_amt_super",
    "orders_amt_transport",
    "orders_amt_shopping",
    "orders_amt_hotel",
    "orders_amt_avia",
]


def run(context) -> None:
    # Работаем с уже собранной SQLite, путь берём из контекста мигратора
    dataset_dir = Path(getattr(context, "dataset_dir", Path.cwd()))
    sqlite_path = getattr(context, "sqlite_path", dataset_dir / "ds.sqlite")

    if not sqlite_path.exists():
        raise FileNotFoundError(f"SQLite database not found: {sqlite_path}")

    conn = sqlite3.connect(sqlite_path)
    try:
        # Каждую колонку приводим к максимуму 1, чтобы убрать аномальные значения
        for col in ORDER_COLS:
            sql = f"""
            UPDATE communications
            SET {col} = CASE WHEN {col} > 1 THEN 1 ELSE {col} END
            """
            conn.execute(sql)
        conn.commit()
        print(f"Capped values >1 to 1 in columns: {', '.join(ORDER_COLS)}")
    finally:
        conn.close()