# Copyright (C) 2026 Hector van der Aa # Copyright (C) 2026 Association Exergie # SPDX-License-Identifier: GPL-3.0-or-later from __future__ import annotations from copy import copy from datetime import datetime, timezone import logging import os from pathlib import Path from tempfile import NamedTemporaryFile from typing import Literal, Self from uuid import UUID from zipfile import ZIP_DEFLATED, ZipFile from pydantic import BaseModel, Field from dynalab_core.buffer import ValueBuffer from dynalab_core.errors import DLPakError from dynalab_core.protocols.packets.handshake import ( JsonSignalDescriptor, SignalDescriptor, ) log = logging.getLogger(__name__) class RecordManifest(BaseModel): format_version: Literal[1] = 1 timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) signals: list[JsonSignalDescriptor] = Field(default_factory=list) comments: str = "" class DLPak: def __init__(self) -> None: self._dir: Path | None = None self._filename: str | None = None self._savetype: Literal["zip", "folder"] = "zip" self._data: ValueBuffer | None = None self._manifest: RecordManifest | None = None def set_data(self, data: ValueBuffer) -> None: self._data = data def set_manifest( self, timestamp: datetime, all_signals: list[SignalDescriptor] ) -> None: if self._data is None: raise DLPakError( "Cannot create a DLPak manifest before setting its data", code="no_data", ) data_signal_ids = self._data.get_signal_ids() manifest_signals: list[JsonSignalDescriptor] = [] for signal in all_signals: if signal.id in data_signal_ids: manifest_signals.append(signal.to_json()) self._manifest = RecordManifest( timestamp=timestamp, signals=manifest_signals, comments="" ) def merge(self, dlpak: DLPak) -> None: for sample in dlpak._data._samples: self._data.append(sample[0], sample[1], sample[2]) self._data.sort() for signal in dlpak._manifest.signals: self._manifest.signals.append(signal) log.debug( "Merged %d sample(s) and %d signal(s) into DLPak", len(dlpak._data), len(dlpak._manifest.signals), extra={ "event": "dlpak.merged", "added_sample_count": len(dlpak._data), "added_signal_count": len(dlpak._manifest.signals), }, ) def remove_derived(self) -> None: if self._manifest is not None: original_sample_count = len(self._data) remove_uuids = [ signal.id for signal in self._manifest.signals if signal.origin == "derived" ] for remove_uuid in remove_uuids: self._data.clear_uuid(remove_uuid) self._manifest.signals = [ signal for signal in self._manifest.signals if signal.origin != "derived" ] log.debug( "Removed %d derived signal(s) and %d sample(s) from DLPak", len(remove_uuids), original_sample_count - len(self._data), extra={ "event": "dlpak.derived_removed", "removed_signal_count": len(remove_uuids), "removed_sample_count": original_sample_count - len(self._data), }, ) def get_all_signal_descriptors(self) -> list[SignalDescriptor]: return copy([signal.to_pure() for signal in self._manifest.signals]) def get_data_range( self, signal_ids: list[UUID] | None = None, start_ns: int | None = None, end_ns: int | None = None, ) -> dict[UUID, list[tuple[int, float]]]: if self._data is None: raise DLPakError("No data is available", code="no_data") return self._data.get_data_range(signal_ids, start_ns, end_ns) def get_time_range(self) -> tuple[int, int] | None: if self._data is None: raise DLPakError("No data is available", code="no_data") return self._data.get_time_range() @staticmethod def _validate_contents(data: ValueBuffer, manifest: RecordManifest) -> None: signal_ids = [signal.id for signal in manifest.signals] if len(signal_ids) != len(set(signal_ids)): raise DLPakError( "The manifest contains duplicate signal IDs", code="invalid_manifest" ) missing = set(data.get_signal_ids()) - set(signal_ids) if missing: raise DLPakError( "The manifest is missing descriptors for signals: " + ", ".join(sorted(str(uuid) for uuid in missing)), code="invalid_manifest", ) def read(self, file: str | Path) -> Self: input_path = Path(file) try: with ZipFile(input_path, mode="r") as archive: manifest = RecordManifest.model_validate_json( archive.read("manifest.json") ) data = ValueBuffer() data.import_csv(archive.read("data.csv").decode("utf-8")) self._validate_contents(data, manifest) except Exception as error: log.exception( "Failed to read DLPak archive from %s", input_path.name, extra={ "event": "dlpak.read_failed", "input_filename": input_path.name, "exception_type": type(error).__name__, }, ) raise DLPakError( f"Could not read DLPak archive {input_path}: {error}", code="read_failed", ) from error self._dir = input_path.parent self._filename = input_path.name self._data = data self._manifest = manifest log.info( "Read DLPak archive from %s", input_path.name, extra={ "event": "dlpak.read", "input_filename": input_path.name, "sample_count": len(data), "signal_count": len(manifest.signals), }, ) return self def write(self, dir: str | Path, filename: str) -> None: output_dir = Path(dir) if self._data is None: raise DLPakError( "Cannot write a DLPak archive before setting its data", code="no_data", ) if self._manifest is None: raise DLPakError( "Cannot write a DLPak archive before creating its manifest", code="no_manifest", ) if output_dir.is_file(): output_dir = output_dir.parent file_path = Path(filename) if file_path.suffix.lower() != ".dlpak": file_path = file_path.with_suffix(".dlpak") output_path = output_dir / file_path temporary_path: Path | None = None try: self._validate_contents(self._data, self._manifest) with NamedTemporaryFile( dir=output_path.parent, prefix=f".{output_path.name}.", suffix=".tmp", delete=False, ) as temporary_file: temporary_path = Path(temporary_file.name) with ZipFile( temporary_file, mode="w", compression=ZIP_DEFLATED ) as archive: archive.writestr( "manifest.json", self._manifest.model_dump_json(indent=2) ) archive.writestr("data.csv", self._data.export_csv()) temporary_file.flush() os.fsync(temporary_file.fileno()) temporary_path.replace(output_path) except Exception as error: log.exception( "Failed to write DLPak archive to %s", output_path.name, extra={ "event": "dlpak.write_failed", "output_filename": output_path.name, "exception_type": type(error).__name__, }, ) raise DLPakError( f"Could not write DLPak archive {output_path}: {error}", code="write_failed", ) from error finally: if temporary_path is not None: temporary_path.unlink(missing_ok=True) log.info( "Wrote DLPak archive to %s", output_path.name, extra={ "event": "dlpak.written", "output_filename": output_path.name, "sample_count": len(self._data), "signal_count": len(self._manifest.signals), }, )