Agent/examples/tencent_realtime_control/storage.py

"""MySQL state and audit storage for real-time control."""

from __future__ import annotations

import os
import json
from contextlib import contextmanager
from datetime import date, datetime
from pathlib import Path
from typing import Any, Iterator

import pymysql
import pymysql.cursors


ROOT = Path(__file__).resolve().parent


def connect() -> pymysql.Connection:
    required = ["DB_HOST", "DB_USER", "DB_NAME"]
    missing = [key for key in required if not os.getenv(key)]
    if missing:
        raise RuntimeError(f"Missing database environment variables: {', '.join(missing)}")
    return pymysql.connect(
        host=os.environ["DB_HOST"],
        port=int(os.getenv("DB_PORT", "3306")),
        user=os.environ["DB_USER"],
        password=os.getenv("DB_PASSWORD", ""),
        database=os.environ["DB_NAME"],
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
        autocommit=True,
        connect_timeout=int(os.getenv("DB_CONNECT_TIMEOUT", "10")),
        read_timeout=int(os.getenv("DB_READ_TIMEOUT", "60")),
        write_timeout=int(os.getenv("DB_WRITE_TIMEOUT", "60")),
    )


def initialize_schema() -> None:
    statements = [
        statement.strip()
        for statement in (ROOT / "schema.sql").read_text(encoding="utf-8").split(";")
        if statement.strip()
    ]
    connection = connect()
    try:
        with connection.cursor() as cursor:
            for statement in statements:
                cursor.execute(statement)
            cursor.execute(
                """
                SELECT COLUMN_NAME
                FROM information_schema.COLUMNS
                WHERE TABLE_SCHEMA=%s
                  AND TABLE_NAME='realtime_control_ad_state'
                """,
                (os.environ["DB_NAME"],),
            )
            existing_columns = {row["COLUMN_NAME"] for row in cursor.fetchall()}
            migrations = {
                "initial_base_bid_fen": "INT DEFAULT NULL",
                "last_roi_scaled_at": "DATETIME DEFAULT NULL",
                "bid_hold": "BOOLEAN NOT NULL DEFAULT FALSE",
                "bid_hold_reason": "VARCHAR(255) DEFAULT NULL",
                "operator_pause_mode": "VARCHAR(32) DEFAULT NULL",
                "operator_resume_at": "DATETIME DEFAULT NULL",
                "operator_command_id": "VARCHAR(64) DEFAULT NULL",
                "operator_paused_from_status": "VARCHAR(50) DEFAULT NULL",
                "operator_paused_at": "DATETIME DEFAULT NULL",
            }
            for column, definition in migrations.items():
                if column not in existing_columns:
                    cursor.execute(
                        f"ALTER TABLE realtime_control_ad_state "
                        f"ADD COLUMN {column} {definition}"
                    )
            cursor.execute(
                """
                UPDATE realtime_control_ad_state
                SET initial_base_bid_fen=base_bid_fen
                WHERE initial_base_bid_fen IS NULL
                """
            )
            cursor.execute(
                """
                SELECT INDEX_NAME
                FROM information_schema.STATISTICS
                WHERE TABLE_SCHEMA=%s
                  AND TABLE_NAME='realtime_control_ad_state'
                  AND INDEX_NAME='idx_operator_pause'
                """,
                (os.environ["DB_NAME"],),
            )
            if not cursor.fetchone():
                cursor.execute(
                    """
                    CREATE INDEX idx_operator_pause
                    ON realtime_control_ad_state
                        (operator_pause_mode, operator_resume_at)
                    """
                )

            operator_migrations = {
                "operator_command": {
                    "raw_text": "MEDIUMTEXT DEFAULT NULL",
                    "parse_source": "VARCHAR(32) NOT NULL DEFAULT 'deterministic'",
                    "intent_json": "MEDIUMTEXT DEFAULT NULL",
                    "preview_cost_fen": "BIGINT NOT NULL DEFAULT 0",
                    "preview_impressions": "BIGINT NOT NULL DEFAULT 0",
                    "preview_clicks": "BIGINT NOT NULL DEFAULT 0",
                    "preview_conversions": "BIGINT NOT NULL DEFAULT 0",
                    "previewed_at": "DATETIME DEFAULT NULL",
                    "resume_at": "DATETIME DEFAULT NULL",
                },
                "operator_command_item": {
                    "preview_cost_fen": "BIGINT NOT NULL DEFAULT 0",
                    "preview_impressions": "BIGINT NOT NULL DEFAULT 0",
                    "preview_clicks": "BIGINT NOT NULL DEFAULT 0",
                    "preview_conversions": "BIGINT NOT NULL DEFAULT 0",
                    "previewed_at": "DATETIME DEFAULT NULL",
                    "updated_at": (
                        "TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP "
                        "ON UPDATE CURRENT_TIMESTAMP"
                    ),
                },
            }
            for table_name, columns in operator_migrations.items():
                cursor.execute(
                    """
                    SELECT COLUMN_NAME
                    FROM information_schema.COLUMNS
                    WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s
                    """,
                    (os.environ["DB_NAME"], table_name),
                )
                table_columns = {row["COLUMN_NAME"] for row in cursor.fetchall()}
                for column, definition in columns.items():
                    if column not in table_columns:
                        cursor.execute(
                            f"ALTER TABLE {table_name} ADD COLUMN {column} {definition}"
                        )

            roi_migrations = {
                "roi_metric_run": {
                    "fission_parameter_version": "VARCHAR(64) DEFAULT NULL",
                    "fission_cohort_date": "DATE DEFAULT NULL",
                },
                "roi_entity_snapshot": {
                    "conversion_goal": "VARCHAR(255) DEFAULT NULL",
                    "fission_parameter_version": "VARCHAR(64) DEFAULT NULL",
                    "fission_cohort_date": "VARCHAR(8) DEFAULT NULL",
                    "fission_multiplier_vs_t0": "DECIMAL(18,8) DEFAULT NULL",
                    "fission_multiplier_vs_first": "DECIMAL(18,8) DEFAULT NULL",
                    "fission_match_level": "VARCHAR(64) DEFAULT NULL",
                    "fission_source": "VARCHAR(512) DEFAULT NULL",
                    "actual_total_revenue": "DECIMAL(20,4) DEFAULT NULL",
                    "actual_roi": "DECIMAL(18,8) DEFAULT NULL",
                    "predicted_tail_revenue": "DECIMAL(20,4) DEFAULT NULL",
                    "predicted_fission_revenue": "DECIMAL(20,4) DEFAULT NULL",
                },
                "roi_fission_parameter_value": {
                    "multiplier_vs_first": "DOUBLE DEFAULT NULL",
                },
                "roi_action_item": {
                    "approval_status": "VARCHAR(32) NOT NULL DEFAULT 'PENDING'",
                    "approval_source": "VARCHAR(32) DEFAULT NULL",
                    "approved_at": "DATETIME DEFAULT NULL",
                    "rejected_at": "DATETIME DEFAULT NULL",
                    "sheet_row_number": "INT DEFAULT NULL",
                    "result_notified_at": "DATETIME DEFAULT NULL",
                    "notification_error": "TEXT DEFAULT NULL",
                },
            }
            for table_name, columns in roi_migrations.items():
                cursor.execute(
                    """
                    SELECT COLUMN_NAME
                    FROM information_schema.COLUMNS
                    WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s
                    """,
                    (os.environ["DB_NAME"], table_name),
                )
                table_columns = {
                    row["COLUMN_NAME"] for row in cursor.fetchall()
                }
                for column, definition in columns.items():
                    if column not in table_columns:
                        cursor.execute(
                            f"ALTER TABLE {table_name} "
                            f"ADD COLUMN {column} {definition}"
                        )
            cursor.execute(
                """
                SELECT INDEX_NAME
                FROM information_schema.STATISTICS
                WHERE TABLE_SCHEMA=%s
                  AND TABLE_NAME='roi_action_item'
                  AND INDEX_NAME='idx_roi_action_approval'
                """,
                (os.environ["DB_NAME"],),
            )
            if not cursor.fetchone():
                cursor.execute(
                    """
                    CREATE INDEX idx_roi_action_approval
                    ON roi_action_item (run_id, approval_status)
                    """
                )
    finally:
        connection.close()


