"""Run the observer through real HA lifecycle and storage APIs."""
from datetime import timedelta
from unittest.mock import patch

import pytest
from homeassistant.util import dt as dt_util
from pytest_homeassistant_custom_component.common import async_fire_time_changed

from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.storage import Store

from custom_components.heating_mpc.const import DOMAIN


async def setup(hass, entry):
    entry.add_to_hass(hass)
    assert await hass.config_entries.async_setup(entry.entry_id)
    await hass.async_block_till_done()
    return hass.data[DOMAIN][entry.entry_id]


async def test_setup_reload_unload(hass, entry, inputs):
    coordinator = await setup(hass, entry)
    assert coordinator.config_entry is entry
    registry = er.async_get(hass)
    entities = er.async_entries_for_config_entry(registry, entry.entry_id)
    assert len(entities) == 11  # 7 core + 4 forecast-error sensors
    ids = {e.unique_id: e.entity_id for e in entities}
    assert coordinator.data["status"] == "learning"
    status = hass.states.get(ids[f"{entry.entry_id}_status"])
    assert status.attributes["input_quality"]["rooms"]["sensor.room"]["max_age_seconds_24h"] == 0
    # Volatile diagnostics stay visible but are not written to the Recorder.
    assert {"input_quality", "weather_comparison"} <= status.state_info["unrecorded_attributes"]
    assert "models" not in status.state_info["unrecorded_attributes"]
    proposal_id = registry.async_get_entity_id("binary_sensor", DOMAIN, f"{entry.entry_id}_proposal")
    assert hass.states.get(proposal_id).state == "unavailable"
    assert await hass.config_entries.async_reload(entry.entry_id)
    await hass.async_block_till_done()
    assert hass.data[DOMAIN][entry.entry_id] is not coordinator
    assert ids == {e.unique_id: e.entity_id for e in er.async_entries_for_config_entry(registry, entry.entry_id)}
    assert await hass.config_entries.async_unload(entry.entry_id)
    await hass.async_block_till_done()
    assert entry.entry_id not in hass.data[DOMAIN]
    assert not coordinator._listeners
    assert coordinator._unsub_refresh is None


async def test_options_reload_removes_only_old_room(hass, entry, inputs):
    await setup(hass, entry)
    registry = er.async_get(hass)
    old = registry.async_get_entity_id("sensor", DOMAIN, f"{entry.entry_id}_sensor.room_predicted_60min")
    hass.states.async_set("sensor.new", "21", {"unit_of_measurement": "°C"})
    hass.config_entries.async_update_entry(entry, options={"room_entities": ["sensor.new"]})
    await hass.async_block_till_done()
    assert hass.data[DOMAIN][entry.entry_id].rooms == ["sensor.new"]
    assert registry.async_get(old) is None
    assert registry.async_get_entity_id("sensor", DOMAIN, f"{entry.entry_id}_sensor.new_predicted_60min")
    assert len(er.async_entries_for_config_entry(registry, entry.entry_id)) == 11


async def test_storage_restore_and_remove(hass, entry, inputs):
    coordinator = await setup(hass, entry)
    model = coordinator.learners["sensor.room"]
    model.count, model.on_count, model.off_count = 250, 125, 125
    await coordinator._store.async_save({"schema_version": 2, "heater_entity": "switch.boiler",
                                       "learners": {"sensor.room": model.as_dict()}})
    assert await hass.config_entries.async_reload(entry.entry_id)
    await hass.async_block_till_done()
    restored = hass.data[DOMAIN][entry.entry_id]
    assert restored.learners["sensor.room"].count == 250
    assert restored.data["status"] == "ready"
    forecast_id = er.async_get(hass).async_get_entity_id(
        "sensor", DOMAIN, f"{entry.entry_id}_sensor.room_predicted_60min")
    forecast = hass.states.get(forecast_id)
    assert len(forecast.attributes["forecast_15min"]) == len(forecast.attributes["plan_15min"]) == 4
    assert forecast.attributes["forecast_15min"][3] == float(forecast.state)
    assert {"forecast_15min", "plan_15min"} <= forecast.state_info["unrecorded_attributes"]
    assert await hass.config_entries.async_remove(entry.entry_id)
    await hass.async_block_till_done()
    assert await Store(hass, 1, f"{DOMAIN}.{entry.entry_id}").async_load() is None
    assert not er.async_entries_for_config_entry(er.async_get(hass), entry.entry_id)


