Compressor/models/common/alarm_store.py

"""Persist pressure / vibration stage segments into ``compressor_alarm``.

Two write modes:

* ``write_full``  – full scan: delete previous pressure/vibration alarms for the
  selected devices, then insert every segment (vibration writes one row per
  side: YSJ_10 and YSJ_11).
* ``merge_realtime`` – realtime tick: a growing segment is merged into the
  existing row with the same identity (device + part + point + type + level)
  when their time ranges overlap or are within ``MERGE_GAP_HOURS``.  Different
  levels are kept as separate rows, matching the stage CSV rows, so a continuous
  same-level fault is never split while a level transition is preserved.
"""
from __future__ import annotations

import json
from datetime import datetime, timedelta
from typing import Any

from .db import connect

OIL_ALARM_TYPE = "润滑油压力故障"
VIB_ALARM_TYPE = "振动故障"
OIL_PART = "润滑油压力"
OIL_POINT = "YSJ_5"
COUPLING_PART = "联轴器端振动"
COUPLING_POINT = "YSJ_10"
CHAIN_PART = "链轮端振动"
CHAIN_POINT = "YSJ_11"
ALARM_STATUS = 1
MERGE_GAP_HOURS = 6
DES_MAX = 4000

INSERT_SQL = """
INSERT INTO compressor_alarm
    (device_code, device_part, device_point, alarm_type, alarm_des,
     alarm_time_start, alarm_time_end, status, alarm_level,
     scan_start_time, scan_end_time, alarm_info_json)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
"""


def _as_dt(value) -> datetime:
    if isinstance(value, datetime):
        return value
    return datetime.strptime(str(value), "%Y-%m-%d %H:%M:%S")


def _json(row: dict) -> str:
    # The explanatory note goes to alarm_des, not to alarm_info_json.
    data = {key: value for key, value in row.items() if not key.startswith("_") and key != "模型预测参考"}
    return json.dumps(data, ensure_ascii=False, separators=(",", ":"), default=str)


def _targets(device_code: str, kind: str, row: dict) -> list[dict]:
    level = row.get("压力阶段") if kind == "pressure" else row.get("振动阶段")
    start, end = row.get("开始小时"), row.get("结束小时")
    des = str(row.get("阶段触发依据") or "")
    note = row.get("模型预测参考")
    if note:
        des = f"{des}；{note}"
    des = des[:DES_MAX]
    info = _json(row)
    if kind == "pressure":
        parts = [(OIL_PART, OIL_POINT, OIL_ALARM_TYPE)]
    else:
        # One row per segment, at the strongest side (matches 振动异常段.csv).
        side = str(row.get("最强异常侧") or "")
        if side == "联轴器端":
            parts = [(COUPLING_PART, COUPLING_POINT, VIB_ALARM_TYPE)]
        else:
            parts = [(CHAIN_PART, CHAIN_POINT, VIB_ALARM_TYPE)]
    result = []
    for part, point, alarm_type in parts:
        result.append({
            "device_code": device_code,
            "device_part": part,
            "device_point": point,
            "alarm_type": alarm_type,
            "alarm_level": level,
            "alarm_des": des,
            "alarm_time_start": _as_dt(start),
            "alarm_time_end": _as_dt(end),
            "alarm_info_json": info,
        })
    return result


def build_alarms(pressure_rows: list[dict], vibration_rows: list[dict], batch_to_device: dict[int, str]) -> list[dict]:
    alarms = []
    for row in pressure_rows:
        device = batch_to_device.get(int(row["批次"]))
        if device:
            alarms.extend(_targets(device, "pressure", row))
    for row in vibration_rows:
        device = batch_to_device.get(int(row["批次"]))
        if device:
            alarms.extend(_targets(device, "vibration", row))
    return alarms