@contextmanager
def advisory_lock(lock_name: str) -> Iterator[bool]:
    connection = connect()
    acquired = False
    try:
        with connection.cursor() as cursor:
            cursor.execute("SELECT GET_LOCK(%s, 0) AS acquired", (lock_name,))
            acquired = bool((cursor.fetchone() or {}).get("acquired"))
        yield acquired
    finally:
        if acquired:
            try:
                with connection.cursor() as cursor:
                    cursor.execute("SELECT RELEASE_LOCK(%s)", (lock_name,))
            except Exception:
                pass
        connection.close()


def load_enabled_accounts() -> list[dict[str, Any]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT c.account_id, c.audience_name, c.bid_scene
                FROM ad_creation_account_config c
                JOIN account_whitelist w ON w.account_id = c.account_id
                WHERE c.enabled = TRUE
                  AND w.enabled = TRUE
                ORDER BY c.account_id
                """
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_automation_spend_accounts() -> list[dict[str, Any]]:
    """Load historical automation accounts that remain whitelisted."""
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT c.account_id
                FROM ad_creation_account_config c
                JOIN account_whitelist w ON w.account_id = c.account_id
                WHERE w.enabled = TRUE
                ORDER BY c.account_id
                """
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_all_spend_accounts() -> list[dict[str, Any]]:
    """Load every enabled account from the local account whitelist."""
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT w.account_id
                FROM account_whitelist w
                WHERE w.enabled = TRUE
                ORDER BY w.account_id
                """
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_realtime_accounts() -> list[dict[str, Any]]:
    """Load automation accounts plus explicit real-time scope overrides."""
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT c.account_id, c.audience_name, c.bid_scene,
                       'FULL' AS control_mode
                FROM ad_creation_account_config c
                JOIN account_whitelist w ON w.account_id = c.account_id
                WHERE w.enabled = TRUE
                ORDER BY c.account_id
                """
            )
            accounts = {
                int(row["account_id"]): row for row in cursor.fetchall()
            }
            cursor.execute(
                """
                SELECT account_id, audience_name, bid_scene, control_mode
                FROM realtime_control_account_scope
                WHERE enabled=TRUE
                ORDER BY account_id
                """
            )
            for row in cursor.fetchall():
                accounts[int(row["account_id"])] = row
            return [accounts[key] for key in sorted(accounts)]
    finally:
        connection.close()


