"""Unload, reload and removal while work is still in flight, in the real HA lifecycle."""
import asyncio
import threading
from datetime import timedelta
from unittest.mock import patch

from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.storage import Store
from homeassistant.util import dt as dt_util
from pytest_homeassistant_custom_component.common import async_fire_time_changed

from custom_components.heating_mpc import coordinator as coordinator_module
from custom_components.heating_mpc.const import DOMAIN

REAL_OPTIMIZE = coordinator_module.optimize


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]


def trained(coordinator):
    model = coordinator.learners["sensor.room"]
    model.count, model.on_count, model.off_count = 250, 125, 125
    coordinator._update_counts()


class BlockedOptimizer:
    """Holds the optimizer in its worker thread until released."""

    def __init__(self):
        self.started, self.release = threading.Event(), threading.Event()

    def __call__(self, *args, **kwargs):
        self.started.set()
        self.release.wait(5)
        return REAL_OPTIMIZE(*args, **kwargs)

    async def wait_started(self):
        # Polled on the event loop: the test executor may run one job at a time.
        for _ in range(500):
            if self.started.is_set():
                return True
            await asyncio.sleep(0.01)
        return False


def status_id(hass, entry):
    return er.async_get(hass).async_get_entity_id("sensor", DOMAIN, f"{entry.entry_id}_status")


async def test_unload_during_optimization_publishes_nothing(hass, entry, inputs):
    coordinator = await setup(hass, entry)
    trained(coordinator)
    shown = coordinator.data
    blocked = BlockedOptimizer()
    with patch.object(coordinator_module, "optimize", blocked):
        refresh = hass.async_create_task(coordinator.async_refresh())
        assert await blocked.wait_started()
        assert await hass.config_entries.async_unload(entry.entry_id)
        state_after_unload = hass.states.get(status_id(hass, entry))
        blocked.release.set()
        await refresh
        await hass.async_block_till_done()
    # The plan finished after unloading began: kept out of the data and the states.
    assert coordinator.data is shown
    assert coordinator.data["status"] == "learning"
    assert hass.states.get(status_id(hass, entry)) == state_after_unload


async def test_reload_during_optimization_shows_only_the_new_coordinator(hass, entry, inputs):
    old = await setup(hass, entry)
    trained(old)  # only in memory: the reloaded coordinator is not ready
    blocked = BlockedOptimizer()
    with patch.object(coordinator_module, "optimize", blocked):
        refresh = hass.async_create_task(old.async_refresh())
        assert await blocked.wait_started()
        assert await hass.config_entries.async_reload(entry.entry_id)
        await hass.async_block_till_done()
        new = hass.data[DOMAIN][entry.entry_id]
        blocked.release.set()
        await refresh
        await hass.async_block_till_done()
    assert new is not old and old.closing
    assert old.data["status"] != "ready"
    assert hass.states.get(status_id(hass, entry)).state == new.data["status"] == "learning"


async def test_removal_during_a_save_leaves_no_model_behind(hass, entry, inputs, freezer):
    coordinator = await setup(hass, entry)
    started, release = asyncio.Event(), asyncio.Event()
    write = coordinator._store.async_save

    async def slow_save(data):
        started.set()
        await release.wait()
        await write(data)
    coordinator._store.async_save = slow_save
    future = dt_util.utcnow() + timedelta(minutes=5, seconds=1)
    freezer.move_to(future)
    async_fire_time_changed(hass, future)  # a learning poll, which saves
    await asyncio.wait_for(started.wait(), 5)
    removal = hass.async_create_task(hass.config_entries.async_remove(entry.entry_id))
    for _ in range(5):
        await asyncio.sleep(0)
    assert not removal.done()  # deletion waits for the save in progress
    release.set()
    assert await removal
    await hass.async_block_till_done()
    assert await Store(hass, 1, f"{DOMAIN}.{entry.entry_id}").async_load() is None


async def test_repeated_setup_and_unload_leave_nothing_running(hass, entry, inputs, freezer):
    entry.add_to_hass(hass)
    coordinators = []
    for _ in range(3):
        assert await hass.config_entries.async_setup(entry.entry_id)
        await hass.async_block_till_done()
        coordinators.append(hass.data[DOMAIN][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 all(c.closing and c._unsub_refresh is None and not c._listeners for c in coordinators)
    counts = [c.learners["sensor.room"].count for c in coordinators]
    hass.states.async_set("switch.boiler", "on")
    future = dt_util.utcnow() + timedelta(minutes=10)
    freezer.move_to(future)
    async_fire_time_changed(hass, future)
    await hass.async_block_till_done()
    assert not any(c._relay_events for c in coordinators)
    assert [c.learners["sensor.room"].count for c in coordinators] == counts