Compressor/backend/app/services/pretrain_service.py

import csv
import threading
from pathlib import Path

_PRETRAIN_CSV_NAME = "预训练之文件头.csv"
_CSV_COLUMNS = ["id", "point_name", "measurement_type", "rpm", "sample_time", "file_name"]


class PretrainService:
    """Append selected wave_file rows to the pre-training CSV (deduplicated by id)."""

    def __init__(self) -> None:
        self._csv_path = Path(__file__).resolve().parents[2] / _PRETRAIN_CSV_NAME
        self._lock = threading.Lock()

    @property
    def csv_path(self) -> Path:
        return self._csv_path

    def _read_records(self) -> list[dict[str, str]]:
        if not self._csv_path.exists():
            return []
        with self._csv_path.open("r", encoding="utf-8-sig", newline="") as handle:
            reader = csv.DictReader(handle)
            return [dict(row) for row in reader]

    def _write_records(self, records: list[dict[str, str]]) -> None:
        with self._csv_path.open("w", encoding="utf-8-sig", newline="") as handle:
            writer = csv.DictWriter(handle, fieldnames=_CSV_COLUMNS)
            writer.writeheader()
            writer.writerows(records)

    def recorded_ids(self, ids: list[int]) -> list[int]:
        wanted = {int(value) for value in ids if value}
        if not wanted:
            return []
        with self._lock:
            records = self._read_records()
        existing: set[int] = set()
        for row in records:
            raw = str(row.get("id") or "").strip()
            if raw.isdigit():
                existing.add(int(raw))
        return sorted(existing & wanted)

    def recorded_counts(self) -> dict[str, int]:
        """每个 point_name 在预训练 CSV 中的记录数量。"""
        with self._lock:
            records = self._read_records()
        counts: dict[str, int] = {}
        for row in records:
            name = str(row.get("point_name") or "").strip()
            if name:
                counts[name] = counts.get(name, 0) + 1
        return counts

    def record(self, record: dict[str, object]) -> tuple[bool, bool]:
        row = {column: str(record.get(column) or "") for column in _CSV_COLUMNS}
        try:
            file_id = int(row["id"])
        except (TypeError, ValueError) as error:
            raise ValueError("id 必须为正整数") from error
        if file_id <= 0:
            raise ValueError("id 必须为正整数")
        with self._lock:
            records = self._read_records()
            existing: set[int] = set()
            for item in records:
                raw = str(item.get("id") or "").strip()
                if raw.isdigit():
                    existing.add(int(raw))
            if file_id in existing:
                return False, True
            records.append(row)
            self._write_records(records)
        return True, False

    def remove(self, file_id: int) -> bool:
        if file_id <= 0:
            raise ValueError("id 必须为正整数")
        with self._lock:
            records = self._read_records()
            remaining = [
                row
                for row in records
                if not str(row.get("id") or "").strip().isdigit()
                or int(str(row.get("id") or "").strip()) != file_id
            ]
            if len(remaining) == len(records):
                return False
            self._write_records(remaining)
        return True


pretrain_service = PretrainService()