def upsert_realtime_account_scope(
    *,
    account_id: int,
    control_mode: str,
    audience_name: str,
    bid_scene: str | None,
    source: str,
    note: str,
) -> None:
    if control_mode not in {"FULL", "PAUSE_ONLY"}:
        raise ValueError(f"Unsupported real-time control mode: {control_mode}")
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO realtime_control_account_scope
                    (account_id, control_mode, audience_name, bid_scene,
                     enabled, source, note)
                VALUES (%s,%s,%s,%s,TRUE,%s,%s)
                ON DUPLICATE KEY UPDATE
                    control_mode=VALUES(control_mode),
                    audience_name=VALUES(audience_name),
                    bid_scene=VALUES(bid_scene),
                    enabled=TRUE,
                    source=VALUES(source),
                    note=VALUES(note)
                """,
                (
                    account_id,
                    control_mode,
                    audience_name,
                    bid_scene,
                    source,
                    note,
                ),
            )
    finally:
        connection.close()


def load_managed_accounts() -> list[dict[str, Any]]:
    """Historical automation accounts still allowed by the account whitelist."""
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT c.account_id, c.audience_name, c.bid_scene, c.enabled
                FROM ad_creation_account_config c
                JOIN account_whitelist w ON w.account_id = c.account_id
                WHERE w.enabled = TRUE
                ORDER BY c.account_id
                """
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_daily_state(control_date: date) -> dict[str, Any]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM realtime_control_daily_state WHERE control_date=%s",
                (control_date,),
            )
            return cursor.fetchone() or {}
    finally:
        connection.close()


