"""Battery-fleet aggregation over the Battery Notes integration.

The user does NOT want one maintenance task per battery (30-70+ devices would
bury the task list). Instead this aggregates every Battery Notes ``battery_plus``
sensor into ONE fleet view: which batteries are low now, grouped by battery
type (so you know *what to buy*), plus a simple deterministic forecast of what
will be needed soon (so you can order in time).

Battery Notes exposes everything we need as ATTRIBUTES on the single
``battery_plus`` sensor (device_class ``battery``): ``battery_type``,
``battery_quantity``, ``battery_low``, ``battery_low_threshold``,
``battery_last_replaced``. We read that one sensor kind — no dependency on the
(optional, often-disabled) battery-low binary.

The pure builder ``build_overview`` takes plain battery dicts + an injected
``today`` so the forecast is unit-testable with synthetic dates; ``read_batteries``
is the thin HA-reading adapter.
"""

from __future__ import annotations

from collections import OrderedDict
from dataclasses import dataclass, field
from datetime import date
from typing import Any

from homeassistant.core import HomeAssistant
from homeassistant.util import dt as dt_util

# Editorial typical service life per battery type, in MONTHS — the forecast
# anchor (battery_last_replaced + lifetime = predicted replacement). These are
# deliberately conservative sensor-use estimates and are meant to be tunable;
# unknown types fall back to DEFAULT_LIFETIME_MONTHS.
TYPICAL_LIFETIME_MONTHS: dict[str, int] = {
    "AAAA": 10,
    "AAA": 10,
    "AA": 12,
    "C": 18,
    "D": 24,
    "9V": 12,
    "CR2": 18,
    "CR123A": 18,
    "CR2032": 18,
    "CR2450": 24,
    "CR2477": 24,
    "CR2016": 18,
    "CR2025": 18,
}
DEFAULT_LIFETIME_MONTHS = 12

# How far ahead "needed soon" looks by default (days).
DEFAULT_HORIZON_DAYS = 28

# A native battery %-sensor without a dedicated low binary is treated as low
# at or below this level (editorial — HA has no universal low threshold).
NATIVE_LOW_PERCENT = 20

# States that mean "no reading" — a removed device leaves nothing; an offline
# one leaves these. We keep an OFFLINE battery visible only when its last-known
# low flag says it needs attention (a dead battery often takes its device
# offline — that's exactly the one you must not hide).
_NO_READING = {"unavailable", "unknown", "none", ""}


def _norm_type(raw: Any) -> str:
    """Canonicalize a battery-type label for grouping (upper, trimmed)."""
    s = str(raw or "").strip()
    return s.upper() if s else "UNKNOWN"


def lifetime_months(battery_type: str) -> int:
    """Typical service life for a (canonicalized) battery type."""
    return TYPICAL_LIFETIME_MONTHS.get(_norm_type(battery_type), DEFAULT_LIFETIME_MONTHS)


@dataclass
class Battery:
    """One battery-powered device — from Battery Notes (rich) or a native
    ``device_class: battery`` entity (degraded: no type/quantity/forecast)."""

    entity_id: str
    device_name: str
    battery_type: str
    quantity: int
    low: bool
    level: float | None
    last_replaced: date | None
    available: bool = True
    source: str = "battery_notes"


@dataclass
class BatteryOverview:
    """The aggregated fleet view backing the single fleet task + its detail."""

    total: int = 0
    low: list[dict[str, Any]] = field(default_factory=list)
    soon: list[dict[str, Any]] = field(default_factory=list)
    # Grouped quantities by canonical type.
    needs_now: OrderedDict[str, int] = field(default_factory=OrderedDict)
    needs_soon: OrderedDict[str, int] = field(default_factory=OrderedDict)
    types: list[str] = field(default_factory=list)

    @property
    def low_count(self) -> int:
        return len(self.low)


def _predicted_date(bat: Battery) -> date | None:
    if bat.last_replaced is None:
        return None
    months = lifetime_months(bat.battery_type)
    # Month arithmetic without dateutil: add whole months, clamp the day.
    y, m = bat.last_replaced.year, bat.last_replaced.month + months
    y += (m - 1) // 12
    m = (m - 1) % 12 + 1
    day = min(bat.last_replaced.day, 28)
    return date(y, m, day)