async def test_missing_sensor_manual_off_and_target(hass, entry, inputs):
    coordinator = await setup(hass, entry)
    hass.states.async_set("sensor.room", "unavailable")
    await coordinator.async_refresh()
    assert coordinator.data["status"] == "waiting_for_sensors"
    hass.states.async_set("sensor.room", "20", {"unit_of_measurement": "°C"})
    hass.states.async_set("climate.reference", "off", {"temperature": 19, "preset_mode": "away"})
    await coordinator.async_refresh()
    assert coordinator.data["status"] == "paused"
    assert coordinator.data["proposal"] is None
    assert hass.states.get("climate.reference").attributes["preset_mode"] == "away"
    hass.states.async_set("climate.reference", "heat", {"temperature": 21, "preset_mode": "sleep"})
    await coordinator.async_refresh()
    assert coordinator.data["target"] == 21
    assert hass.states.get("switch.boiler").state == "off"


async def test_optimizer_failure_is_unavailable(hass, entry, inputs):
    coordinator = await setup(hass, entry)
    model = coordinator.learners["sensor.room"]
    model.count, model.on_count, model.off_count = 250, 125, 125
    coordinator._update_counts()
    with patch("custom_components.heating_mpc.coordinator.optimize", side_effect=RuntimeError("test")):
        await coordinator.async_refresh()
    assert coordinator.data["status"] == "prediction_error"
    assert coordinator.data["proposal"] is None


async def test_scheduled_sample_saves_and_unload_stops_polling(hass, entry, inputs, freezer, hass_storage):
    coordinator = await setup(hass, entry)
    future = dt_util.utcnow() + timedelta(minutes=5, seconds=1)
    freezer.move_to(future)
    async_fire_time_changed(hass, future)
    await hass.async_block_till_done()
    assert coordinator.learners["sensor.room"].count == 1
    payload = hass_storage[f"{DOMAIN}.{entry.entry_id}"]["data"]
    assert payload["learners"]["sensor.room"]["off_count"] == 1
    assert await hass.config_entries.async_unload(entry.entry_id)
    future += timedelta(minutes=5)
    freezer.move_to(future)
    async_fire_time_changed(hass, future)
    await hass.async_block_till_done()
    assert coordinator.learners["sensor.room"].count == 1


@pytest.mark.parametrize("payload", ["broken", {"learners": []}, {"learners": {"sensor.room": {"theta": []}}}])
async def test_corrupt_saved_payload_starts_learning(hass, entry, inputs, hass_storage, payload):
    hass_storage[f"{DOMAIN}.{entry.entry_id}"] = {"version": 1, "data": payload}
    coordinator = await setup(hass, entry)
    assert coordinator.data["status"] == "learning"
    assert coordinator.data["samples"] == 0


async def test_startup_preserves_manual_off_with_trained_model(hass, entry, inputs, hass_storage):
    from custom_components.heating_mpc.physics import Learner
    model = Learner(count=250, on_count=125, off_count=125)
    hass_storage[f"{DOMAIN}.{entry.entry_id}"] = {"version": 1, "data": {
        "schema_version": 2, "heater_entity": "switch.boiler",
        "learners": {"sensor.room": model.as_dict()}}}
    hass.states.async_set("climate.reference", "off", {"temperature": 19, "preset_mode": "away"})
    coordinator = await setup(hass, entry)
    assert coordinator.data["status"] == "paused"
    assert coordinator.data["proposal"] is None
    assert hass.states.get("climate.reference").state == "off"
    assert hass.states.get("climate.reference").attributes["temperature"] == 19


async def test_startup_missing_inputs_then_recovery(hass, entry):
    coordinator = await setup(hass, entry)
    assert coordinator.data["status"] == "waiting_for_sensors"
    hass.states.async_set("sensor.room", "20", {"unit_of_measurement": "°C"})
    hass.states.async_set("sensor.outside", "5", {"unit_of_measurement": "°C"})
    hass.states.async_set("switch.boiler", "off")
    hass.states.async_set("climate.reference", "heat", {"temperature": 22})
    await coordinator.async_refresh()
    assert coordinator.data["status"] == "learning"
    assert coordinator.data["samples"] == 0


async def test_freshly_written_relay_state_defers_learning(hass, entry, inputs, freezer):
    """After an HA restart the relay's last_changed is new, so radiator lag is unknown."""
    hass.states.async_set("switch.boiler", "on")
    coordinator = await setup(hass, entry)
    assert coordinator.data["heat_lag"] == 0.5
    assert coordinator.data["heat_lag_uncertainty"] == 0.5
    future = dt_util.utcnow() + timedelta(minutes=5, seconds=1)
    freezer.move_to(future)
    async_fire_time_changed(hass, future)
    await hass.async_block_till_done()
    assert coordinator.data["sample_result"] == "lag_uncertain"
    assert coordinator.learners["sensor.room"].count == 0


