Compressor/backend/app/services/label_import_service.py

from __future__ import annotations

import copy
import csv
import hashlib
import json
import re
import threading
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Callable

from ..db import ensure_label_table, get_connection

ROOT = Path(__file__).resolve().parents[3]

LABEL_FILES = (
    {
        "key": "vibration",
        "label": "振动异常段",
        "relativePath": "models/datas/reference/振动异常段.csv",
    },
    {
        "key": "oil_pressure",
        "label": "润滑油压力小时阶段段",
        "relativePath": "models/datas/reference/润滑油压力小时阶段段.csv",
    },
    {
        "key": "temperature",
        "label": "温度异常段",
        "relativePath": "_tmp_temperature/out/温度异常段.csv",
    },
)

SEVERITY = {"轻微": 1, "异常": 2, "严重异常": 3}
VIBRATION_SIDE_TO_METRIC = {"联轴器端": "YSJ_10", "链轮端": "YSJ_11"}

INSERT_SQL = """
INSERT INTO compressor_label (
    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
) VALUES (
    %s, %s, %s, %s, %s, %s, %s,
    %s, %s, %s, %s, %s, %s,
    %s, %s, %s, %s,
    %s, %s, %s, %s, %s,
    %s, %s, %s, %s
)
ON DUPLICATE KEY UPDATE
    label_type = IF(compressor_label.origin = 'manual', compressor_label.label_type, VALUES(label_type)),
    label_subtype = IF(compressor_label.origin = 'manual', compressor_label.label_subtype, VALUES(label_subtype)),
    severity = IF(compressor_label.origin = 'manual', compressor_label.severity, VALUES(severity)),
    confirm_status = IF(compressor_label.origin = 'manual', compressor_label.confirm_status, VALUES(confirm_status)),
    usable_for_train = IF(compressor_label.origin = 'manual', compressor_label.usable_for_train, VALUES(usable_for_train)),
    features_json = IF(compressor_label.origin = 'manual', compressor_label.features_json, VALUES(features_json)),
    updated_at = CURRENT_TIMESTAMP
"""

ProgressCallback = Callable[[int, str], None]


def _now() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def read_csv(path: Path) -> list[dict[str, str]]:
    if not path.exists():
        return []
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        return list(csv.DictReader(handle))


def parse_unit(raw: str) -> str | None:
    match = re.match(r"\s*(\d+)\s*号机", raw or "")
    return match.group(1) if match else None


def parse_time(raw: str) -> datetime | None:
    text = (raw or "").strip()
    if not text:
        return None
    for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M", "%Y-%m-%d"):
        try:
            return datetime.strptime(text, fmt)
        except ValueError:
            continue
    return None


def time_text(value: datetime) -> str:
    return value.strftime("%Y-%m-%d %H:%M:%S")


def metric_label(site_points: dict[str, str], unit: str, metric: str) -> str:
    return site_points.get(f"YSJ{unit}_{metric.split('_')[-1]}", metric)


def confirm_status(row: dict[str, str]) -> str:
    return "已确认" if (row.get("是否形成确认等级") or "").strip() == "是" else "未确认"


def make_segment(
    *,
    source_type: str,
    unit: str,
    metric: str,
    start: datetime,
    end: datetime,
    label_type: str,
    label_subtype: str | None,
    row: dict[str, str],
    source_file: str,
    site_points: dict[str, str],
) -> tuple[Any, ...]:
    device_code = f"{unit}#"
    target_key = f"pks:{unit}:{device_code}:{metric}"
    source_row = (row.get("周期编号") or "").strip()
    segment_key = hashlib.md5(
        f"{source_type}|{unit}|{metric}|{source_row}|{time_text(start)}|{time_text(end)}".encode("utf-8")
    ).hexdigest()
    label = metric_label(site_points, unit, metric)
    return (
        source_type,
        "algorithm",
        "pks",
        None,
        target_key,
        target_key,
        segment_key,
        unit,
        device_code,
        (row.get("机组") or "").strip(),
        label,
        metric,
        label,
        time_text(start),
        time_text(end),
        time_text(start),
        time_text(end),
        label_type,
        label_subtype,
        SEVERITY.get(label_type),
        confirm_status(row),
        1,
        source_file,
        source_row,
        "csv-v1",
        json.dumps(row, ensure_ascii=False),
    )