def _insert(cursor, alarm: dict) -> int:
    cursor.execute(INSERT_SQL, (
        alarm["device_code"], alarm["device_part"], alarm["device_point"], alarm["alarm_type"],
        alarm["alarm_des"], alarm["alarm_time_start"], alarm["alarm_time_end"], ALARM_STATUS,
        alarm["alarm_level"], alarm["scan_start_time"], alarm["scan_end_time"], alarm["alarm_info_json"],
    ))
    return int(cursor.lastrowid)


def write_full(alarms: list[dict], devices: list[str], scan_start: datetime, scan_end: datetime, dry_run: bool = False) -> dict:
    for alarm in alarms:
        alarm["scan_start_time"] = scan_start
        alarm["scan_end_time"] = scan_end
    if dry_run:
        return {"mode": "full", "dry_run": True, "delete_scope": [devices, OIL_ALARM_TYPE, VIB_ALARM_TYPE], "inserted": len(alarms)}
    connection = connect()
    deleted = 0
    inserted_ids = []
    try:
        with connection.cursor() as cursor:
            placeholders = ",".join(["%s"] * len(devices))
            cursor.execute(
                f"DELETE FROM compressor_alarm WHERE device_code IN ({placeholders}) AND alarm_type IN (%s, %s)",
                (*devices, OIL_ALARM_TYPE, VIB_ALARM_TYPE),
            )
            deleted = int(cursor.rowcount)
            for alarm in alarms:
                inserted_ids.append(_insert(cursor, alarm))
    finally:
        connection.close()
    return {"mode": "full", "deleted": deleted, "inserted": len(inserted_ids), "ids": inserted_ids}


def merge_realtime(alarms: list[dict], scan_start: datetime, scan_end: datetime, journal: list | None = None, dry_run: bool = False) -> dict:
    """Merge each segment into the matching existing row (same identity+level)."""
    inserted = updated = merged_extra = 0
    if dry_run:
        return {"mode": "realtime", "dry_run": True, "segments": len(alarms)}
    connection = connect()
    gap = timedelta(hours=MERGE_GAP_HOURS)
    try:
        with connection.cursor() as cursor:
            for alarm in alarms:
                start, end = alarm["alarm_time_start"], alarm["alarm_time_end"]
                cursor.execute(
                    # Identity ignores part/point so a segment whose strongest side
                    # flips (coupling <-> chain) still merges into one row instead
                    # of being split.  part/point are refreshed to the latest side.
                    """
                    SELECT id, alarm_time_start, alarm_time_end, alarm_des, alarm_info_json,
                           scan_start_time, scan_end_time, status
                    FROM compressor_alarm
                    WHERE device_code = %s AND alarm_type = %s AND alarm_level <=> %s
                      AND alarm_time_start <= %s AND alarm_time_end >= %s
                    ORDER BY alarm_time_start ASC, id ASC
                    """,
                    (
                        alarm["device_code"], alarm["alarm_type"],
                        alarm["alarm_level"], end + gap, start - gap,
                    ),
                )
                existing = cursor.fetchall()
                if not existing:
                    new_id = _insert(cursor, {**alarm, "scan_start_time": scan_start, "scan_end_time": scan_end})
                    inserted += 1
                    if journal is not None:
                        journal.append({"op": "insert", "id": new_id})
                    continue
                keep = existing[0]
                new_start = min(start, *(row["alarm_time_start"] for row in existing))
                new_end = max(end, *(row["alarm_time_end"] for row in existing))
                new_scan_start = min([scan_start] + [row["scan_start_time"] for row in existing if row["scan_start_time"]])
                new_scan_end = max([scan_end] + [row["scan_end_time"] for row in existing if row["scan_end_time"]])
                if journal is not None:
                    journal.append({
                        "op": "update", "id": int(keep["id"]),
                        "alarm_time_start": keep["alarm_time_start"], "alarm_time_end": keep["alarm_time_end"],
                        "alarm_des": keep["alarm_des"], "alarm_info_json": keep["alarm_info_json"],
                        "scan_start_time": keep["scan_start_time"], "scan_end_time": keep["scan_end_time"],
                        "status": keep["status"],
                    })
                if len(existing) > 1:
                    extra_ids = [int(row["id"]) for row in existing[1:]]
                    placeholders = ",".join(["%s"] * len(extra_ids))
                    cursor.execute(f"DELETE FROM compressor_alarm WHERE id IN ({placeholders})", extra_ids)
                    merged_extra += len(extra_ids)
                    if journal is not None:
                        for row in existing[1:]:
                            journal.append({"op": "delete_restore", "row": dict(row)})
                cursor.execute(
                    """
                    UPDATE compressor_alarm
                    SET device_part = %s, device_point = %s,
                        alarm_time_start = %s, alarm_time_end = %s, alarm_des = %s,
                        scan_start_time = %s, scan_end_time = %s, status = %s,
                        alarm_info_json = %s
                    WHERE id = %s
                    """,
                    (alarm["device_part"], alarm["device_point"],
                     new_start, new_end, alarm["alarm_des"], new_scan_start, new_scan_end,
                     ALARM_STATUS, alarm["alarm_info_json"], keep["id"]),
                )
                updated += 1
    finally:
        connection.close()
    return {"mode": "realtime", "inserted": inserted, "updated": updated, "merged_extra": merged_extra}


