from __future__ import annotations

import json
import logging
import os
from pathlib import Path
import re
from threading import RLock
from collections.abc import Iterable

from app.ml.training import TrainedArtifact

logger = logging.getLogger(__name__)

_ARTIFACT_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")


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

    def register(self, artifact: TrainedArtifact) -> TrainedArtifact:
        registered, _ = self.register_with_status(artifact)
        return registered

    def register_with_status(self, artifact: TrainedArtifact) -> tuple[TrainedArtifact, bool]:
        self._validate_artifact_id(artifact.artifact_id)
        with self._lock:
            replaced = artifact.artifact_id in self._artifacts
            self._persist(artifact)
            self._artifacts[artifact.artifact_id] = artifact
        return artifact, replaced

    def load_artifact(self, artifact_id: str) -> TrainedArtifact:
        self._validate_artifact_id(artifact_id)
        with self._lock:
            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]:
        with self._lock:
            return [self._artifacts[key] for key in sorted(self._artifacts)]

    def _load_existing(self) -> None:
        for source in sorted(self._root.glob("*.json")):
            try:
                raw = json.loads(source.read_text(encoding="utf-8"))
                artifact_id = raw["artifact_id"]
                supported_sensors = raw["supported_sensors"]
                if not isinstance(artifact_id, str) or not isinstance(supported_sensors, list):
                    raise ValueError("invalid artifact structure")
                self._validate_artifact_id(artifact_id)
                if source.name != f"{artifact_id}.json":
                    raise ValueError("artifact id does not match filename")
                if not all(isinstance(sensor, str) for sensor in supported_sensors):
                    raise ValueError("supported_sensors must contain strings")
            except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
                raise ValueError(f"Ungültiges Modell-Artefakt: {source.name}") from exc

            self._artifacts[artifact_id] = TrainedArtifact(
                artifact_id=artifact_id,
                supported_sensors=tuple(supported_sensors),
            )

    def _persist(self, artifact: TrainedArtifact) -> None:
        target = self._root / f"{artifact.artifact_id}.json"
        temporary = target.with_suffix(".json.tmp")
        payload = {
            "artifact_id": artifact.artifact_id,
            "supported_sensors": list(artifact.supported_sensors),
        }
        temporary.write_text(
            json.dumps(payload, ensure_ascii=True, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        os.replace(temporary, target)
        logger.info("Modell gespeichert: %s", target)

    @staticmethod
    def _validate_artifact_id(artifact_id: str) -> None:
        if not _ARTIFACT_ID_PATTERN.fullmatch(artifact_id) or ".." in artifact_id:
            raise ValueError(
                "artifact_id darf nur Buchstaben, Ziffern, Punkt, Unterstrich "
                "und Bindestrich enthalten."
            )