def save_daily_state(control_date: date, **values: Any) -> None:
    allowed = {
        "morning_recovery_done",
        "cutoff_done",
        "last_observed_partition",
        "last_observed_cpm",
        "last_decision",
        "last_inventory_refresh_at",
        "last_evaluated_at",
    }
    unknown = set(values) - allowed
    if unknown:
        raise ValueError(f"Unsupported daily-state fields: {sorted(unknown)}")
    columns = ["control_date", *values]
    params = [control_date, *values.values()]
    updates = ", ".join(f"{column}=VALUES({column})" for column in values)
    placeholders = ", ".join(["%s"] * len(columns))
    sql = (
        f"INSERT INTO realtime_control_daily_state ({', '.join(columns)}) "
        f"VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {updates}"
    )
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(sql, params)
    finally:
        connection.close()


def load_ad_states(account_id: int) -> dict[int, dict[str, Any]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM realtime_control_ad_state WHERE account_id=%s",
                (account_id,),
            )
            return {int(row["adgroup_id"]): row for row in cursor.fetchall()}
    finally:
        connection.close()


def upsert_ad_state(
    *,
    account_id: int,
    adgroup_id: int,
    adgroup_name: str,
    bid_field: str,
    base_bid_fen: int,
    boosted_date: date | None,
    paused_by_strategy: bool,
    pause_reason: str | None,
    last_action: str,
    action_at: datetime,
) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO realtime_control_ad_state
                    (account_id, adgroup_id, adgroup_name, bid_field, base_bid_fen,
                     initial_base_bid_fen,
                     boosted_date, paused_by_strategy, pause_reason, last_action,
                     last_action_at, last_seen_at)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                ON DUPLICATE KEY UPDATE
                    adgroup_name=VALUES(adgroup_name),
                    bid_field=VALUES(bid_field),
                    initial_base_bid_fen=COALESCE(
                        initial_base_bid_fen,
                        VALUES(initial_base_bid_fen)
                    ),
                    boosted_date=VALUES(boosted_date),
                    paused_by_strategy=VALUES(paused_by_strategy),
                    pause_reason=VALUES(pause_reason),
                    last_action=VALUES(last_action),
                    last_action_at=VALUES(last_action_at),
                    last_seen_at=VALUES(last_seen_at)
                """,
                (
                    account_id,
                    adgroup_id,
                    adgroup_name,
                    bid_field,
                    base_bid_fen,
                    base_bid_fen,
                    boosted_date,
                    paused_by_strategy,
                    pause_reason,
                    last_action,
                    action_at,
                    action_at,
                ),
            )
    finally:
        connection.close()


def sync_roi_base_bid(
    *,
    account_id: int,
    adgroup_id: int,
    adgroup_name: str,
    bid_field: str,
    initial_base_bid_fen: int,
    new_base_bid_fen: int,
    boosted_date: date | None,
    action_at: datetime,
) -> None:
    """Persist a permanent ROI base bid without losing intraday boost state."""

    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO realtime_control_ad_state
                    (account_id, adgroup_id, adgroup_name, bid_field,
                     base_bid_fen, initial_base_bid_fen, boosted_date,
                     paused_by_strategy, pause_reason, last_action,
                     last_action_at, last_seen_at, last_roi_scaled_at)
                VALUES (%s,%s,%s,%s,%s,%s,%s,FALSE,NULL,'ROI_SCALE',%s,%s,%s)
                ON DUPLICATE KEY UPDATE
                    adgroup_name=VALUES(adgroup_name),
                    bid_field=VALUES(bid_field),
                    base_bid_fen=VALUES(base_bid_fen),
                    initial_base_bid_fen=COALESCE(
                        initial_base_bid_fen,
                        VALUES(initial_base_bid_fen)
                    ),
                    boosted_date=VALUES(boosted_date),
                    last_action='ROI_SCALE',
                    last_action_at=VALUES(last_action_at),
                    last_seen_at=VALUES(last_seen_at),
                    last_roi_scaled_at=VALUES(last_roi_scaled_at)
                """,
                (
                    account_id,
                    adgroup_id,
                    adgroup_name,
                    bid_field,
                    new_base_bid_fen,
                    initial_base_bid_fen,
                    boosted_date,
                    action_at,
                    action_at,
                    action_at,
                ),
            )
    finally:
        connection.close()