def build_overview(
    batteries: list[Battery],
    *,
    today: date,
    horizon_days: int = DEFAULT_HORIZON_DAYS,
) -> BatteryOverview:
    """Aggregate batteries into the fleet view.

    * ``low`` = reported low right now (Battery Notes' own threshold).
    * ``soon`` = NOT low yet but predicted to reach end-of-life within
      ``horizon_days`` (deterministic last_replaced + typical-lifetime forecast).
      A battery already low is never double-counted into soon.
    * ``needs_now`` / ``needs_soon`` = summed quantities per type — the shopping
      grouping ("2× AA, 4× AAA").
    """
    ov = BatteryOverview(total=len(batteries))
    types_seen: OrderedDict[str, None] = OrderedDict()

    for bat in sorted(batteries, key=lambda b: b.device_name.lower()):
        t = _norm_type(bat.battery_type)
        types_seen[t] = None
        if bat.low:
            ov.low.append(_row(bat, t, None))
            ov.needs_now[t] = ov.needs_now.get(t, 0) + bat.quantity
            continue
        pred = _predicted_date(bat)
        if pred is not None:
            days = (pred - today).days
            if days <= horizon_days:
                ov.soon.append(_row(bat, t, days))
                ov.needs_soon[t] = ov.needs_soon.get(t, 0) + bat.quantity

    ov.soon.sort(key=lambda r: r["days_until"] if r["days_until"] is not None else 1 << 30)
    ov.types = sorted(types_seen)
    ov.needs_now = OrderedDict(sorted(ov.needs_now.items()))
    ov.needs_soon = OrderedDict(sorted(ov.needs_soon.items()))
    return ov


def _row(bat: Battery, canon_type: str, days_until: int | None) -> dict[str, Any]:
    return {
        "entity_id": bat.entity_id,
        "device_name": bat.device_name,
        "battery_type": canon_type,
        "quantity": bat.quantity,
        "level": bat.level,
        "days_until": days_until,
        "available": bat.available,
    }


def _parse_last_replaced(raw: Any) -> date | None:
    if not raw:
        return None
    try:
        parsed = dt_util.parse_datetime(str(raw))
        if parsed is not None:
            return parsed.date()
        return date.fromisoformat(str(raw)[:10])
    except (ValueError, TypeError):
        return None


def _level_of(state_val: str) -> float | None:
    try:
        return float(state_val)
    except (ValueError, TypeError):
        return None


