Compressor/backend/app/services/label_service.py

from __future__ import annotations

import hashlib
import json
from datetime import datetime
from typing import Any

from ..db import get_connection

LEVELS = ("轻微", "异常", "严重异常")
SEVERITY_BY_LEVEL = {"轻微": 1, "异常": 2, "严重异常": 3}

_SELECT_FIELDS = """
    id, source_type, origin, target_type, target_id, target_key, channel_key,
    unit, device_code, device_part, device_point, metric_key, metric_label,
    start_time, end_time, algorithm_start, algorithm_end, label_type, label_subtype,
    severity, confirm_status, usable_for_train, source_file, source_row,
    algorithm_version, features_json, note
"""


def _time_string(value: Any) -> str:
    if isinstance(value, datetime):
        return value.strftime("%Y-%m-%d %H:%M:%S")
    return str(value or "")


def _parse_time(value: str) -> datetime:
    text = str(value or "").strip().replace("T", " ")
    for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M"):
        try:
            return datetime.strptime(text[:19] if fmt.endswith("%S") else text[:16], fmt)
        except ValueError:
            continue
    raise ValueError("时间格式不正确，应为 YYYY-MM-DD HH:MM:SS")


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


def _row_key(row: dict[str, Any]) -> str:
    if row.get("target_type") == "pks":
        unit = str(row.get("unit") or "")
        metric = str(row.get("metric_key") or "")
        suffix = metric.split("_")[-1] if "_" in metric else metric
        return f"YSJ{unit}_{suffix}" if unit and suffix else metric
    return str(row.get("device_point") or "")


def _row_dict(row: dict[str, Any]) -> dict[str, Any]:
    features: Any = None
    raw = row.get("features_json")
    if raw:
        try:
            features = json.loads(raw)
        except (TypeError, ValueError):
            features = raw
    return {
        "id": int(row["id"]),
        "sourceType": row.get("source_type") or "",
        "origin": row.get("origin") or "",
        "targetType": row.get("target_type") or "",
        "targetId": int(row["target_id"]) if row.get("target_id") is not None else None,
        "targetKey": row.get("target_key") or "",
        "channelKey": row.get("channel_key") or "",
        "unit": row.get("unit") or "",
        "deviceCode": row.get("device_code") or "",
        "devicePart": row.get("device_part") or "",
        "devicePoint": row.get("device_point") or "",
        "metricKey": row.get("metric_key") or "",
        "metricLabel": row.get("metric_label") or "",
        "startTime": _time_string(row.get("start_time")),
        "endTime": _time_string(row.get("end_time")),
        "labelType": row.get("label_type") or "",
        "labelSubtype": row.get("label_subtype"),
        "severity": int(row["severity"]) if row.get("severity") is not None else None,
        "confirmStatus": row.get("confirm_status") or "",
        "usableForTrain": bool(row.get("usable_for_train")),
        "sourceRow": row.get("source_row") or "",
        "features": features,
        "note": row.get("note"),
        "rowKey": _row_key(row),
    }


def _fetch_label(cursor: Any, label_id: int) -> dict[str, Any]:
    cursor.execute(f"SELECT {_SELECT_FIELDS} FROM compressor_label WHERE id = %s", (label_id,))
    row = cursor.fetchone()
    if not row:
        raise ValueError("标注不存在")
    return _row_dict(row)


def query_labels(
    target_type: str | None = None,
    unit: str | None = None,
    device_code: str | None = None,
    metric_keys: list[str] | None = None,
    channel_keys: list[str] | None = None,
    wave_file_ids: list[int] | None = None,
    start: str | None = None,
    end: str | None = None,
) -> dict[str, Any]:
    clauses = ["deleted_at IS NULL"]
    params: list[Any] = []
    if target_type:
        clauses.append("target_type = %s")
        params.append(target_type)
    if unit:
        clauses.append("unit = %s")
        params.append(unit)
    if device_code:
        clauses.append("device_code = %s")
        params.append(device_code)
    if metric_keys:
        placeholders = ", ".join(["%s"] * len(metric_keys))
        clauses.append(f"metric_key IN ({placeholders})")
        params.extend(metric_keys)
    if channel_keys:
        placeholders = ", ".join(["%s"] * len(channel_keys))
        clauses.append(f"channel_key IN ({placeholders})")
        params.extend(channel_keys)
    if wave_file_ids:
        placeholders = ", ".join(["%s"] * len(wave_file_ids))
        clauses.append(f"target_id IN ({placeholders})")
        params.extend(wave_file_ids)
    if start:
        clauses.append("end_time > %s")
        params.append(start)
    if end:
        clauses.append("start_time < %s")
        params.append(end)

    sql = f"""
        SELECT {_SELECT_FIELDS}
        FROM compressor_label
        WHERE {' AND '.join(clauses)}
        ORDER BY channel_key ASC, start_time ASC, id ASC
    """
    with get_connection() as connection:
        with connection.cursor() as cursor:
            cursor.execute(sql, params)
            rows = cursor.fetchall()
    return {
        "source": "database",
        "labels": [_row_dict(row) for row in rows],
        "notice": None,
    }


