dy-dashboard/services/planning_snapshot.py

"""Daily planning suggestion snapshots and consecutive-day counters."""
from datetime import date, datetime, timedelta
from zoneinfo import ZoneInfo

from services.db import get_db, query
from services.cache import bump_data_version


TABLE = 'planning_suggestion_snapshots'
LOCK_NAME = 'dy_planning_suggestion_snapshot'
_schema_ready = False


def shanghai_today():
    return datetime.now(ZoneInfo('Asia/Shanghai')).date()


def as_date(value):
    if isinstance(value, datetime):
        return value.date()
    if isinstance(value, date):
        return value
    return date.fromisoformat(str(value))


def ensure_schema():
    global _schema_ready
    if _schema_ready:
        return
    query(f"""
        CREATE TABLE IF NOT EXISTS {TABLE} (
            id BIGINT NOT NULL AUTO_INCREMENT,
            snapshot_date DATE NOT NULL,
            product_code VARCHAR(100) NOT NULL,
            product_name VARCHAR(300) NOT NULL DEFAULT '',
            suggested_qty INT NOT NULL DEFAULT 0,
            consecutive_days INT NOT NULL DEFAULT 0,
            created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (id),
            UNIQUE KEY uk_planning_snapshot_date_product (snapshot_date, product_code),
            KEY idx_planning_snapshot_product_date (product_code, snapshot_date)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """)
    _schema_ready = True


def snapshot_exists(day=None):
    ensure_schema()
    day = as_date(day or shanghai_today())
    rows = query(f"SELECT 1 AS found FROM {TABLE} WHERE snapshot_date=%s LIMIT 1", (day,))
    return bool(rows)


def _next_counter(suggested_qty, previous):
    if suggested_qty <= 0:
        return 0
    if previous and int(previous.get('suggested_qty') or 0) > 0:
        return int(previous.get('consecutive_days') or 0) + 1
    return 1


def store_snapshot(rows, day=None):
    """Persist one immutable snapshot for a natural day, using the prior day's counter."""
    ensure_schema()
    day = as_date(day or shanghai_today())
    previous_day = day - timedelta(days=1)
    conn = get_db()
    locked = False
    try:
        with conn.cursor() as cur:
            cur.execute('SELECT GET_LOCK(%s, 30) AS acquired', (LOCK_NAME,))
            locked = bool((cur.fetchone() or {}).get('acquired'))
            if not locked:
                raise RuntimeError('Another planning snapshot job is running')
            cur.execute(f"SELECT 1 AS found FROM {TABLE} WHERE snapshot_date=%s LIMIT 1", (day,))
            if cur.fetchone():
                conn.rollback()
                return {'date': day.isoformat(), 'created': False, 'rows': 0,
                        'message': 'Snapshot already exists; unchanged'}
            cur.execute(
                f"SELECT product_code, suggested_qty, consecutive_days FROM {TABLE} "
                "WHERE snapshot_date=%s",
                (previous_day,),
            )
            previous = {row['product_code']: row for row in cur.fetchall()}
            payload = []
            for row in rows:
                code = str(row.get('product_code') or '').strip()
                if not code:
                    continue
                suggested = max(0, int(float(row.get('建议下单数量') or 0)))
                payload.append((
                    day, code, str(row.get('product_name') or '')[:300], suggested,
                    _next_counter(suggested, previous.get(code)),
                ))
            if not payload:
                raise RuntimeError('No planning rows available; snapshot not written')
            cur.executemany(
                f"INSERT INTO {TABLE} "
                "(snapshot_date, product_code, product_name, suggested_qty, consecutive_days) "
                "VALUES (%s,%s,%s,%s,%s)",
                payload,
            )
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        if locked:
            try:
                with conn.cursor() as cur:
                    cur.execute('SELECT RELEASE_LOCK(%s)', (LOCK_NAME,))
            except Exception:
                pass
        conn.close()
    bump_data_version()
    positive = sum(1 for item in payload if item[3] > 0)
    return {'date': day.isoformat(), 'created': True, 'rows': len(payload),
            'positive': positive, 'message': 'Snapshot created'}


def consecutive_days_map(codes, day=None):
    ensure_schema()
    day = as_date(day or shanghai_today())
    result = {}
    unique = list(dict.fromkeys(str(code) for code in codes if code))
    for start in range(0, len(unique), 800):
        chunk = unique[start:start + 800]
        placeholders = ','.join(['%s'] * len(chunk))
        rows = query(
            f"SELECT product_code, consecutive_days FROM {TABLE} "
            f"WHERE snapshot_date=%s AND product_code IN ({placeholders})",
            (day, *chunk),
        )
        result.update({row['product_code']: int(row.get('consecutive_days') or 0) for row in rows})
    return result


def attach_consecutive_days(rows, day=None):
    """Attach the immutable snapshot counter to planning rows as read-only data."""
    values = consecutive_days_map([row.get('product_code') for row in rows], day)
    for row in rows:
        row['连续建议天数'] = values.get(str(row.get('product_code') or ''), 0)
    return rows