# WashData - Home Assistant integration for appliance cycle monitoring via smart plugs.
# Copyright (C) 2026 Lukas Bandura
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Opt-in ML scoring bridge for WashData (experimental).

This package holds compact, NumPy-only models trained offline in the
``ml_washdata`` lab and embedded here as base64 blobs (see
``promoted_manifest.json`` for provenance). The integration runtime stays
NumPy-only; no sklearn/torch/scipy are imported.

The runtime entry points are :func:`resolve_scorer`, which returns the shipped
embedded baseline's scoring callable for a classifier capability, and
:func:`resolve_regressor`, which returns an on-device trained regressor. All live
ML consumers go through them (the panel's ``ml_health`` shadow comparison in
``ws_api``), and any new runtime consumer should too - feature extraction lives in ``feature_extraction``
and gating in :func:`ml_models_enabled`, so there is no separate engine object.

Each model consumes a feature mapping whose keys are the model's
``FEATURE_COLUMNS``; the integration computes those from live data per the
``*_feature_contract.json`` files shipped alongside the model modules.
"""

from __future__ import annotations

import importlib
import json
import logging
from pathlib import Path
from typing import Any, Mapping

_LOGGER = logging.getLogger(__name__)

CONF_ENABLE_ML_MODELS = "enable_ml_models"

# Logical capability -> generated model module name (without the _model suffix).
_MODEL_MODULES = {
    "quality": "hybrid_curve_quality_model",
    "end": "cycle_end_detector_model",
}

# Modules the live ML paths import lazily besides the baselines themselves.
_SIBLING_MODULES = ("trainer", "feature_extraction")

# Imported baseline model modules, keyed by module name. Importing a module is a
# blocking call Home Assistant forbids inside the event loop, and
# resolve_scorer() consumers (end detection, the panel's cycle health) run
# there - so the modules are imported once from an import executor at setup
# (:func:`preload_models`) and every later resolution is a dict lookup. A failed
# import is cached as ``None`` so a broken install warns once instead of retrying
# the import on every inference.
_MODULE_CACHE: dict[str, object | None] = {}


def _load_model_module(module_name: str) -> Any | None:
    """Return the embedded model module, importing it at most once.

    Safe to call from the event loop *after* :func:`preload_models` has run (the
    import is then already satisfied from ``sys.modules``); the first call should
    happen in an executor thread.
    """
    if module_name in _MODULE_CACHE:
        return _MODULE_CACHE[module_name]
    try:
        module = importlib.import_module(f"{__package__}.{module_name}")
    except Exception as exc:  # noqa: BLE001 - a missing model must not break setup
        _LOGGER.warning(
            "Failed to load embedded model module %r: %s", module_name, exc
        )
        module = None
    _MODULE_CACHE[module_name] = module
    return module


def _sibling_attr(module_name: str, attr: str) -> Any | None:
    """Fetch ``attr`` from an embedded sibling module via the cache, or None.

    Resolution paths run in the event loop, so they must never re-import: this hits
    the ``_MODULE_CACHE`` warmed by :func:`preload_models` (which stores ``None`` on a
    failed import). A missing module or attribute returns None, and the caller falls
    back to the baseline / inert path rather than triggering a blocking loop import.
    """
    module = _load_model_module(module_name)
    return getattr(module, attr, None) if module is not None else None


def preload_models() -> None:
    """Import everything the live ML paths touch. Call from an executor thread.

    ``resolve_scorer`` / ``resolve_regressor`` are called from the event loop, so
    the imports they need (the embedded baselines plus ``trainer`` /
    ``feature_extraction``) must already be in ``sys.modules`` by then - Home
    Assistant flags a blocking ``importlib.import_module`` in the loop (issue
    #328). Also warms the manifest cache, which reads a file. Never raises;
    idempotent, so calling it once per config entry is cheap.
    """
    for module_name in _MODEL_MODULES.values():
        _load_model_module(module_name)
    # Cache the siblings too (module-or-None), so a failed import is recorded once
    # here and the event-loop resolvers read it from the cache instead of retrying
    # a blocking import (issue #328).
    for sibling in _SIBLING_MODULES:
        _load_model_module(sibling)
    # No guard needed: available_models() carries its own outer try/except and caches
    # [] on every failure path, so it cannot raise here (and a try/except/pass around it
    # would be unreachable code that Ruff flags as S110/SIM105).
    available_models()


def ml_models_enabled(options: Mapping[str, object] | None) -> bool:
    """True when the user has opted into experimental ML models."""
    if not options:
        return False
    return bool(options.get(CONF_ENABLE_ML_MODELS, False))


def resolve_scorer(capability: str):
    """Return ``(score_fn, source)`` for a classifier capability's shipped baseline.

    ``score_fn`` maps a feature mapping -> float in [0,1]; ``source`` is
    ``"baseline"``. Returns ``(None, None)`` when the capability has no embedded
    model or it failed to import. The lookup hits the module cache warmed by
    :func:`preload_models`, so no import happens in the event loop. (The
    on-device classifier specs this used to prefer can no longer exist: their
    training was removed in 0.5.8 and storage v17 drops their records.)
    """
    module_name = _MODEL_MODULES.get(capability)
    if module_name is None:
        return (None, None)
    module = _load_model_module(module_name)
    if module is None:
        return (None, None)

    def _baseline_score(feats, _m=module):
        # The embedded baseline must never raise into live inference: on any
        # scoring error log and return a neutral 0.0 so a gate treats the signal
        # as absent rather than letting the exception reach detection/matching.
        try:
            return float(_m.score(feats))
        except Exception as exc:  # noqa: BLE001 - never raise into live inference
            _LOGGER.warning(
                "Embedded baseline scorer for capability %r failed at call "
                "time, returning neutral 0.0: %s", capability, exc,
            )
            return 0.0

    return (_baseline_score, "baseline")


def resolve_regressor(capability: str, store: object | None):
    """Return ``(predict_fn, source)`` for a regression capability.

    Regression models (since 0.5.8 only ``"total_energy"``) have **no**
    shipped embedded baseline - they are trained purely on-device (Stage 4) and
    stored as ``standardized_linear`` specs. This returns ``(None, None)`` until
    on-device training promotes one, so live behaviour is unchanged until then.

    ``predict_fn`` maps a feature mapping -> float in the model's target units
    (for ``total_energy``, the energy fraction so far in ~[0, 1]).
    """
    if store is None:
        return (None, None)
    try:
        versions = store.get_ml_model_versions() or {}  # type: ignore[attr-defined]
        record = versions.get(capability)
        spec = record.get("spec") if isinstance(record, dict) else None
        if isinstance(spec, dict) and spec.get("kind") == "standardized_linear":
            # Feature-column schema guard for regression specs.
            try:
                _prog_cols = _sibling_attr("feature_extraction", "PROGRESS_FEATURE_COLUMNS")
                _expected_r = list(_prog_cols or [])
                _stored_r = list(spec.get("feature_columns") or [])
                if _expected_r and _stored_r and _stored_r != _expected_r:
                    _LOGGER.warning(
                        "Promoted regression spec for %r has stale feature schema "
                        "(%d cols vs current %d); reverting to inert.",
                        capability, len(_stored_r), len(_expected_r),
                    )
                    return (None, None)
            except Exception:  # noqa: BLE001 - schema check must not break inference
                pass

            predict_value_spec = _sibling_attr("trainer", "predict_value_spec")
            if predict_value_spec is None:
                return (None, None)

            def _on_device_predict(feats, _s=spec):
                # A malformed / incompatible promoted regression spec must never
                # raise into the live remaining-time / energy estimates: on any
                # call-time error return NaN so the (isfinite-guarded) consumers
                # treat this capability as inert.
                try:
                    return float(predict_value_spec(_s, feats))
                except Exception as exc:  # noqa: BLE001 - never raise into live inference
                    _LOGGER.warning(
                        "Trained regressor for capability %r failed at call time, "
                        "returning inert value: %s", capability, exc,
                    )
                    return float("nan")

            return (_on_device_predict, "on_device")
    except Exception as exc:  # noqa: BLE001 - never let a bad store break inference
        _LOGGER.warning(
            "Failed to load trained regression spec for capability %r, capability will be inert: %s",
            capability, exc,
        )
    return (None, None)


_MANIFEST_MODELS_CACHE: list[dict[str, object]] | None = None


def available_models() -> list[dict[str, object]]:
    """Return provenance for the embedded models, or [] if none are shipped.

    The manifest is a shipped baseline file that never changes at runtime
    (on-device training writes specs into the store, not this file), so the parsed
    result is cached module-side after the first read.
    """
    global _MANIFEST_MODELS_CACHE
    if _MANIFEST_MODELS_CACHE is not None:
        return _MANIFEST_MODELS_CACHE
    # Outer guard so EVERY failure caches a result: an unhandled exception here (from
    # warm-up in preload_models, whose caller swallows it) would leave the cache cold,
    # and the next event-loop caller would retry Path.exists()/read_text() - re-creating
    # the blocking-call warning preload exists to prevent (#328).
    try:
        manifest = Path(__file__).resolve().parent / "promoted_manifest.json"
        # A missing manifest is cached as [] too: this is a shipped file that cannot
        # appear at runtime, and the read is a blocking open() some callers make on
        # the event loop.
        if not manifest.exists():
            result: list[dict[str, object]] = []
        else:
            payload = json.loads(manifest.read_text(encoding="utf-8"))
            # A manifest decoding to a list/scalar would make .get() raise; keep only
            # dict model entries so the return honours its list[dict] contract even for
            # a malformed manifest like {"models": [null]}.
            models = payload.get("models") if isinstance(payload, dict) else None
            result = [m for m in models if isinstance(m, dict)] if isinstance(models, list) else []
    except Exception as exc:  # noqa: BLE001 - never raise / never leave the cache cold
        _LOGGER.debug("Could not read the promoted model manifest (%s); caching empty", exc)
        result = []
    _MANIFEST_MODELS_CACHE = result
    return result