def _insert_row(cursor: Any, row: dict[str, Any]) -> int:
    fields = (
        "source_type", "origin", "target_type", "target_id", "target_key", "channel_key",
        "segment_key", "unit", "device_code", "device_part", "device_point", "metric_key",
        "metric_label", "start_time", "end_time", "algorithm_start", "algorithm_end",
        "label_type", "label_subtype", "severity", "confirm_status", "usable_for_train",
        "source_file", "source_row", "algorithm_version", "features_json", "note",
    )
    placeholders = ", ".join(["%s"] * len(fields))
    cursor.execute(
        f"INSERT INTO compressor_label ({', '.join(fields)}) VALUES ({placeholders})",
        tuple(row.get(field) for field in fields),
    )
    return int(cursor.lastrowid)


def _trim_algorithm_row(cursor: Any, row: dict[str, Any], start: datetime, end: datetime) -> None:
    row_start = _parse_time(_time_string(row["start_time"]))
    row_end = _parse_time(_time_string(row["end_time"]))
    if row_end <= start or row_start >= end:
        return
    keep_before = row_start < start
    keep_after = row_end > end
    if keep_before and keep_after:
        cursor.execute(
            "UPDATE compressor_label SET end_time = %s WHERE id = %s",
            (_time_string(start), row["id"]),
        )
        clone = dict(row)
        clone["start_time"] = _time_string(end)
        clone["end_time"] = _time_string(row_end)
        clone["algorithm_start"] = row.get("algorithm_start") or row["start_time"]
        clone["algorithm_end"] = row.get("algorithm_end") or row["end_time"]
        base = str(row.get("segment_key") or row["id"])
        clone["segment_key"] = hashlib.md5(f"{base}|split|{_time_string(end)}".encode("utf-8")).hexdigest()
        clone["origin"] = "algorithm"
        _insert_row(cursor, clone)
    elif keep_before:
        cursor.execute(
            "UPDATE compressor_label SET end_time = %s WHERE id = %s",
            (_time_string(start), row["id"]),
        )
    elif keep_after:
        cursor.execute(
            "UPDATE compressor_label SET start_time = %s WHERE id = %s",
            (_time_string(end), row["id"]),
        )
    else:
        cursor.execute(
            "UPDATE compressor_label SET deleted_at = NOW() WHERE id = %s",
            (row["id"],),
        )