def read_batteries(hass: HomeAssistant) -> list[Battery]:
    """Read the battery fleet from HA state — Battery Notes AND native.

    * **Battery Notes** ``battery_plus`` sensors (device_class ``battery`` + a
      ``battery_type`` attribute) give the rich view: type, quantity, low,
      last-replaced. When the source goes offline the sensor reads
      unavailable/unknown but RETAINS its last-known ``battery_low`` — so a
      dead battery that took its device offline stays visible.
    * **Native** ``device_class: battery`` entities (a %-sensor and/or a
      battery-low binary), grouped per device, give a degraded view (type
      "Unknown", quantity 1, no forecast). A device already covered by a
      Battery Notes note is skipped (dedup by the note's source entity + its
      device) so it isn't counted twice.
    """
    from homeassistant.helpers import device_registry as dr
    from homeassistant.helpers import entity_registry as er

    ent_reg = er.async_get(hass)
    dev_reg = dr.async_get(hass)

    out: list[Battery] = []
    covered_sources: set[str] = set()
    covered_devices: set[str] = set()

    # ── Pass 1: Battery Notes battery_plus ──────────────────────────────────
    for state in hass.states.async_all("sensor"):
        attrs = state.attributes
        if attrs.get("device_class") != "battery" or "battery_type" not in attrs:
            continue
        src = attrs.get("source_entity_id")
        if src:
            covered_sources.add(src)
        reg = ent_reg.async_get(state.entity_id)
        if reg and reg.device_id:
            covered_devices.add(reg.device_id)
        level = _level_of(state.state)
        available = state.state not in _NO_READING and level is not None
        low = bool(attrs.get("battery_low"))
        # Offline AND not last-known-low = pure connectivity noise → drop it.
        if not available and not low:
            continue
        out.append(
            Battery(
                entity_id=state.entity_id,
                device_name=attrs.get("device_name") or attrs.get("friendly_name") or state.entity_id,
                battery_type=str(attrs.get("battery_type") or "Unknown"),
                quantity=int(attrs.get("battery_quantity") or 1),
                low=low,
                level=level,
                last_replaced=_parse_last_replaced(attrs.get("battery_last_replaced")),
                available=available,
                source="battery_notes",
            )
        )

    # ── Pass 2: native battery entities, grouped per device ─────────────────
    # {group_key: {"level_state": s, "low_state": s, "name": ..., "device_id": ..., "eid": ...}}
    native: dict[str, dict[str, Any]] = {}
    for domain in ("sensor", "binary_sensor"):
        for state in hass.states.async_all(domain):
            if state.attributes.get("device_class") != "battery":
                continue
            eid = state.entity_id
            if "battery_type" in state.attributes:  # Battery Notes battery_plus — handled above
                continue
            if eid in covered_sources:
                continue
            reg = ent_reg.async_get(eid)
            dev_id = reg.device_id if reg else None
            if dev_id and dev_id in covered_devices:
                continue
            key = dev_id or eid
            rec = native.setdefault(
                key,
                {"level_state": None, "low_state": None, "device_id": dev_id, "eid": eid, "name": None},
            )
            friendly = state.attributes.get("friendly_name")
            if domain == "sensor":
                rec["level_state"] = state.state
                rec["eid"] = eid
            else:
                rec["low_state"] = state.state
            if rec["name"] is None and friendly:
                rec["name"] = friendly

    for rec in native.values():
        level = _level_of(rec["level_state"]) if rec["level_state"] is not None else None
        low_state = rec["low_state"]
        level_available = rec["level_state"] not in _NO_READING if rec["level_state"] is not None else False
        low_available = low_state not in _NO_READING if low_state is not None else False
        available = level_available or low_available
        if low_state is not None:
            low = low_available and str(low_state).lower() in ("on", "true", "1")
        else:
            low = level is not None and level <= NATIVE_LOW_PERCENT
        if not available and not low:
            continue
        name = rec["name"]
        if not name and rec["device_id"] and (dev := dev_reg.async_get(rec["device_id"])):
            name = dev.name_by_user or dev.name
        out.append(
            Battery(
                entity_id=rec["eid"],
                device_name=name or rec["eid"],
                battery_type="Unknown",
                quantity=1,
                low=low,
                level=level,
                last_replaced=None,
                available=available,
                source="native",
            )
        )
    return out


def has_battery_notes(hass: HomeAssistant) -> bool:
    """Whether the Battery Notes integration is present (any battery_plus)."""
    for state in hass.states.async_all("sensor"):
        a = state.attributes
        if a.get("device_class") == "battery" and "battery_type" in a:
            return True
    return False


def has_batteries(hass: HomeAssistant) -> bool:
    """Whether ANY battery is trackable — Battery Notes OR native. Gates setup."""
    for domain in ("sensor", "binary_sensor"):
        for state in hass.states.async_all(domain):
            if state.attributes.get("device_class") == "battery":
                return True
    return False


def compute_overview(hass: HomeAssistant, *, horizon_days: int = DEFAULT_HORIZON_DAYS) -> BatteryOverview:
    """Read + aggregate in one call (HA-side entry point)."""
    today = dt_util.now().date()
    return build_overview(read_batteries(hass), today=today, horizon_days=horizon_days)


def discover_battery_types(hass: HomeAssistant) -> OrderedDict[str, int]:
    """Battery types present across the fleet → total quantity, for part setup."""
    totals: OrderedDict[str, int] = OrderedDict()
    for bat in read_batteries(hass):
        t = _norm_type(bat.battery_type)
        totals[t] = totals.get(t, 0) + bat.quantity
    return OrderedDict(sorted(totals.items()))


__all__ = [
    "DEFAULT_HORIZON_DAYS",
    "NATIVE_LOW_PERCENT",
    "TYPICAL_LIFETIME_MONTHS",
    "Battery",
    "BatteryOverview",
    "build_overview",
    "compute_overview",
    "discover_battery_types",
    "has_batteries",
    "has_battery_notes",
    "lifetime_months",
    "read_batteries",
]
