"""Read existing HA entities, learn thermal response, produce shadow MPC decisions."""
from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from functools import partial
from math import exp, isfinite
from typing import Any

from homeassistant.core import HomeAssistant
from homeassistant.config_entries import ConfigEntry
from homeassistant.helpers.storage import Store
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator

from .const import (
    CONF_ROOM_MAX_AGE, CONF_WEATHER_MAX_AGE,
    DOMAIN, CONF_CLIMATE, CONF_HEATER, CONF_ROOMS, CONF_WEIGHTS, WEATHER_SOURCES,
    CONF_HORIZON, CONF_MIN_ON, CONF_MIN_OFF, CONF_SWITCH_COST, CONF_HEAT_COST,
    CONF_BAND, DEFAULTS,
)
from .data_quality import freshness
from .physics import Learner, optimize
from .weather_sources import choose_weather_sources, sources_for, normalize

_LOGGER = logging.getLogger(__name__)


class HeatingMPCCoordinator(DataUpdateCoordinator[dict[str, Any]]):
    """The existing generic thermostat stays the only actuator owner."""

    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
        super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=timedelta(minutes=5))
        self.entry = entry
        self.config = {**DEFAULTS, **entry.data, **entry.options}
        self.rooms: list[str] = self.config[CONF_ROOMS]
        self.learners: dict[str, Learner] = {room: Learner() for room in self.rooms}
        self._store = Store(hass, 1, f"{DOMAIN}.{entry.entry_id}")
        self._previous: dict[str, Any] | None = None
        self._heat_lag: float | None = None
        self._on_samples = 0
        self._off_samples = 0

    async def async_restore_model(self) -> None:
        stored = await self._store.async_load() or {}
        if not isinstance(stored, dict) or stored.get("schema_version", 1) not in (1, 2):
            _LOGGER.warning("Ignoring unsupported Heating MPC model data")
            return
        # Replacing the actual heating system invalidates heating coefficients.
        # Switching among weather providers does not discard months of learning.
        if stored.get("heater_entity", self.config[CONF_HEATER]) != self.config[CONF_HEATER]:
            _LOGGER.warning("Heating MPC heater changed; starting a fresh thermal model")
            return
        # Models carry room identity, so changing the room list is safe.
        saved = stored.get("learners", {})
        if not isinstance(saved, dict):
            return
        for room, data in saved.items():
            if room in self.learners:
                self.learners[room] = Learner.from_dict(data)
        self._update_counts()

    def _update_counts(self) -> None:
        self._on_samples = min(m.on_count for m in self.learners.values())
        self._off_samples = min(m.off_count for m in self.learners.values())

    def _sensor_quality(self, entity_id: str, now: datetime, limit: float) -> dict:
        state = self.hass.states.get(entity_id)
        if state is None:
            return {"status": "missing", "age_seconds": None}
        if state.state in ("unknown", "unavailable"):
            return {"status": state.state, "age_seconds": None}
        # last_updated is a conservative compatibility fallback for older HA.
        return freshness(getattr(state, "last_reported", getattr(state, "last_updated", None)),
                         now, limit * 60)

    async def _async_update_data(self) -> dict[str, Any]:
        now = datetime.now(timezone.utc)
        heater = self.hass.states.get(self.config[CONF_HEATER])
        climate = self.hass.states.get(self.config[CONF_CLIMATE])
        # Choose each weather variable independently from ordered sensor sources.
        # Sensor priorities are editable in Options without touching integration code.
        weather_entity_ids = {
            entity_id for metric in WEATHER_SOURCES
            for entity_id in sources_for(self.config, metric)
        }
        watched = set(self.rooms) | weather_entity_ids | {
            self.config[CONF_CLIMATE], self.config[CONF_HEATER]}
        snapshot = {entity: self.hass.states.get(entity) for entity in watched}
        room_quality = {}
        temps = {}
        for room in self.rooms:
            quality = self._sensor_quality(room, now, self.config[CONF_ROOM_MAX_AGE])
            value = None
            if quality["status"] == "valid":
                state = snapshot[room]
                value = normalize("temperature", state.state, state.attributes.get("unit_of_measurement"))
                if value is None or not 0 <= value <= 45:
                    quality["status"] = "invalid_value_or_unit"
                    value = None
            room_quality[room] = quality
            temps[room] = value
        weather_states = {}
        weather_quality = {}
        for entity_id in weather_entity_ids:
            quality = self._sensor_quality(entity_id, now, self.config[CONF_WEATHER_MAX_AGE])
            weather_quality[entity_id] = quality
            if quality["status"] == "valid":
                state = snapshot[entity_id]
                weather_states[entity_id] = (state.state, state.attributes.get("unit_of_measurement"))
        weather, active_sources, weather_comparisons = choose_weather_sources(
            self.config, weather_states
        )
        metric_quality = {}
        for metric in WEATHER_SOURCES:
            metric_quality[metric] = {}
            for entity_id in sources_for(self.config, metric):
                quality = dict(weather_quality[entity_id])
                if quality["status"] == "valid" and entity_id not in weather_comparisons[metric]:
                    quality["status"] = "invalid_value_or_unit"
                metric_quality[metric][entity_id] = quality
        outside = weather["temperature"]
        status: dict[str, Any] = {
            "input_quality": {"rooms": room_quality, "weather": metric_quality},
            "status": "waiting_for_sensors", "proposal": None, "prediction_60m": {},
            "temperatures": {r: t for r, t in temps.items() if t is not None},
            "on_samples": self._on_samples, "off_samples": self._off_samples,
            "samples": min(learner.count for learner in self.learners.values()),
            "mode": climate.state if climate else "unavailable",
            "weather": {k: v for k, v in weather.items() if v is not None},
            "weather_active_sources": {k: v for k, v in active_sources.items() if v},
            "weather_comparison": {k: v for k, v in weather_comparisons.items() if v},
            "weather_model_inputs": ["temperature"],
        }
        if heater is None or heater.state not in ("on", "off") or climate is None:
            self._previous = None
            self._heat_lag = None
            return status
        if outside is None or any(temp is None for temp in temps.values()):
            self._previous = None
            self._heat_lag = None
            return status

        # No casting surprises: all rooms are valid here.
        vals: dict[str, float] = {room: float(temps[room]) for room in self.rooms}
        heater_on = heater.state == "on"
        if self._heat_lag is None:
            self._heat_lag = float(heater_on)
        if self._previous is not None:
            dt = (now - self._previous["time"]).total_seconds()
            # Do not fit one transition across a provider/fallback switch.
            if (270 <= dt <= 330 and
                    heater.last_changed <= self._previous["time"] and
                    heater_on == self._previous["heater_on"] and
                    self._previous["outside_source"] == active_sources["temperature"]):
                next_lag = self._heat_lag + (1 - exp(-dt / 1200)) * (float(self._previous["heater_on"]) - self._heat_lag)
                for room in self.rooms:
                    self.learners[room].learn(
                        self._previous["temps"][room], vals[room],
                        self._previous["outside"], next_lag, dt, self._previous["heater_on"]
                    )
                self._heat_lag = next_lag
                self._update_counts()
                # Persist at each sample so restart does not erase hours of training.
                await self._store.async_save({
                    "schema_version": 2,
                    "learners": {room: learner.as_dict() for room, learner in self.learners.items()},
                    "on_samples": self._on_samples, "off_samples": self._off_samples,
                    "heater_entity": self.config[CONF_HEATER],
                    "outside_source": active_sources["temperature"],
                })
            else:
                # Model lag state becomes uncertain after long gaps.
                self._heat_lag = float(heater_on)

        self._previous = {
            "time": now, "temps": vals, "outside": outside, "heater_on": heater_on,
            "outside_source": active_sources["temperature"],
        }
        weights = self.config.get(CONF_WEIGHTS, {})
        try:
            weights = {r: float(weights.get(r, 1.0)) for r in self.rooms}
            if any(not isfinite(w) or w <= 0 for w in weights.values()) or not isfinite(sum(weights.values())):
                raise ValueError
        except (AttributeError, TypeError, ValueError, OverflowError):
            status["status"] = "invalid_weights"
            return status
        weighted_temp = sum(weights[r] * vals[r] for r in self.rooms) / sum(weights.values())
        status.update({
            "weighted_temperature": round(weighted_temp, 3),
            "outside_temperature": outside,
            "reference_heat": heater_on,
            "on_samples": self._on_samples,
            "off_samples": self._off_samples,
            "samples": min(learner.count for learner in self.learners.values()),
            "models": {
                room: {"loss": round(model.theta[0], 6),
                       "heat": round(model.theta[1], 6),
                       "bias": round(model.theta[2], 6),
                       "samples": model.count,
                       "on_samples": model.on_count, "off_samples": model.off_count,
                       "mae_5m": round(model.mae, 4) if model.mae is not None else None}
                for room, model in self.learners.items()
            },
        })
        if climate.state != "heat":
            status["status"] = "paused"
            return status
        try:
            target = normalize("temperature", climate.attributes["temperature"],
                               climate.attributes.get("temperature_unit",
                                   self.hass.config.units.temperature_unit))
            if target is None or not 5 <= target <= 35:
                raise ValueError
        except (TypeError, KeyError, ValueError):
            status["status"] = "waiting_for_target"
            return status
        status["target"] = target
        # Do not act on a cold model: need ~17 hours of samples plus both conditions.
        if status["samples"] < 200 or self._on_samples < 30 or self._off_samples < 30:
            status["status"] = "learning"
            return status
        try:
            # 180-minute exhaustive search can take several seconds on RPi5;
            # never block Home Assistant's event loop with CPU work.
            job = partial(
                optimize,
                vals, outside, self._heat_lag, target, weights, self.learners,
                heater_on, (now - heater.last_changed).total_seconds(),
                horizon_minutes=int(self.config[CONF_HORIZON]),
                minimum_on_minutes=int(self.config[CONF_MIN_ON]),
                minimum_off_minutes=int(self.config[CONF_MIN_OFF]),
                switching_cost=float(self.config[CONF_SWITCH_COST]),
                heating_cost=float(self.config[CONF_HEAT_COST]),
                comfort_band=float(self.config[CONF_BAND]),
            )
            # HA State objects are immutable snapshots. Discard a result if any
            # input changed while the worker was calculating the hypothetical plan.
            output = await self.hass.async_add_executor_job(job)
            if any(self.hass.states.get(entity) is not state for entity, state in snapshot.items()):
                status["status"] = "inputs_changed"
                return status
        except (ValueError, RuntimeError, OverflowError) as err:
            _LOGGER.warning("MPC prediction failed: %s", err)
            status["status"] = "prediction_error"
            return status
        status.update({
            "status": "ready", "proposal": output["proposal"],
            "prediction_60m": output["predicted_60m"],
            "prediction_end": output["predicted_end"],
            "plan_15min": output["plan_15min"], "objective": output["cost"],
            "model_quality_note": "5-minute in-sample MAE; not a validated 60-minute forecast error",
        })
        return status