def build_rows_for_file(source_type: str, path: Path, site_points: dict[str, str]) -> list[tuple[Any, ...]]:
    rows: list[tuple[Any, ...]] = []
    source_file = str(path.relative_to(ROOT))
    for row in read_csv(path):
        unit = parse_unit(row.get("机组", ""))
        start = parse_time(row.get("开始小时", ""))
        end = parse_time(row.get("结束小时", ""))
        if not unit or not start or not end:
            continue
        if source_type == "vibration":
            side = (row.get("最强异常侧") or "").strip()
            metric = VIBRATION_SIDE_TO_METRIC.get(side)
            if not metric:
                continue
            rows.append(
                make_segment(
                    source_type=source_type,
                    unit=unit,
                    metric=metric,
                    start=start,
                    end=end + timedelta(hours=1),
                    label_type=(row.get("振动阶段") or "").strip(),
                    label_subtype=side or None,
                    row=row,
                    source_file=source_file,
                    site_points=site_points,
                )
            )
        elif source_type == "oil_pressure":
            metric = "YSJ_5" if unit in {"7", "8", "9"} else "YSJ_1"
            rows.append(
                make_segment(
                    source_type=source_type,
                    unit=unit,
                    metric=metric,
                    start=start,
                    end=end + timedelta(hours=1),
                    label_type=(row.get("压力阶段") or "").strip(),
                    label_subtype="润滑油压力",
                    row=row,
                    source_file=source_file,
                    site_points=site_points,
                )
            )
        elif source_type == "temperature":
            metric = (row.get("阶段点位") or "").strip()
            if not re.fullmatch(r"YSJ_\d+", metric):
                continue
            rows.append(
                make_segment(
                    source_type=source_type,
                    unit=unit,
                    metric=metric,
                    start=start,
                    end=end + timedelta(hours=1),
                    label_type=(row.get("温度阶段") or "").strip(),
                    label_subtype=(row.get("阶段部位") or row.get("特征类型") or "").strip() or None,
                    row=row,
                    source_file=source_file,
                    site_points=site_points,
                )
            )
    return [row for row in rows if row[17] in SEVERITY]


def build_rows(site_points: dict[str, str]) -> list[tuple[Any, ...]]:
    rows: list[tuple[Any, ...]] = []
    for item in LABEL_FILES:
        path = ROOT / item["relativePath"]
        rows.extend(build_rows_for_file(item["key"], path, site_points))
    return rows


def load_site_points() -> dict[str, str]:
    with get_connection() as connection:
        with connection.cursor() as cursor:
            cursor.execute("SELECT ItemName, ItemDescription FROM site_point")
            return {str(row["ItemName"]): str(row["ItemDescription"] or "") for row in cursor.fetchall()}


def file_status() -> list[dict[str, Any]]:
    files: list[dict[str, Any]] = []
    for item in LABEL_FILES:
        path = ROOT / item["relativePath"]
        exists = path.exists()
        files.append(
            {
                "key": item["key"],
                "label": item["label"],
                "relativePath": item["relativePath"],
                "exists": exists,
                "size": path.stat().st_size if exists else None,
                "modifiedAt": (
                    datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d %H:%M:%S")
                    if exists
                    else None
                ),
            }
        )
    return files


def run_import(progress: ProgressCallback | None = None) -> dict[str, Any]:
    started = time.monotonic()
    ensure_label_table()
    if progress:
        progress(5, "准备")
    site_points = load_site_points()
    rows: list[tuple[Any, ...]] = []
    by_source: dict[str, int] = {}
    skipped: list[str] = []
    for index, item in enumerate(LABEL_FILES):
        path = ROOT / item["relativePath"]
        if progress:
            progress(15 + index * 18, f"读取{item['label']}")
        if not path.exists():
            skipped.append(item["relativePath"])
            continue
        file_rows = build_rows_for_file(item["key"], path, site_points)
        by_source[item["key"]] = len(file_rows)
        rows.extend(file_rows)
    if progress:
        progress(75, "写入数据库")
    if rows:
        with get_connection() as connection:
            with connection.cursor() as cursor:
                cursor.executemany(INSERT_SQL, rows)
    if progress:
        progress(90, "统计结果")
    with get_connection() as connection:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT COUNT(*) AS n FROM compressor_label WHERE origin = 'algorithm' AND deleted_at IS NULL"
            )
            active = int(cursor.fetchone()["n"])
    duration_ms = int((time.monotonic() - started) * 1000)
    if progress:
        progress(100, "完成")
    return {
        "total": len(rows),
        "bySource": by_source,
        "skippedFiles": skipped,
        "activeLabels": active,
        "durationMs": duration_ms,
    }


_job_lock = threading.Lock()
_job_state: dict[str, Any] = {
    "running": False,
    "progress": 0,
    "stage": "空闲",
    "message": "",
    "startedAt": None,
    "finishedAt": None,
    "result": None,
    "error": None,
}


def job_status() -> dict[str, Any]:
    with _job_lock:
        return copy.deepcopy(_job_state)


def _update_job(**kwargs: Any) -> None:
    with _job_lock:
        _job_state.update(kwargs)


def _progress_callback(progress: int, stage: str) -> None:
    _update_job(progress=progress, stage=stage)


def _run_job() -> None:
    try:
        result = run_import(progress=_progress_callback)
        _update_job(
            running=False,
            progress=100,
            stage="完成",
            result=result,
            finishedAt=_now(),
            error=None,
        )
    except Exception as error:  # noqa: BLE001
        _update_job(
            running=False,
            stage="失败",
            error=str(error),
            finishedAt=_now(),
        )


def start_import() -> None:
    with _job_lock:
        if _job_state["running"]:
            raise ValueError("导入正在进行中")
        _job_state.update(
            {
                "running": True,
                "progress": 0,
                "stage": "准备",
                "message": "",
                "startedAt": _now(),
                "finishedAt": None,
                "result": None,
                "error": None,
            }
        )
    threading.Thread(target=_run_job, daemon=True).start()