def rollback(journal: list) -> dict:
    """Undo a realtime smoke run using the journal produced by ``merge_realtime``."""
    if not journal:
        return {"undone": 0}
    connection = connect()
    undone = 0
    try:
        with connection.cursor() as cursor:
            for entry in reversed(journal):
                if entry["op"] == "insert":
                    cursor.execute("DELETE FROM compressor_alarm WHERE id = %s", (entry["id"],))
                    undone += 1
                elif entry["op"] == "update":
                    cursor.execute(
                        """
                        UPDATE compressor_alarm
                        SET alarm_time_start = %s, alarm_time_end = %s, alarm_des = %s,
                            alarm_info_json = %s, scan_start_time = %s, scan_end_time = %s, status = %s
                        WHERE id = %s
                        """,
                        (entry["alarm_time_start"], entry["alarm_time_end"], entry["alarm_des"],
                         entry["alarm_info_json"], entry["scan_start_time"], entry["scan_end_time"],
                         entry["status"], entry["id"]),
                    )
                    undone += 1
                elif entry["op"] == "delete_restore":
                    row = entry["row"]
                    cursor.execute(INSERT_SQL, (
                        row["device_code"], row["device_part"], row["device_point"], row["alarm_type"],
                        row["alarm_des"], row["alarm_time_start"], row["alarm_time_end"], row["status"],
                        row["alarm_level"] if "alarm_level" in row else None,
                        row["scan_start_time"], row["scan_end_time"], row["alarm_info_json"],
                    ))
                    undone += 1
    finally:
        connection.close()
    return {"undone": undone}


def fetch_alarms(devices: list[str]) -> list[dict]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            placeholders = ",".join(["%s"] * len(devices))
            cursor.execute(
                f"""
                SELECT id, device_code, device_part, device_point, alarm_type, alarm_level,
                       alarm_time_start, alarm_time_end, status, alarm_des, alarm_info_json
                FROM compressor_alarm
                WHERE device_code IN ({placeholders}) AND alarm_type IN (%s, %s)
                ORDER BY device_code, alarm_time_start, device_part
                """,
                (*devices, OIL_ALARM_TYPE, VIB_ALARM_TYPE),
            )
            return [dict(row) for row in cursor.fetchall()]
    finally:
        connection.close()


__all__ = [
    "OIL_ALARM_TYPE", "VIB_ALARM_TYPE", "OIL_PART", "OIL_POINT",
    "COUPLING_PART", "COUPLING_POINT", "CHAIN_PART", "CHAIN_POINT",
    "MERGE_GAP_HOURS", "build_alarms", "write_full", "merge_realtime",
    "rollback", "fetch_alarms",
]