def _apply_manual(
    cursor: Any,
    payload: dict[str, Any],
    exclude_id: int | None = None,
    carry: dict[str, Any] | None = None,
) -> int:
    label_type = str(payload.get("label_type") or "").strip()
    if label_type not in LEVELS:
        raise ValueError("标签类型只能是 轻微 / 异常 / 严重异常")
    start = _snap_hour(_parse_time(payload["start_time"]))
    end = _snap_hour(_parse_time(payload["end_time"]))
    if end <= start:
        raise ValueError("结束时间必须大于开始时间")
    channel_key = str(payload.get("channel_key") or "").strip()
    if not channel_key:
        raise ValueError("channel_key 不能为空")

    cursor.execute(
        f"""
        SELECT {_SELECT_FIELDS}
        FROM compressor_label
        WHERE channel_key = %s
          AND deleted_at IS NULL
          AND start_time <= %s
          AND end_time >= %s
          AND (%s IS NULL OR id <> %s)
        """,
        (channel_key, _time_string(end), _time_string(start), exclude_id, exclude_id),
    )
    rows = cursor.fetchall()

    for row in rows:
        if row["origin"] == "manual" and row["label_type"] != label_type:
            raise ValueError("同一子图内不同类型不能重叠或相邻，请先修改或删除相邻标注")

    same = [row for row in rows if row["label_type"] == label_type]
    merged_start = min([start] + [_parse_time(_time_string(row["start_time"])) for row in same])
    merged_end = max([end] + [_parse_time(_time_string(row["end_time"])) for row in same])

    for row in rows:
        if row["label_type"] != label_type and row["origin"] != "manual":
            _trim_algorithm_row(cursor, row, merged_start, merged_end)

    same_ids = [row["id"] for row in same]
    if same_ids:
        placeholders = ", ".join(["%s"] * len(same_ids))
        cursor.execute(
            f"UPDATE compressor_label SET deleted_at = NOW() WHERE id IN ({placeholders})",
            same_ids,
        )

    carry = carry or {}
    source_type = "pks_manual" if payload.get("target_type") == "pks" else "wave_manual"
    new_row = {
        "source_type": source_type,
        "origin": "manual",
        "target_type": payload.get("target_type") or ("pks" if payload.get("metric_key") else "wave_file"),
        "target_id": payload.get("target_id"),
        "target_key": channel_key,
        "channel_key": channel_key,
        "segment_key": None,
        "unit": str(payload.get("unit") or carry.get("unit") or ""),
        "device_code": payload.get("device_code") or carry.get("device_code"),
        "device_part": payload.get("device_part") or carry.get("device_part"),
        "device_point": payload.get("device_point") or carry.get("device_point"),
        "metric_key": payload.get("metric_key") or carry.get("metric_key"),
        "metric_label": payload.get("metric_label") or carry.get("metric_label"),
        "start_time": _time_string(merged_start),
        "end_time": _time_string(merged_end),
        "algorithm_start": carry.get("algorithm_start"),
        "algorithm_end": carry.get("algorithm_end"),
        "label_type": label_type,
        "label_subtype": payload.get("label_subtype") or carry.get("label_subtype"),
        "severity": SEVERITY_BY_LEVEL.get(label_type),
        "confirm_status": "人工确认",
        "usable_for_train": 1,
        "source_file": carry.get("source_file"),
        "source_row": carry.get("source_row"),
        "algorithm_version": carry.get("algorithm_version"),
        "features_json": carry.get("features_json"),
        "note": payload.get("note"),
    }
    return _insert_row(cursor, new_row)


def create_label(payload: dict[str, Any]) -> dict[str, Any]:
    with get_connection() as connection:
        with connection.cursor() as cursor:
            connection.begin()
            try:
                label_id = _apply_manual(cursor, payload)
                result = _fetch_label(cursor, label_id)
                connection.commit()
                return result
            except Exception:
                connection.rollback()
                raise


def update_label(label_id: int, payload: dict[str, Any]) -> dict[str, Any]:
    if label_id <= 0:
        raise ValueError("标注 id 必须为正整数")
    with get_connection() as connection:
        with connection.cursor() as cursor:
            connection.begin()
            try:
                cursor.execute(
                    f"SELECT {_SELECT_FIELDS} FROM compressor_label WHERE id = %s AND deleted_at IS NULL",
                    (label_id,),
                )
                old = cursor.fetchone()
                if not old:
                    raise ValueError("标注不存在")
                new_id = _apply_manual(cursor, payload, exclude_id=label_id, carry=old)
                result = _fetch_label(cursor, new_id)
                if old["origin"] == "manual":
                    cursor.execute(
                        "UPDATE compressor_label SET deleted_at = NOW() WHERE id = %s",
                        (label_id,),
                    )
                else:
                    _trim_algorithm_row(
                        cursor,
                        old,
                        _parse_time(result["startTime"]),
                        _parse_time(result["endTime"]),
                    )
                connection.commit()
                return result
            except Exception:
                connection.rollback()
                raise


def delete_label(label_id: int) -> dict[str, Any]:
    if label_id <= 0:
        raise ValueError("标注 id 必须为正整数")
    with get_connection() as connection:
        with connection.cursor() as cursor:
            cursor.execute(
                "UPDATE compressor_label SET deleted_at = NOW() WHERE id = %s AND deleted_at IS NULL",
                (label_id,),
            )
            if cursor.rowcount == 0:
                raise ValueError("标注不存在")
    return {"deleted": label_id}