def set_bid_experiment_hold(
    *,
    account_id: int,
    adgroup_id: int,
    adgroup_name: str,
    bid_field: str,
    base_bid_fen: int,
    action_at: datetime,
    reason: str,
) -> None:
    """Persist an experiment base bid and prevent real-time bid overrides."""

    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO realtime_control_ad_state
                    (account_id, adgroup_id, adgroup_name, bid_field,
                     base_bid_fen, initial_base_bid_fen, boosted_date,
                     bid_hold, bid_hold_reason, paused_by_strategy,
                     last_action, last_action_at, last_seen_at)
                VALUES (%s,%s,%s,%s,%s,%s,%s,TRUE,%s,FALSE,
                        'BID_EXPERIMENT_HOLD',%s,%s)
                ON DUPLICATE KEY UPDATE
                    adgroup_name=VALUES(adgroup_name),
                    bid_field=VALUES(bid_field),
                    base_bid_fen=VALUES(base_bid_fen),
                    initial_base_bid_fen=COALESCE(
                        initial_base_bid_fen,
                        VALUES(initial_base_bid_fen)
                    ),
                    boosted_date=VALUES(boosted_date),
                    bid_hold=TRUE,
                    bid_hold_reason=VALUES(bid_hold_reason),
                    last_action='BID_EXPERIMENT_HOLD',
                    last_action_at=VALUES(last_action_at),
                    last_seen_at=VALUES(last_seen_at)
                """,
                (
                    account_id,
                    adgroup_id,
                    adgroup_name,
                    bid_field,
                    base_bid_fen,
                    base_bid_fen,
                    action_at.date(),
                    reason,
                    action_at,
                    action_at,
                ),
            )
    finally:
        connection.close()


def clear_bid_experiment_hold(
    account_id: int,
    adgroup_id: int,
    *,
    action_at: datetime,
) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE realtime_control_ad_state
                SET bid_hold=FALSE,
                    bid_hold_reason=NULL,
                    last_action='BID_EXPERIMENT_RELEASE',
                    last_action_at=%s,
                    last_seen_at=%s
                WHERE account_id=%s AND adgroup_id=%s
                """,
                (action_at, action_at, account_id, adgroup_id),
            )
            if cursor.rowcount != 1:
                raise ValueError(
                    "未找到实验出价状态: "
                    f"account={account_id} adgroup={adgroup_id}"
                )
    finally:
        connection.close()

def insert_action_log(record: dict[str, Any]) -> None:
    columns = [
        "run_id",
        "control_date",
        "observed_partition",
        "observed_cpm",
        "decision",
        "account_id",
        "adgroup_id",
        "adgroup_name",
        "bid_field",
        "base_bid_fen",
        "before_bid_fen",
        "target_bid_fen",
        "before_status",
        "target_status",
        "apply_mode",
        "execution_status",
        "error_message",
    ]
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"INSERT INTO realtime_control_action_log ({', '.join(columns)}) "
                f"VALUES ({', '.join(['%s'] * len(columns))})",
                [record.get(column) for column in columns],
            )
    finally:
        connection.close()


def create_operator_command(record: dict[str, Any]) -> dict[str, Any]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            try:
                cursor.execute(
                    """
                    INSERT INTO operator_command
                        (command_id, source_message_id, chat_id, sender_open_id,
                         sender_name, action, scope_type, target_account_ids,
                         status, preview_account_count, preview_ad_count, expires_at)
                    VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                    """,
                    (
                        record["command_id"],
                        record["source_message_id"],
                        record["chat_id"],
                        record["sender_open_id"],
                        record.get("sender_name"),
                        record["action"],
                        record["scope_type"],
                        json.dumps(record["target_account_ids"]),
                        record["status"],
                        record.get("preview_account_count", 0),
                        record.get("preview_ad_count", 0),
                        record.get("expires_at"),
                    ),
                )
            except pymysql.err.IntegrityError:
                cursor.execute(
                    "SELECT * FROM operator_command WHERE source_message_id=%s",
                    (record["source_message_id"],),
                )
                existing = cursor.fetchone()
                if existing:
                    existing["target_account_ids"] = json.loads(
                        existing.get("target_account_ids") or "[]"
                    )
                    return existing
                raise
        return load_operator_command(record["command_id"]) or {}
    finally:
        connection.close()


