"""Wren: a typed command in, values for the devices wired to it out.

`Wren <https://github.com/stroblme/wren>`_ is a small language model trained on
one job: read a house — its areas, the devices in each and their states — and
answer a command with calls on those devices. This node is the bridge. Its
output ports *are* the devices, named ``<area>_<kind>``, so the wiring is the
allowlist: the model can reach nothing that is not wired to the node.
"""

from __future__ import annotations

import time
from collections.abc import Iterable
from typing import Any

from pydantic import BaseModel, ConfigDict, Field
from threadpoolctl import threadpool_limits

from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes.base import Node


def _wren() -> Any:
    """The model package, which is optional.

    Imported when a node is built rather than at boot, so an instance without
    it still offers every other type, and a node placed without it says so
    on the canvas straight away.
    """
    try:
        import wren
    except ImportError:
        raise RuntimeError("the wren node needs the model: pip install wren") from None
    return wren


#: The last word of a target port, and the kind of device it stands for.
KIND_WORDS = {
    "light": "light",
    "lamp": "light",
    "brightness": "light",
    "setpoint": "heat",
    "heating": "heat",
    "heat": "heat",
    "boost": "heat",
    "media": "media",
    "music": "media",
    "tv": "media",
    "shutter": "shutter",
    "blind": "shutter",
    "cover": "shutter",
    "water": "water",
    "watering": "water",
    "sprinkler": "water",
}

#: What the node reads, and the types each may be. Only the prompt wakes it.
INPUTS = {
    "prompt": {DType.STR},
    "presence": {DType.STR},
    "outdoor": {DType.FLOAT, DType.INT},
    "dark": {DType.BOOL},
}

#: The output the calls are reported on. Every other output is a device.
ACTIONS = "actions"

_TARGET_TYPES = {DType.BOOL, DType.INT, DType.FLOAT, DType.STR}

#: The tools that act on each kind.
_TOOLS = {
    "light": "lights",
    "media": "media",
    "heat": "heating",
    "shutter": "shutters",
    "water": "watering",
}

#: What a call sets a target to, by the target's type: (bool, number, text).
#: None is a call that type has no value for, which is then not applied.
#: `heating.set` carries its own number.
_VALUES: dict[str, tuple[bool | None, float | None, str | None]] = {
    "lights.on": (True, 100, "on"),
    "lights.off": (False, 0, "off"),
    "media.play": (True, None, "play"),
    "media.stop": (False, None, "stop"),
    "heating.boost": (True, None, "boost"),
    "heating.off": (False, None, "off"),
    "shutters.up": (True, 100, "up"),
    "shutters.down": (False, 0, "down"),
    "shutters.stop": (None, None, "stop"),
    "watering.start": (True, None, "on"),
    "watering.stop": (False, None, "off"),
}

#: What every area held in training: a light, and in the living room, which
#: always existed, media and heating besides. A house missing them reads as
#: unfamiliar and the confidence drops with it, so they are filled in; a call
#: on one that nothing is wired to is reported and not applied.
_TRAINED = {"living room": ("light", "media", "heat")}


def _device(port: str, areas: dict[str, str]) -> tuple[str, str] | None:
    """The (area, kind) a port name stands for, or None."""
    for slug, area in areas.items():
        word = port.removeprefix(f"{slug}_")
        if word != port and word in KIND_WORDS:
            return area, KIND_WORDS[word]
    return None


def _on(value: Any) -> bool:
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return value > 0
    return str(value).strip().lower() in {"on", "up", "open", "play", "true"}


def _state(kind: str, values: list[Any]) -> str:
    """A device's state as the model reads it, from what its targets hold."""
    if kind == "heat":
        temps = [
            v for v in values if isinstance(v, (int, float)) and not isinstance(v, bool)
        ]
        return f"{round(temps[0])}C" if temps else "20C"
    on = bool(values) and _on(values[0])
    if kind == "shutter":
        return "up" if on else "down"
    return "on" if on else "off"


def _value(call: Any, dtype: DType) -> Any:
    """What ``call`` sets a target of this type to, or None."""
    if call.tool == "heating.set":
        temp = call.args.get("temp")
        flag, number, text = (
            None,
            temp if isinstance(temp, (int, float)) else None,
            None,
        )
    else:
        flag, number, text = _VALUES.get(call.tool, (None, None, None))
    if dtype is DType.BOOL:
        return flag
    if dtype is DType.STR:
        return text
    if number is None:
        return None
    return int(number) if dtype is DType.INT else float(number)


def _text(call: Any) -> str:
    """The call as the chat shows it, less the device it names."""
    args = ", ".join(f"{k}={v}" for k, v in call.args.items() if k != "id")
    return f"{call.tool}({args})"


