from __future__ import annotations

import logging
from typing import Sequence

from app.ml.feature_store import FeatureVector
from app.ml.registry.model_registry import ModelRegistry
from app.ml.training import TrainedArtifact, TrainingPipeline

logger = logging.getLogger(__name__)


class Predictor:
    def __init__(
        self,
        pipeline: TrainingPipeline | None = None,
        registry: ModelRegistry | None = None,
    ) -> None:
        if isinstance(pipeline, ModelRegistry) and registry is None:
            registry = pipeline
            pipeline = None
        if pipeline is None and registry is None:
            raise ValueError("Predictor erfordert TrainingPipeline oder ModelRegistry.")
        self._pipeline = pipeline
        self._registry = registry

    def predict(self, artifact_id: str, entity: FeatureVector) -> str:
        artifact = self._get_artifact(artifact_id)
        if entity.sensor_id not in artifact.supported_sensors:
            raise ValueError(
                f"Sensor '{entity.sensor_id}' wird vom Modell '{artifact_id}' nicht unterstützt."
            )
        return f"{artifact_id}:{entity.sensor_id}:{entity.values}"

    def predict_batch(self, artifact_id: str, entities: Sequence[FeatureVector]) -> list[str]:
        return [self.predict(artifact_id, entity) for entity in entities]

    @staticmethod
    def default_artifact(pipeline: TrainingPipeline) -> TrainedArtifact:
        artifacts = list(pipeline._artifacts)
        if not artifacts:
            raise ValueError("Kein trainiertes Modell gefunden.")
        return pipeline.export(artifacts[-1])

    def _get_artifact(self, artifact_id: str) -> TrainedArtifact:
        if self._registry is not None:
            return self._registry.load_artifact(artifact_id)
        if self._pipeline is not None:
            return self._pipeline.export(artifact_id)
        raise RuntimeError("Predictor nicht initialisiert.")