def create_operator_command_with_items(
    record: dict[str, Any],
    items: list[dict[str, Any]],
) -> dict[str, Any]:
    """Atomically persist an immutable command preview and its ad snapshots."""
    connection = connect()
    try:
        connection.begin()
        with connection.cursor() as cursor:
            try:
                cursor.execute(
                    """
                    INSERT INTO operator_command
                        (command_id, source_message_id, chat_id, sender_open_id,
                         sender_name, raw_text, parse_source, intent_json,
                         action, scope_type, target_account_ids, status,
                         preview_account_count, preview_ad_count,
                         preview_cost_fen, preview_impressions, preview_clicks,
                         preview_conversions, previewed_at, resume_at, expires_at)
                    VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                    """,
                    (
                        record["command_id"], record["source_message_id"],
                        record["chat_id"], record["sender_open_id"],
                        record.get("sender_name"), record.get("raw_text"),
                        record.get("parse_source", "deterministic"),
                        json.dumps(record.get("intent") or {}, ensure_ascii=False),
                        record["action"], record["scope_type"],
                        json.dumps(record["target_account_ids"]), record["status"],
                        record.get("preview_account_count", 0),
                        record.get("preview_ad_count", 0),
                        record.get("preview_cost_fen", 0),
                        record.get("preview_impressions", 0),
                        record.get("preview_clicks", 0),
                        record.get("preview_conversions", 0),
                        record.get("previewed_at"), record.get("resume_at"),
                        record.get("expires_at"),
                    ),
                )
            except pymysql.err.IntegrityError:
                connection.rollback()
                existing = load_operator_command_by_source(record["source_message_id"])
                if existing:
                    return existing
                raise

            columns = [
                "command_id", "account_id", "audience_name", "adgroup_id",
                "adgroup_name", "before_status", "target_status",
                "preview_cost_fen", "preview_impressions", "preview_clicks",
                "preview_conversions", "previewed_at", "execution_status",
            ]
            for item in items:
                cursor.execute(
                    f"INSERT INTO operator_command_item ({', '.join(columns)}) "
                    f"VALUES ({', '.join(['%s'] * len(columns))})",
                    [item.get(column) for column in columns],
                )
        connection.commit()
        return load_operator_command(record["command_id"]) or {}
    except Exception:
        connection.rollback()
        raise
    finally:
        connection.close()


def load_operator_command_by_source(source_message_id: str) -> dict[str, Any] | None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT command_id FROM operator_command WHERE source_message_id=%s",
                (source_message_id,),
            )
            row = cursor.fetchone()
        return load_operator_command(row["command_id"]) if row else None
    finally:
        connection.close()


def load_operator_command(command_id: str) -> dict[str, Any] | None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM operator_command WHERE command_id=%s",
                (command_id,),
            )
            row = cursor.fetchone()
            if row:
                row["target_account_ids"] = json.loads(
                    row.get("target_account_ids") or "[]"
                )
                row["intent"] = json.loads(row.get("intent_json") or "{}")
            return row
    finally:
        connection.close()


def update_operator_command(command_id: str, status: str, **values: Any) -> None:
    allowed = {"confirmed_at", "executed_at", "error_message"}
    unknown = set(values) - allowed
    if unknown:
        raise ValueError(f"Unsupported operator-command fields: {sorted(unknown)}")
    assignments = ["status=%s"]
    params: list[Any] = [status]
    for column, value in values.items():
        assignments.append(f"{column}=%s")
        params.append(value)
    params.append(command_id)
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"UPDATE operator_command SET {', '.join(assignments)} "
                "WHERE command_id=%s",
                params,
            )
    finally:
        connection.close()