class WrenNode(Node):
    """Turn the command on ``prompt`` into values for the devices wired to it.

    Inputs: ``prompt`` (the command, and the only one that wakes the node),
    ``presence`` (the area the user is in), ``outdoor`` (°C) and ``dark``
    (otherwise read off the clock: before 07:00 or from 18:00).

    Outputs: ``actions``, a list of records, one per call and target, and
    any number of targets named ``<area>_<kind>``. Targets of one area and kind
    are one device to the model; each takes what the call means for its type.
    A delayed call is held by the engine's work queue, like the Trigger
    node's second value.
    """

    class Params(BaseModel):
        model_config = ConfigDict(extra="allow")

        apply: bool = Field(
            default=False,
            description="Write to the targets. Off, the node only reports what it would do.",
        )
        min_confidence: float = Field(
            default=0.93,
            ge=0,
            le=1,
            description="Below this, nothing is written; the calls are still reported.",
        )

    __slots__ = ("cfg", "devices")

    # A call is a command to a device; doing it twice is not doing it once.
    idempotent = False

    def __init__(
        self,
        requires: MessageSpec | Iterable[MessageSpec] = (),
        provides: MessageSpec | Iterable[MessageSpec] = (),
        params: dict[str, Any] | None = None,
        name: str | None = None,
    ):
        self.cfg = self.Params.model_validate(params or {})
        wren = _wren()
        inputs = self._normalize_ports(requires)
        outputs = self._normalize_ports(provides)

        for spec in inputs:
            allowed = INPUTS.get(spec.port)
            if allowed is None:
                raise ValueError(
                    f"'{spec.port}' is not an input of wren. It reads "
                    f"{', '.join(INPUTS)}."
                )
            if spec.dtype not in allowed:
                raise ValueError(
                    f"'{spec.port}' is {spec.dtype.value}; wren reads it as "
                    f"{' or '.join(sorted(d.value for d in allowed))}."
                )
        if not any(spec.port == "prompt" for spec in inputs):
            raise ValueError("wren needs a 'prompt' input: the command to act on.")

        areas = {area.replace(" ", "_"): area for area in wren.AREAS}
        #: (area, kind) to the ports standing for that device.
        self.devices: dict[tuple[str, str], list[str]] = {}
        for spec in outputs:
            if spec.port == ACTIONS:
                if spec.dtype is not DType.LIST or spec.item_dtype is not DType.RECORD:
                    raise ValueError(f"'{ACTIONS}' is a list of records.")
                continue
            device = _device(spec.port, areas)
            if device is None:
                raise ValueError(
                    f"'{spec.port}' names no device wren knows. A target is "
                    f"<area>_<kind>, the area one of {', '.join(areas)}, the kind "
                    f"one of {', '.join(KIND_WORDS)}."
                )
            if spec.dtype not in _TARGET_TYPES:
                raise ValueError(
                    f"'{spec.port}' is {spec.dtype.value}; a target is bool, int, "
                    "float or str."
                )
            self.devices.setdefault(device, []).append(spec.port)

        super().__init__(
            self._act,
            # The prompt is the event; the rest is read when it arrives.
            requires=[
                spec.model_copy(update={"trigger": spec.port == "prompt"})
                for spec in inputs
            ],
            provides=outputs,
            params=params,
            name=name or "wren",
        )

    def _house(
        self,
        wren: Any,
        specs: dict[str, MessageSpec],
        presence: str | None,
        outdoor: float | None,
        dark: bool | None,
    ) -> Any:
        """The house as the model reads it, each device in its current state."""
        names = [specs[port].name for ports in self.devices.values() for port in ports]
        current = (
            self._pipeline.state.get_present(names)
            if self._pipeline is not None
            else {}
        )
        areas: dict[str, list[Any]] = {}
        for area in wren.AREAS:
            wired = {kind for (a, kind) in self.devices if a == area}
            if not wired and area not in _TRAINED:
                continue
            kinds = wired | set(_TRAINED.get(area, ("light",)))
            areas[area] = []
            for kind in (k for k in wren.KINDS if k in kinds):
                ports = self.devices.get((area, kind), [])
                held = [
                    current[specs[p].name] for p in ports if specs[p].name in current
                ]
                # The id is ours to choose: the model is shown aliases anyway.
                areas[area].append(
                    wren.Device(kind, f"{area}:{kind}", _state(kind, held))
                )
        clock = time.localtime()
        where = str(presence or "").strip().lower().replace("_", " ")
        return wren.House(
            time=time.strftime("%H:%M", clock),
            # The range the model was trained on; past it a number is unfamiliar.
            outdoor_c=min(max(float(20 if outdoor is None else outdoor), -5), 30),
            dark=clock.tm_hour < 7 or clock.tm_hour >= 18 if dark is None else dark,
            user_area=where if where in wren.AREAS else None,
            areas=areas,
        )

    def _act(
        self,
        params: dict[str, Any],
        prompt: str = "",
        presence: str | None = None,
        outdoor: float | None = None,
        dark: bool | None = None,
    ) -> dict[str, Any] | None:
        if not prompt.strip():
            # An empty command is nothing to do, which is what a flow input
            # starts as.
            return None
        wren = _wren()
        specs = {spec.port: spec for spec in self.output_ports}
        house = self._house(wren, specs, presence, outdoor, dark)
        # One BLAS thread answers in about half the time for a model this
        # small. Scoped to the call, so numpy elsewhere in the process keeps
        # its threads.
        with threadpool_limits(1):
            calls, confidence = wren.act(prompt, house)
        confidence = round(float(confidence), 3)
        write = self.cfg.apply and confidence >= self.cfg.min_confidence
        writes: dict[str, Any] = {}
        actions: list[dict[str, Any]] = []
        for call in calls:
            area, _, kind = str(call.args.get("id", "")).partition(":")
            ports = self.devices.get((area, kind), [])
            if call.tool.split(".")[0] != _TOOLS.get(kind):
                ports = []
            hits = [
                (specs[port], value)
                for port in ports
                if (value := _value(call, specs[port].dtype)) is not None
            ]
            record = {
                "message": "",
                "value": None,
                "applied": False,
                "delay_s": call.delay_s,
                "call": _text(call),
                "confidence": confidence,
            }
            if not hits:
                actions.append(record)
            for spec, value in hits:
                applied = write
                if applied and call.delay_s:
                    applied = self._pipeline is not None and self._pipeline.defer(
                        self, self._to_messages({spec.port: value}) or {}, call.delay_s
                    )
                elif applied:
                    writes[spec.port] = value
                actions.append(
                    record | {"message": spec.name, "value": value, "applied": applied}
                )

        if any(spec.port == ACTIONS for spec in self.output_ports):
            writes[ACTIONS] = actions
        return writes or None