from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass

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


@dataclass(frozen=True)
class RetrainingResult:
    artifact: TrainedArtifact
    replaced: bool


class RetrainingService:
    """Runs one retraining cycle without owning scheduling or background threads."""

    def __init__(self, registry: ModelRegistry) -> None:
        self._registry = registry

    def retrain(
        self,
        artifact_id: str,
        vectors: Iterable[FeatureVector],
    ) -> RetrainingResult:
        store = FeatureStore()
        store.add_batch(vectors)
        pipeline = TrainingPipeline(store)
        artifact = pipeline.run(artifact_id)
        _, replaced = self._registry.register_with_status(artifact)
        return RetrainingResult(artifact=artifact, replaced=replaced)


def retrain_model(
    registry: ModelRegistry,
    artifact_id: str,
    vectors: Iterable[FeatureVector],
) -> RetrainingResult:
    """Scheduler-compatible entry point for exactly one retraining run."""

    return RetrainingService(registry).retrain(artifact_id, vectors)