def transition_operator_command(
    command_id: str,
    *,
    expected_statuses: set[str],
    target_status: str,
    **values: Any,
) -> bool:
    allowed = {"confirmed_at", "executed_at", "error_message"}
    unknown = set(values) - allowed
    if unknown:
        raise ValueError(f"Unsupported operator-command fields: {sorted(unknown)}")
    if not expected_statuses:
        raise ValueError("expected_statuses must not be empty")
    assignments = ["status=%s"]
    params: list[Any] = [target_status]
    for column, value in values.items():
        assignments.append(f"{column}=%s")
        params.append(value)
    placeholders = ", ".join(["%s"] * len(expected_statuses))
    params.extend([command_id, *sorted(expected_statuses)])
    connection = connect()
    try:
        with connection.cursor() as cursor:
            affected = cursor.execute(
                f"""
                UPDATE operator_command
                SET {', '.join(assignments)}
                WHERE command_id=%s
                  AND status IN ({placeholders})
                """,
                params,
            )
            return affected == 1
    finally:
        connection.close()


def insert_operator_command_item(record: dict[str, Any]) -> None:
    columns = [
        "command_id",
        "account_id",
        "audience_name",
        "adgroup_id",
        "adgroup_name",
        "before_status",
        "target_status",
        "readback_status",
        "execution_status",
        "error_message",
    ]
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"INSERT INTO operator_command_item ({', '.join(columns)}) "
                f"VALUES ({', '.join(['%s'] * len(columns))})",
                [record.get(column) for column in columns],
            )
    finally:
        connection.close()


def load_operator_command_items(command_id: str) -> list[dict[str, Any]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT * FROM operator_command_item
                WHERE command_id=%s
                ORDER BY account_id, adgroup_id, id
                """,
                (command_id,),
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def update_operator_command_item(item_id: int, **values: Any) -> None:
    allowed = {
        "adgroup_id", "adgroup_name", "before_status", "target_status",
        "readback_status", "execution_status", "error_message",
    }
    unknown = set(values) - allowed
    if unknown:
        raise ValueError(f"Unsupported operator-command item fields: {sorted(unknown)}")
    if not values:
        return
    assignments = [f"{column}=%s" for column in values]
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"UPDATE operator_command_item SET {', '.join(assignments)} WHERE id=%s",
                [*values.values(), item_id],
            )
    finally:
        connection.close()


def find_pending_command_conflict(
    targets: list[tuple[int, int]],
    *,
    now: datetime,
) -> dict[str, Any] | None:
    if not targets:
        return None
    connection = connect()
    try:
        with connection.cursor() as cursor:
            for start in range(0, len(targets), 200):
                chunk = targets[start:start + 200]
                conditions = " OR ".join(
                    ["(i.account_id=%s AND i.adgroup_id=%s)"] * len(chunk)
                )
                params: list[Any] = []
                for account_id, adgroup_id in chunk:
                    params.extend([account_id, adgroup_id])
                params.append(now.replace(tzinfo=None))
                cursor.execute(
                    f"""
                    SELECT c.command_id, c.action, i.account_id, i.adgroup_id
                    FROM operator_command c
                    JOIN operator_command_item i ON i.command_id=c.command_id
                    WHERE ({conditions})
                      AND (
                          (c.status='PENDING_CONFIRMATION' AND c.expires_at >= %s)
                          OR c.status='EXECUTING'
                      )
                    ORDER BY c.created_at
                    LIMIT 1
                    """,
                    params,
                )
                conflict = cursor.fetchone()
                if conflict:
                    return conflict
            return None
    finally:
        connection.close()


def list_pending_operator_commands(
    chat_id: str,
    sender_open_id: str,
    now: datetime,
) -> list[dict[str, Any]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT command_id, action, preview_account_count, preview_ad_count,
                       preview_cost_fen, expires_at
                FROM operator_command
                WHERE chat_id=%s AND sender_open_id=%s
                  AND status='PENDING_CONFIRMATION' AND expires_at >= %s
                ORDER BY created_at
                """,
                (chat_id, sender_open_id, now.replace(tzinfo=None)),
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_active_operator_draft(
    chat_id: str,
    sender_open_id: str,
    now: datetime,
) -> dict[str, Any] | None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT * FROM operator_command_draft
                WHERE chat_id=%s AND sender_open_id=%s
                  AND status='ACTIVE' AND expires_at >= %s
                """,
                (chat_id, sender_open_id, now.replace(tzinfo=None)),
            )
            row = cursor.fetchone()
            if row:
                for field in ("account_ids", "missing_fields", "source_message_ids"):
                    row[field] = json.loads(row.get(field) or "[]")
            return row
    finally:
        connection.close()


def save_operator_draft(record: dict[str, Any]) -> dict[str, Any]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO operator_command_draft
                    (draft_id, chat_id, sender_open_id, raw_text, action, scope_type,
                     account_ids, missing_fields, source_message_ids, status, expires_at)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,'ACTIVE',%s)
                ON DUPLICATE KEY UPDATE
                    draft_id=VALUES(draft_id), raw_text=VALUES(raw_text),
                    action=VALUES(action),
                    scope_type=VALUES(scope_type), account_ids=VALUES(account_ids),
                    missing_fields=VALUES(missing_fields),
                    source_message_ids=VALUES(source_message_ids),
                    status='ACTIVE', expires_at=VALUES(expires_at)
                """,
                (
                    record["draft_id"], record["chat_id"], record["sender_open_id"],
                    record.get("raw_text"), record.get("action"),
                    record.get("scope_type", "MISSING"),
                    json.dumps(record.get("account_ids") or []),
                    json.dumps(record.get("missing_fields") or []),
                    json.dumps(record.get("source_message_ids") or []),
                    record["expires_at"],
                ),
            )
        return record
    finally:
        connection.close()


