from __future__ import annotations

import pytest

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


def _vector(sensor_id: str, temperature: float, label: str | None = None) -> FeatureVector:
    return FeatureVector(sensor_id=sensor_id, values={"temperature": temperature}, label=label)


def store_with_data() -> TrainingPipeline:
    store = FeatureStore()
    store.add_batch(
        [
            _vector("sensor.kitchen", 19.0),
            _vector("sensor.kitchen", 20.0),
            _vector("sensor.bedroom", 18.5),
        ]
    )
    return TrainingPipeline(store)


def test_run_returns_trained_artifact() -> None:
    pipeline = store_with_data()
    artifact = pipeline.run("artifact_v1")
    assert artifact.artifact_id == "artifact_v1"
    assert artifact.supported_sensors == ("sensor.bedroom", "sensor.kitchen")


def test_run_without_data_raises_value_error() -> None:
    pipeline = TrainingPipeline(FeatureStore())
    with pytest.raises(ValueError):
        pipeline.run("artifact_v1")


def test_export_returns_registered_artifact() -> None:
    pipeline = store_with_data()
    pipeline.run("artifact_v1")
    exported = pipeline.export("artifact_v1")
    assert exported == pipeline.export("artifact_v1")


def test_export_missing_artifact_raises_key_error() -> None:
    pipeline = store_with_data()
    with pytest.raises(KeyError):
        pipeline.export("artifact_v1")