Compressor/models/common/devices.py

"""Device-code resolution for PKS scans.

``--device`` accepts compressor codes such as ``7#`` / ``8#`` / ``9#`` (also
tolerates ``7`` or ``7号机``).  A device maps to one or more PKS import batches
through ``compressor_data_import.directory_path``; the batch ids are then used
for the (index-friendly) hourly aggregation so results stay identical to the
original batch-based scanner output.

4#/5#/6# exist in the database but are intentionally not part of the default
scan set.
"""
from __future__ import annotations

import re
from datetime import datetime, timedelta

from .db import connect

DEVICE_CODES = ("7#", "8#", "9#")
DEVICE_TO_UNIT = {"7#": "7号机", "8#": "8号机", "9#": "9号机"}


def normalize_device(value) -> str:
    text = str(value).strip()
    if text in DEVICE_TO_UNIT:
        return text
    digits = re.sub(r"\D", "", text)
    candidate = f"{digits}#" if digits else ""
    if candidate in DEVICE_TO_UNIT:
        return candidate
    raise ValueError(f"无法识别的压缩机编号：{value!r}（应为 7# / 8# / 9#）")


def normalize_devices(values=None) -> list[str]:
    devices = []
    for value in (values or DEVICE_CODES):
        device = normalize_device(value)
        if device not in devices:
            devices.append(device)
    return devices


def load_device_batches(devices=None) -> dict[str, list[int]]:
    """Resolve ``device -> [import_batch_id, ...]`` from the import metadata."""
    selected = normalize_devices(devices)
    placeholders = ",".join(["%s"] * len(selected))
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                SELECT id, directory_path
                FROM compressor_data_import
                WHERE data_type = 'PKS'
                  AND directory_path IN ({placeholders})
                ORDER BY id
                """,
                selected,
            )
            result: dict[str, list[int]] = {device: [] for device in selected}
            for row in cursor.fetchall():
                device = str(row["directory_path"]).strip()
                if device in result:
                    result[device].append(int(row["id"]))
    finally:
        connection.close()
    missing = [device for device, batches in result.items() if not batches]
    if missing:
        raise RuntimeError(f"以下设备没有找到 PKS 导入批次：{missing}")
    return result


def batch_to_device(device_batches: dict[str, list[int]]) -> dict[int, str]:
    return {batch: device for device, batches in device_batches.items() for batch in batches}


def batch_to_unit(device_batches: dict[str, list[int]]) -> dict[int, str]:
    return {batch: DEVICE_TO_UNIT[device] for device, batches in device_batches.items() for batch in batches}


def floor_hour(moment: datetime) -> datetime:
    return moment.replace(minute=0, second=0, microsecond=0)


def device_data_range(device_batches: dict[str, list[int]]) -> dict[str, tuple[datetime, datetime]]:
    """Return ``device -> (min_sample_time, max_sample_time)`` using the batch index."""
    ranges: dict[str, tuple[datetime, datetime]] = {}
    connection = connect()
    try:
        with connection.cursor() as cursor:
            for device, batches in device_batches.items():
                placeholders = ",".join(["%s"] * len(batches))
                cursor.execute(
                    f"""
                    SELECT MIN(sample_time) AS mn, MAX(sample_time) AS mx
                    FROM pks_long_sample
                    WHERE import_batch_id IN ({placeholders})
                    """,
                    batches,
                )
                row = cursor.fetchone()
                if row and row["mn"] and row["mx"]:
                    ranges[device] = (row["mn"], row["mx"])
    finally:
        connection.close()
    if not ranges:
        raise RuntimeError("所选设备没有可用样本时间范围")
    return ranges


def auto_range(device_batches: dict[str, list[int]]) -> tuple[datetime, datetime]:
    """Settled scan range: floor(min) .. floor(max)+1h (last complete hour included)."""
    ranges = device_data_range(device_batches)
    start = min(floor_hour(mn) for mn, _ in ranges.values())
    end = max(floor_hour(mx) + timedelta(hours=1) for _, mx in ranges.values())
    return start, end


__all__ = [
    "DEVICE_CODES", "DEVICE_TO_UNIT", "normalize_device", "normalize_devices",
    "load_device_batches", "batch_to_device", "batch_to_unit", "floor_hour",
    "device_data_range", "auto_range",
]