from __future__ import annotations

import logging
from pathlib import Path
from typing import Iterable

from app.ml.training import TrainedArtifact

logger = logging.getLogger(__name__)


class ModelRegistry:
    def __init__(self, root: str | Path) -> None:
        self._root = Path(root)
        self._root.mkdir(parents=True, exist_ok=True)
        self._artifacts: dict[str, TrainedArtifact] = {}

    def register(self, artifact: TrainedArtifact) -> TrainedArtifact:
        self._artifacts[artifact.artifact_id] = artifact
        self._persist(artifact)
        return artifact

    def load_artifact(self, artifact_id: str) -> TrainedArtifact:
        if artifact_id not in self._artifacts:
            raise KeyError(f"Artifact '{artifact_id}' nicht registriert.")
        return self._artifacts[artifact_id]

    def list_models(self) -> Iterable[TrainedArtifact]:
        return list(self._artifacts.values())

    def _persist(self, artifact: TrainedArtifact) -> None:
        target = self._root / f"{artifact.artifact_id}.json"
        target.write_text(
            f"{artifact.artifact_id}\t{','.join(artifact.supported_sensors)}\n",
            encoding="utf-8",
        )
        logger.info("Modell gespeichert: %s", target)