from __future__ import annotations

import pytest

from app.ml.feature_store import FeatureStore, FeatureVector
from app.ml.predictor import Predictor
from app.ml.training import TrainingPipeline


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


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


def test_predict_returns_statistical_forecast() -> None:
    p = predictor()
    result = p.predict("artifact_v1", _vector("sensor.kitchen", 21.0))
    assert result.artifact_id == "artifact_v1"
    assert result.sensor_id == "sensor.kitchen"
    assert result.predictions == {"temperature": 22.0}
    assert 0.0 < result.confidence <= 1.0
    assert result.model_type == "statistical_baseline"


def test_predict_rejects_unknown_sensor() -> None:
    p = predictor()
    with pytest.raises(ValueError):
        p.predict("artifact_v1", _vector("sensor.unknown", 10.0))


def test_predict_batch_matches_single_calls() -> None:
    p = predictor()
    entities = [_vector("sensor.kitchen", 21.0), _vector("sensor.bedroom", 19.0)]
    assert p.predict_batch("artifact_v1", entities) == [
        p.predict("artifact_v1", item) for item in entities
    ]


def test_default_artifact_returns_last_registered() -> None:
    store = FeatureStore()
    store.add_batch([_vector("sensor.kitchen", 19.0), _vector("sensor.bedroom", 18.5)])
    pipeline = TrainingPipeline(store)
    pipeline.run("first")
    pipeline.run("second")
    assert Predictor.default_artifact(pipeline).artifact_id == "second"