def close_operator_draft(chat_id: str, sender_open_id: str, status: str) -> None:
    if status not in {"COMPLETED", "CANCELLED", "SUPERSEDED", "EXPIRED"}:
        raise ValueError(f"Unsupported draft status: {status}")
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE operator_command_draft SET status=%s
                WHERE chat_id=%s AND sender_open_id=%s AND status='ACTIVE'
                """,
                (status, chat_id, sender_open_id),
            )
    finally:
        connection.close()


def set_operator_pause(
    *,
    account_id: int,
    adgroup_id: int,
    mode: str,
    resume_at: datetime | None,
    command_id: str,
    paused_from_status: str,
    paused_at: datetime,
) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE realtime_control_ad_state
                SET operator_pause_mode=%s,
                    operator_resume_at=%s,
                    operator_command_id=%s,
                    operator_paused_from_status=%s,
                    operator_paused_at=%s,
                    last_action='OPERATOR_PAUSE',
                    last_action_at=%s,
                    last_seen_at=%s
                WHERE account_id=%s AND adgroup_id=%s
                """,
                (
                    mode,
                    resume_at,
                    command_id,
                    paused_from_status,
                    paused_at,
                    paused_at,
                    paused_at,
                    account_id,
                    adgroup_id,
                ),
            )
    finally:
        connection.close()


def clear_operator_pause(
    account_id: int,
    adgroup_id: int,
    *,
    action: str,
    action_at: datetime,
) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE realtime_control_ad_state
                SET operator_pause_mode=NULL,
                    operator_resume_at=NULL,
                    operator_command_id=NULL,
                    operator_paused_from_status=NULL,
                    operator_paused_at=NULL,
                    last_action=%s,
                    last_action_at=%s,
                    last_seen_at=%s
                WHERE account_id=%s AND adgroup_id=%s
                """,
                (action, action_at, action_at, account_id, adgroup_id),
            )
    finally:
        connection.close()


def load_operator_pauses(
    account_ids: list[int] | None = None,
) -> list[dict[str, Any]]:
    params: list[Any] = []
    where = "WHERE operator_pause_mode IS NOT NULL"
    if account_ids is not None:
        if not account_ids:
            return []
        where += f" AND account_id IN ({', '.join(['%s'] * len(account_ids))})"
        params.extend(account_ids)
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                SELECT *
                FROM realtime_control_ad_state
                {where}
                ORDER BY account_id, adgroup_id
                """,
                params,
            )
            return list(cursor.fetchall())
    finally:
        connection.close()