"""Probe identity: learn hot-in vs tub ROM ids, setup state machine."""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal

SetupState = Literal["need_hot_in", "need_water", "ready", "need_reinstall"]

TUB_COUNT = 4


@dataclass
class ProbeMap:
    hot_in: str | None = None
    tub: list[str] = field(default_factory=list)

    def is_complete(self) -> bool:
        return self.hot_in is not None and len(self.tub) == TUB_COUNT

    def known_ids(self) -> set[str]:
        ids: set[str] = set()
        if self.hot_in:
            ids.add(self.hot_in)
        ids.update(self.tub)
        return ids


def load_probe_map(path: Path) -> ProbeMap:
    if not path.exists():
        return ProbeMap()
    try:
        data = json.loads(path.read_text())
    except (OSError, json.JSONDecodeError):
        return ProbeMap()
    hot_in = data.get("hot_in")
    tub = data.get("tub") or []
    if not isinstance(hot_in, str) and hot_in is not None:
        hot_in = None
    if not isinstance(tub, list):
        tub = []
    tub = [x for x in tub if isinstance(x, str)]
    return ProbeMap(hot_in=hot_in, tub=tub)


def save_probe_map(path: Path, probe_map: ProbeMap) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {"hot_in": probe_map.hot_in, "tub": sorted(probe_map.tub)}
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(payload, indent=2) + "\n")
    tmp.replace(path)


@dataclass
class IdentityResult:
    setup_state: SetupState
    probe_map: ProbeMap
    map_changed: bool
    detail: str
    visible_ids: list[str]


def _try_learn(
    visible: list[str],
    pmap: ProbeMap,
    *,
    setup_mode: bool,
) -> tuple[ProbeMap, bool, str]:
    """Return updated map, whether it changed, and a detail string."""
    n = len(visible)
    vis = set(visible)
    changed = False
    detail = ""

    if pmap.is_complete():
        return pmap, False, detail

    learning = setup_mode or not pmap.is_complete()
    if not learning:
        return pmap, False, detail

    hot_known = pmap.hot_in is not None
    tub_complete = len(pmap.tub) == TUB_COUNT

    if not hot_known and n == 1:
        pmap = ProbeMap(hot_in=visible[0], tub=list(pmap.tub))
        return pmap, True, f"Learned hot_in={visible[0]}"

    if not tub_complete and not hot_known and n == TUB_COUNT and setup_mode:
        pmap = ProbeMap(hot_in=pmap.hot_in, tub=sorted(visible))
        return pmap, True, f"Learned {TUB_COUNT} tub probes"

    if hot_known and not tub_complete and n == TUB_COUNT + 1:
        tub_ids = sorted(vis - {pmap.hot_in})
        if len(tub_ids) == TUB_COUNT:
            pmap = ProbeMap(hot_in=pmap.hot_in, tub=tub_ids)
            return pmap, True, "Derived tub ids (hot_in already known)"

    if tub_complete and not hot_known and n == TUB_COUNT + 1:
        unknown = vis - set(pmap.tub)
        if len(unknown) == 1:
            hid = next(iter(unknown))
            pmap = ProbeMap(hot_in=hid, tub=list(pmap.tub))
            return pmap, True, f"Derived hot_in={hid}"

    if hot_known and not tub_complete and n == TUB_COUNT and pmap.hot_in not in vis:
        pmap = ProbeMap(hot_in=pmap.hot_in, tub=sorted(visible))
        return pmap, True, "Learned 4 tub probes (water bundle alone)"

    return pmap, changed, detail


def _compute_setup_state(visible: list[str], pmap: ProbeMap) -> SetupState:
    n = len(visible)
    if n == TUB_COUNT + 1 and not pmap.hot_in and len(pmap.tub) == 0:
        return "need_reinstall"
    if pmap.is_complete():
        return "ready"
    if pmap.hot_in and len(pmap.tub) < TUB_COUNT:
        return "need_water"
    if len(pmap.tub) == TUB_COUNT and not pmap.hot_in:
        return "need_hot_in"
    if not pmap.hot_in and not pmap.tub:
        if n >= TUB_COUNT + 1:
            return "need_reinstall"
        if n == TUB_COUNT:
            return "need_hot_in"
        if n == 1:
            return "need_hot_in"
        return "need_hot_in"
    return "need_water"


def evaluate_identity(
    visible_ids: list[str],
    probe_map: ProbeMap,
    *,
    setup_mode: bool,
) -> IdentityResult:
    visible = sorted(set(visible_ids))
    pmap, changed, detail = _try_learn(visible, probe_map, setup_mode=setup_mode)
    state = _compute_setup_state(visible, pmap)
    if state == "need_reinstall" and not detail:
        detail = "All 5 probes visible but no stored ids; connect bundles one at a time"
    if not detail:
        detail = state
    return IdentityResult(
        setup_state=state,
        probe_map=pmap,
        map_changed=changed,
        detail=detail,
        visible_ids=visible,
    )


def classify_readings(
    probes_f: dict[str, float | None],
    probe_map: ProbeMap,
) -> tuple[dict[str, float | None], float | None, float | None, int, int, float | None]:
    """Return labeled probes, tub_avg, hot_in_f, ok/fail counts, spread."""
    labeled: dict[str, float | None] = dict(probes_f)
    tub_temps: list[float] = []
    hot_in_f: float | None = None

    tub_set = set(probe_map.tub)
    for device_id, temp_f in probes_f.items():
        if device_id == probe_map.hot_in and temp_f is not None:
            hot_in_f = temp_f
        elif device_id in tub_set and temp_f is not None:
            tub_temps.append(temp_f)

    ok_count = len(tub_temps)
    fail_count = max(0, len(probe_map.tub) - ok_count) if probe_map.tub else 0
    tub_avg: float | None = sum(tub_temps) / len(tub_temps) if tub_temps else None
    spread: float | None = None
    if len(tub_temps) >= 2:
        spread = max(tub_temps) - min(tub_temps)

    return labeled, tub_avg, hot_in_f, ok_count, fail_count, spread