Compressor/models/common/db.py

"""Database connection, PKS batch mapping and alarm-threshold configuration."""
from __future__ import annotations

import re

import pymysql

from .settings import settings

SITE_POINT_COLUMNS = (
    "ItemName, ItemDescription, Units, "
    "AlarmType1, AlarmType2, AlarmType3, AlarmType4, "
    "AlarmLimit1, AlarmLimit2, AlarmLimit3, AlarmLimit4"
)


def connect():
    return pymysql.connect(
        host=settings.db_host,
        port=settings.db_port,
        user=settings.db_user,
        password=settings.db_password,
        database=settings.db_name,
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
        connect_timeout=settings.db_connect_timeout,
        read_timeout=600,
        write_timeout=60,
        autocommit=True,
    )


def number(value):
    """Coerce a value to a finite float or ``None``."""
    import math

    if value is None:
        return None
    try:
        value = float(value)
    except (TypeError, ValueError):
        return None
    return value if math.isfinite(value) else None


def load_pks_batches() -> dict[int, str]:
    """Resolve PKS import batches to units from the import metadata table."""
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT id, directory_path
                FROM compressor_data_import
                WHERE data_type = %s
                  AND directory_path IS NOT NULL
                ORDER BY id
                """,
                ("PKS",),
            )
            result = {}
            for row in cursor.fetchall():
                match = re.search(r"(?<!\d)([789])\s*#", str(row["directory_path"]))
                if match:
                    result[int(row["id"])] = f"{match.group(1)}号机"
    finally:
        connection.close()
    if not result:
        raise RuntimeError("compressor_data_import 中没有可识别 7#/8#/9# 的 PKS 导入批次")
    return result


def alarm_entry(item: dict, alarm_type: str) -> tuple[float | None, str]:
    for index in range(1, 5):
        configured_type = item.get(f"AlarmType{index}")
        if configured_type is not None and str(configured_type).strip() == alarm_type:
            return number(item.get(f"AlarmLimit{index}")), f"AlarmType{index}/AlarmLimit{index}"
    return None, ""


def load_alarm_config(batch_to_unit: dict[int, str]) -> dict[int, dict]:
    """Read alarm limits by AlarmType, because AlarmLimit positions differ."""
    result = {}
    connection = connect()
    try:
        with connection.cursor() as cursor:
            unit_numbers = sorted({unit[0] for unit in batch_to_unit.values()})
            names = [f"YSJ{unit}_{point}" for unit in unit_numbers for point in (5, 10, 11)]
            placeholders = ",".join(["%s"] * len(names))
            cursor.execute(
                f"SELECT {SITE_POINT_COLUMNS} FROM site_point WHERE ItemName IN ({placeholders})",
                names,
            )
            items = {row["ItemName"]: row for row in cursor.fetchall()}
    finally:
        connection.close()

    unit_configs = {}
    for unit in sorted(set(batch_to_unit.values())):
        unit_number = unit[0]
        oil = items.get(f"YSJ{unit_number}_5")
        coupling = items.get(f"YSJ{unit_number}_10")
        chain = items.get(f"YSJ{unit_number}_11")
        if not oil or not coupling or not chain:
            raise RuntimeError(f"site_point 缺少 {unit} 的 YSJ_5/10/11 配置")
        oil_entries = {
            "low": alarm_entry(oil, "PVLow"),
            "low_low": alarm_entry(oil, "PVLowLow"),
            "high": alarm_entry(oil, "PVHigh"),
            "high_high": alarm_entry(oil, "PVHighHigh"),
        }
        coupling_entries = {
            "high": alarm_entry(coupling, "PVHigh"),
            "high_high": alarm_entry(coupling, "PVHighHigh"),
        }
        chain_entries = {
            "high": alarm_entry(chain, "PVHigh"),
            "high_high": alarm_entry(chain, "PVHighHigh"),
        }
        oil_limits = {key: value for key, (value, _source) in oil_entries.items()}
        coupling_limits = {key: value for key, (value, _source) in coupling_entries.items()}
        chain_limits = {key: value for key, (value, _source) in chain_entries.items()}
        if any(value is None for value in oil_limits.values()):
            raise RuntimeError(f"{unit} YSJ_5 缺少完整低压报警阈值")
        if any(value is None for value in (*coupling_limits.values(), *chain_limits.values())):
            raise RuntimeError(f"{unit} YSJ_10/11 缺少完整振动报警阈值")
        unit_configs[unit] = {
            "unit": unit,
            "oil": oil_limits,
            "coupling": coupling_limits,
            "chain": chain_limits,
            "alarm_sources": {
                "oil": {key: source for key, (_value, source) in oil_entries.items()},
                "coupling": {key: source for key, (_value, source) in coupling_entries.items()},
                "chain": {key: source for key, (_value, source) in chain_entries.items()},
            },
            "metadata": {
                "oil_description": oil.get("ItemDescription") or "",
                "oil_units": oil.get("Units") or "",
                "coupling_description": coupling.get("ItemDescription") or "",
                "coupling_units": coupling.get("Units") or "",
                "chain_description": chain.get("ItemDescription") or "",
                "chain_units": chain.get("Units") or "",
            },
        }
    for batch, unit in batch_to_unit.items():
        result[batch] = {**unit_configs[unit], "batch": batch}
    return result


__all__ = [
    "SITE_POINT_COLUMNS", "connect", "number", "load_pks_batches",
    "alarm_entry", "load_alarm_config",
]