async def test_relay_listener_integrates_changes_between_polls(hass, entry, inputs, freezer):
    coordinator = await setup(hass, entry)
    start = dt_util.utcnow()
    for offset, state in ((60, "on"), (200, "off")):  # on/off round trip inside one poll
        freezer.move_to(start + timedelta(seconds=offset))
        hass.states.async_set("switch.boiler", state)
        await hass.async_block_till_done()
    for minutes, expected in ((5, "relay_changed"), (10, "learned")):
        future = start + timedelta(minutes=minutes, seconds=1)
        freezer.move_to(future)
        async_fire_time_changed(hass, future)
        await hass.async_block_till_done()
        assert coordinator.data["sample_result"] == expected
        # Without the listener this round trip would restart the estimate (about 0.46).
        assert coordinator.data["heat_lag_uncertainty"] < 0.05
    assert await hass.config_entries.async_unload(entry.entry_id)
    hass.states.async_set("switch.boiler", "on")
    await hass.async_block_till_done()
    assert not coordinator._relay_events


async def test_removed_room_model_survives_options_round_trip(hass, entry, inputs, hass_storage, freezer):
    from custom_components.heating_mpc.physics import Learner
    key = f"{DOMAIN}.{entry.entry_id}"
    hass_storage[key] = {"version": 1, "data": {
        "schema_version": 2, "heater_entity": "switch.boiler",
        "learners": {"sensor.room": Learner(count=250, on_count=125, off_count=125).as_dict()}}}
    hass.states.async_set("sensor.new", "21", {"unit_of_measurement": "°C"})
    await setup(hass, entry)
    hass.config_entries.async_update_entry(entry, options={"room_entities": ["sensor.new"]})
    await hass.async_block_till_done()
    future = dt_util.utcnow() + timedelta(minutes=5, seconds=1)
    freezer.move_to(future)
    async_fire_time_changed(hass, future)
    await hass.async_block_till_done()
    saved = hass_storage[key]["data"]
    assert set(saved["learners"]) == {"sensor.new"}
    assert saved["retired_learners"]["sensor.room"]["count"] == 250
    hass.config_entries.async_update_entry(
        entry, options={"room_entities": ["sensor.room", "sensor.new"]})
    await hass.async_block_till_done()
    assert hass.data[DOMAIN][entry.entry_id].learners["sensor.room"].count == 250


async def test_forecast_error_sensors_score_scheduled_polls(hass, entry, inputs, freezer):
    coordinator = await setup(hass, entry)
    registry = er.async_get(hass)
    ids = {horizon: registry.async_get_entity_id("sensor", DOMAIN, f"{entry.entry_id}_forecast_error_{horizon}min")
           for horizon in (15, 30, 60, 120)}
    assert all(hass.states.get(entity).state == "unknown" for entity in ids.values())
    future = dt_util.utcnow()
    for _ in range(3):  # 15 minutes of regular polls
        future += timedelta(minutes=5, seconds=1)
        freezer.move_to(future)
        async_fire_time_changed(hass, future)
        await hass.async_block_till_done()
    state = hass.states.get(ids[15])
    assert coordinator.data["forecast_accuracy"][15]["samples"] == 1
    assert float(state.state) == coordinator.data["forecast_accuracy"][15]["mae"] > 0
    assert state.attributes["samples"] == 1
    assert state.attributes["rooms"]["sensor.room"]["samples"] == 1
    assert state.attributes["unit_of_measurement"] == "°C"
    assert "rooms" in state.state_info["unrecorded_attributes"]
    assert hass.states.get(ids[30]).state == "unknown"


async def test_learning_progress_and_last_sample_sensors(hass, entry, inputs, freezer):
    await setup(hass, entry)
    registry = er.async_get(hass)
    progress_id = registry.async_get_entity_id("sensor", DOMAIN, f"{entry.entry_id}_learning_progress")
    last_id = registry.async_get_entity_id("sensor", DOMAIN, f"{entry.entry_id}_last_sample")
    assert hass.states.get(progress_id).state == "0.0"
    last = hass.states.get(last_id)
    assert last.state == "first_sample"
    assert "lag_uncertain" in last.attributes["options"]
    future = dt_util.utcnow() + timedelta(minutes=5, seconds=1)
    freezer.move_to(future)
    async_fire_time_changed(hass, future)
    await hass.async_block_till_done()
    assert hass.states.get(last_id).state == "learned"
    # One OFF sample: (1/200 + 0/30 + 1/30) / 3
    assert float(hass.states.get(progress_id).state) == 1.3
    progress = hass.states.get(progress_id).attributes
    assert (progress["samples"], progress["off_samples"], progress["required"]["samples"]) == (1, 1, 200)
    hass.states.async_set("switch.boiler", "unavailable")
    future += timedelta(minutes=5)
    freezer.move_to(future)
    async_fire_time_changed(hass, future)
    await hass.async_block_till_done()
    last = hass.states.get(last_id)
    assert last.state == "waiting_for_inputs"
    assert last.attributes["counts_since_start"] == {"first_sample": 1, "learned": 1, "waiting_for_inputs": 1}