led-controller/src/controllers/device.py

from microdot import Microdot
from models.device import (
    Device,
    derive_device_mac,
    normalize_mac,
    validate_device_transport,
    validate_device_type,
)
from models.group import Group
from models.transport import get_current_sender
from settings import Settings
from util.brightness_combine import effective_brightness_for_mac
from models.wifi_ws_clients import (
    normalize_tcp_peer_ip,
    send_json_line_to_ip,
    tcp_client_connected,
)
from util.driver_patterns import driver_patterns_dir
from util.espnow_message import build_message
import asyncio
import json
import os
import socket
from urllib.parse import quote

# Ephemeral driver preset name (never written to Pi preset store; ``save`` not set on wire).
_IDENTIFY_PRESET_KEY = "__identify"

# Short-key payload: 10 Hz full cycle = 50 ms on + 50 ms off (driver ``blink`` toggles each ``d`` ms).
_IDENTIFY_DRIVER_PRESET = {
    "p": "blink",
    "c": ["#ff0000"],
    "d": 50,
    "b": 128,
    "a": True,
    "n1": 0,
    "n2": 0,
    "n3": 0,
    "n4": 0,
    "n5": 0,
    "n6": 0,
}


def _compact_v1_json(*, presets=None, select=None, save=False):
    """Single-line v1 object; compact so serial/ESP-NOW stays small."""
    body = {"v": "1"}
    if presets is not None:
        body["presets"] = presets
        if save:
            body["save"] = True
    if select is not None:
        body["select"] = select
    return json.dumps(body, separators=(",", ":"))

# Seconds after identify blink before selecting built-in ``off`` (tests may monkeypatch).
IDENTIFY_OFF_DELAY_S = 2.0


def _validate_output_brightness(value):
    if value is None:
        return None
    try:
        b = int(value)
    except (TypeError, ValueError):
        raise ValueError("output_brightness must be an integer 0–255")
    if b < 0 or b > 255:
        raise ValueError("output_brightness must be between 0 and 255")
    return b


def _brightness_save_message_json(b_val: int) -> str:
    b_val = max(0, min(255, int(b_val)))
    return json.dumps({"v": "1", "b": b_val, "save": True}, separators=(",", ":"))


controller = Microdot()
devices = Device()
_group_registry = Group()
_pi_settings = Settings()


def _device_live_connected(dev_dict):
    """
    Wi-Fi: whether the controller has an outbound WebSocket to this device's IP.
    ESP-NOW: None (no Wi-Fi session on the Pi for that transport).
    """
    tr = (dev_dict.get("transport") or "espnow").strip().lower()
    if tr != "wifi":
        return None
    ip = normalize_tcp_peer_ip(dev_dict.get("address") or "")
    if not ip:
        return False
    return tcp_client_connected(ip)


def _device_json_with_live_status(dev_dict):
    row = dict(dev_dict)
    row["connected"] = _device_live_connected(dev_dict)
    return row


def _safe_pattern_filename(name):
    if not isinstance(name, str):
        return False
    if not name.endswith(".py"):
        return False
    if "/" in name or "\\" in name or ".." in name:
        return False
    return True


def _http_post_pattern_source(ip, filename, code_text, reload_patterns=True, timeout_s=10.0):
    """POST source to driver /patterns/upload?name=...&reload=...; return True on 2xx."""
    if not isinstance(ip, str) or not ip.strip():
        return False
    if not isinstance(filename, str) or not filename:
        return False
    if not isinstance(code_text, str):
        return False

    name_q = quote(filename, safe="")
    reload_q = "1" if reload_patterns else "0"
    path = "/patterns/upload?name=%s&reload=%s" % (name_q, reload_q)
    body = code_text.encode("utf-8")
    req = (
        "POST %s HTTP/1.1\r\n"
        "Host: %s\r\n"
        "Content-Type: text/plain; charset=utf-8\r\n"
        "Content-Length: %d\r\n"
        "Connection: close\r\n"
        "\r\n" % (path, ip, len(body))
    ).encode("utf-8") + body

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        sock.settimeout(timeout_s)
        sock.connect((ip.strip(), 80))
        sock.sendall(req)
        data = b""
        while True:
            chunk = sock.recv(1024)
            if not chunk:
                break
            data += chunk
    except OSError:
        return False
    finally:
        try:
            sock.close()
        except Exception:
            pass

    first_line = data.split(b"\r\n", 1)[0] if data else b""
    return b" 2" in first_line


async def _identify_send_off_after_delay(sender, transport, wifi_ip, dev_id, name):
    try:
        await asyncio.sleep(IDENTIFY_OFF_DELAY_S)
        off_msg = build_message(select={name: ["off"]})
        if transport == "wifi":
            await send_json_line_to_ip(wifi_ip, off_msg)
        else:
            await sender.send(off_msg, addr=dev_id)
    except Exception:
        pass


async def send_identify_to_device(dev_id: str) -> tuple[int, str]:
    """
    Send the same identify blink as ``POST /devices/<id>/identify``.

    Returns ``(http_status, "")`` on success, or ``(status, error_message)`` on failure
    (status matches the single-device route).
    """
    dev = devices.read(dev_id)
    if not dev:
        return 404, "Device not found"
    sender = get_current_sender()
    if not sender:
        return 503, "Transport not configured"
    name = str(dev.get("name") or "").strip()
    if not name:
        return 400, "Device must have a name to identify"

    transport = dev.get("transport") or "espnow"
    wifi_ip = None
    if transport == "wifi":
        wifi_ip = dev.get("address")
        if not wifi_ip:
            return 400, "Device has no IP address"

    try:
        msg = _compact_v1_json(
            presets={_IDENTIFY_PRESET_KEY: dict(_IDENTIFY_DRIVER_PRESET)},
            select={name: [_IDENTIFY_PRESET_KEY]},
        )
        if transport == "wifi":
            ok = await send_json_line_to_ip(wifi_ip, msg)
            if not ok:
                return 503, "Wi-Fi driver not connected"
        else:
            await sender.send(msg, addr=dev_id)

        asyncio.create_task(
            _identify_send_off_after_delay(sender, transport, wifi_ip, dev_id, name)
        )
    except Exception as e:
        return 503, str(e)
    return 200, ""


async def send_identify_to_group_devices(macs: list[str]) -> tuple[int, list[dict]]:
    """
    Identify every listed registry MAC in one delivery round: merged ``select`` and a single
    ESP-NOW split envelope when multiple peers share the serial bridge (avoids per-device
    ``SerialSender`` lock serialisation). Wi-Fi peers are sent in parallel as in
    ``deliver_json_messages``.
    """
    from util.driver_delivery import deliver_json_messages

    errors: list[dict] = []
    sender = get_current_sender()
    if not sender:
        return 0, [{"mac": "*", "error": "Transport not configured"}]

    merged_select: dict[str, list[str]] = {}
    valid_macs: list[str] = []
    for dev_id in macs:
        dev = devices.read(dev_id)
        if not dev:
            errors.append({"mac": dev_id, "error": "Device not found"})
            continue
        name = str(dev.get("name") or "").strip()
        if not name:
            errors.append({"mac": dev_id, "error": "Device must have a name to identify"})
            continue
        transport = (dev.get("transport") or "espnow").strip().lower()
        if transport == "wifi":
            if not dev.get("address"):
                errors.append({"mac": dev_id, "error": "Device has no IP address"})
                continue
        merged_select[name] = [_IDENTIFY_PRESET_KEY]
        valid_macs.append(dev_id)

    if not merged_select:
        return 0, errors

    try:
        msg = _compact_v1_json(
            presets={_IDENTIFY_PRESET_KEY: dict(_IDENTIFY_DRIVER_PRESET)},
            select=merged_select,
        )
        await deliver_json_messages(sender, [msg], valid_macs, devices, delay_s=0)
    except Exception as e:
        return 0, errors + [{"mac": "*", "error": str(e)}]

    for dev_id in valid_macs:
        dev = devices.read(dev_id) or {}
        name = str(dev.get("name") or "").strip()
        transport = (dev.get("transport") or "espnow").strip().lower()
        wifi_ip = dev.get("address") if transport == "wifi" else None
        asyncio.create_task(
            _identify_send_off_after_delay(sender, transport, wifi_ip, dev_id, name)
        )

    return len(valid_macs), errors


@controller.get("")
async def list_devices(request):
    """List all devices (includes ``connected`` for live Wi-Fi WebSocket presence)."""
    devices_data = {}
    for dev_id in devices.list():
        d = devices.read(dev_id)
        if d:
            devices_data[dev_id] = _device_json_with_live_status(d)
    return json.dumps(devices_data), 200, {"Content-Type": "application/json"}


@controller.post("/resolve-brightness")
async def resolve_brightness_batch(request):
    """
    POST JSON ``{ \"macs\": [\"..\"], \"zone_brightness\": optional 0–255 }``.
    Returns ``{ \"values\": { mac: combined_int } }`` — global × group(s) × device × zone (optional).
    """
    try:
        data = request.json or {}
    except Exception:
        data = {}
    macs = data.get("macs")
    if not isinstance(macs, list):
        return json.dumps({"error": "macs must be an array"}), 400, {
            "Content-Type": "application/json",
        }
    zb = None
    if isinstance(data, dict) and data.get("zone_brightness") is not None:
        try:
            zb = _validate_output_brightness(data.get("zone_brightness"))
        except ValueError as e:
            return json.dumps({"error": str(e)}), 400, {"Content-Type": "application/json"}
    values = {}
    for raw in macs:
        m = normalize_mac(str(raw))
        if not m:
            continue
        values[m] = effective_brightness_for_mac(
            _pi_settings,
            _group_registry,
            devices,
            m,
            zone_brightness=zb,
        )
    return json.dumps({"values": values}), 200, {"Content-Type": "application/json"}


@controller.get("/<id>")
async def get_device(request, id):
    """Get a device by ID (includes ``connected`` for live Wi-Fi WebSocket presence)."""
    dev = devices.read(id)
    if dev:
        return json.dumps(_device_json_with_live_status(dev)), 200, {
            "Content-Type": "application/json",
        }
    return json.dumps({"error": "Device not found"}), 404, {
        "Content-Type": "application/json",
    }


@controller.post("")
async def create_device(request):
    """Create a new device."""
    try:
        data = request.json or {}
        name = data.get("name", "").strip()
        if not name:
            return json.dumps({"error": "name is required"}), 400, {
                "Content-Type": "application/json",
            }
        try:
            device_type = validate_device_type(data.get("type", "led"))
            transport = validate_device_transport(data.get("transport", "espnow"))
        except ValueError as e:
            return json.dumps({"error": str(e)}), 400, {
                "Content-Type": "application/json",
            }
        address = data.get("address")
        mac = data.get("mac")
        if derive_device_mac(mac=mac, address=address, transport=transport) is None:
            return json.dumps(
                {
                    "error": "mac is required (12 hex digits); for Wi-Fi include mac plus IP in address"
                }
            ), 400, {"Content-Type": "application/json"}
        default_pattern = data.get("default_pattern")
        zl = data.get("zones")
        if isinstance(zl, list):
            zl = [str(t) for t in zl]
        else:
            zl = []
        dev_id = devices.create(
            name=name,
            address=address,
            mac=mac,
            default_pattern=default_pattern,
            zones=zl,
            device_type=device_type,
            transport=transport,
        )
        dev = devices.read(dev_id)
        return json.dumps({dev_id: dev}), 201, {"Content-Type": "application/json"}
    except ValueError as e:
        msg = str(e)
        code = 409 if "already exists" in msg.lower() else 400
        return json.dumps({"error": msg}), code, {"Content-Type": "application/json"}
    except Exception as e:
        return json.dumps({"error": str(e)}), 400, {"Content-Type": "application/json"}


@controller.put("/<id>")
async def update_device(request, id):
    """Update a device."""
    try:
        raw = request.json or {}
        data = dict(raw)
        data.pop("id", None)
        data.pop("addresses", None)
        data.pop("connected", None)
        if "name" in data:
            n = (data.get("name") or "").strip()
            if not n:
                return json.dumps({"error": "name cannot be empty"}), 400, {
                    "Content-Type": "application/json",
                }
            data["name"] = n
        if "type" in data:
            data["type"] = validate_device_type(data.get("type"))
        if "transport" in data:
            data["transport"] = validate_device_transport(data.get("transport"))
        if "zones" in data and isinstance(data["zones"], list):
            data["zones"] = [str(t) for t in data["zones"]]
        if "output_brightness" in data:
            data["output_brightness"] = _validate_output_brightness(data.get("output_brightness"))
        prev_doc = devices.read(id)
        if devices.update(id, data):
            if prev_doc and "name" in data:
                on = str(prev_doc.get("name") or "").strip()
                nn = str(data.get("name") or "").strip()
                if on and nn and on != nn:
                    from util.beat_driver_route import remap_beat_route_device_name

                    remap_beat_route_device_name(on, nn)
            return json.dumps(devices.read(id)), 200, {"Content-Type": "application/json"}
        return json.dumps({"error": "Device not found"}), 404, {
            "Content-Type": "application/json",
        }
    except ValueError as e:
        return json.dumps({"error": str(e)}), 400, {"Content-Type": "application/json"}
    except Exception as e:
        return json.dumps({"error": str(e)}), 400, {"Content-Type": "application/json"}


@controller.delete("/<id>")
async def delete_device(request, id):
    """Delete a device."""
    if devices.delete(id):
        return (
            json.dumps({"message": "Device deleted successfully"}),
            200,
            {"Content-Type": "application/json"},
        )
    return json.dumps({"error": "Device not found"}), 404, {
        "Content-Type": "application/json",
    }


@controller.post("/<id>/identify")
async def identify_device(request, id):
    """
    One v1 JSON object: ``presets.__identify`` (``d``=50 ms → 10 Hz blink) plus ``select`` for
    this device name — same combined shape as profile sends the driver already accepts over TCP
    / ESP-NOW. No ``save``. After ``IDENTIFY_OFF_DELAY_S``, a background task selects ``off``.
    """
    status, err = await send_identify_to_device(id)
    if status == 200:
        return json.dumps({"message": "Identify sent"}), 200, {
            "Content-Type": "application/json",
        }
    return json.dumps({"error": err}), status, {"Content-Type": "application/json"}


@controller.post("/<id>/brightness")
async def push_device_output_brightness(request, id):
    """
    Push combined brightness to the driver: global × group(s) × device × optional ``zone_brightness``
    in JSON body — single ``b`` (``v``/``b``/``save``). Wi‑Fi or ESP‑NOW.
    """
    dev = devices.read(id)
    if not dev:
        return json.dumps({"error": "Device not found"}), 404, {
            "Content-Type": "application/json",
        }
    body = request.json or {}
    zb = None
    if isinstance(body, dict) and body.get("zone_brightness") is not None:
        try:
            zb = _validate_output_brightness(body.get("zone_brightness"))
        except ValueError as e:
            return json.dumps({"error": str(e)}), 400, {"Content-Type": "application/json"}
    b_val = effective_brightness_for_mac(
        _pi_settings,
        _group_registry,
        devices,
        id,
        zone_brightness=zb,
    )

    msg = _brightness_save_message_json(b_val)
    transport = (dev.get("transport") or "espnow").strip().lower()

    if transport == "wifi":
        ip = normalize_tcp_peer_ip(str(dev.get("address") or ""))
        if not ip:
            return json.dumps({"error": "Device has no IP address"}), 400, {
                "Content-Type": "application/json",
            }
        ok = await send_json_line_to_ip(ip, msg)
        if not ok:
            return json.dumps({"error": "Wi-Fi driver not connected"}), 503, {
                "Content-Type": "application/json",
            }
    else:
        sender = get_current_sender()
        if not sender:
            return json.dumps({"error": "Transport not configured"}), 503, {
                "Content-Type": "application/json",
            }
        try:
            await sender.send(msg, addr=id)
        except Exception as e:
            return json.dumps({"error": str(e)}), 503, {"Content-Type": "application/json"}

    return json.dumps({"message": "brightness sent", "brightness": b_val}), 200, {
        "Content-Type": "application/json",
    }


@controller.post("/<id>/driver-config")
async def push_driver_config(request, id):
    """
    Push ``device_config`` to a Wi‑Fi LED driver over WebSocket.
    Body JSON: optional ``name``, ``num_leds``, ``color_order``, ``startup_mode`` (default|last|off).
    """
    dev = devices.read(id)
    if not dev:
        return json.dumps({"error": "Device not found"}), 404, {
            "Content-Type": "application/json",
        }
    if (dev.get("transport") or "").lower() != "wifi":
        return json.dumps({"error": "driver-config is only for Wi-Fi devices"}), 400, {
            "Content-Type": "application/json",
        }
    wifi_ip = str(dev.get("address") or "").strip()
    if not wifi_ip:
        return json.dumps({"error": "Device has no IP address"}), 400, {
            "Content-Type": "application/json",
        }
    body = request.json or {}
    dc = {}
    if isinstance(body.get("name"), str) and body["name"].strip():
        dc["name"] = body["name"].strip()
    if "num_leds" in body:
        try:
            n = int(body["num_leds"])
            if 1 <= n <= 2048:
                dc["num_leds"] = n
        except (TypeError, ValueError):
            pass
    if isinstance(body.get("color_order"), str):
        co = body["color_order"].strip().lower()
        if co in ("rgb", "rbg", "grb", "gbr", "brg", "bgr"):
            dc["color_order"] = co
    if isinstance(body.get("startup_mode"), str):
        sm = body["startup_mode"].strip().lower()
        if sm in ("default", "last", "off"):
            dc["startup_mode"] = sm
    if not dc:
        return json.dumps(
            {
                "error": "Provide at least one of name, num_leds, color_order, startup_mode"
            }
        ), 400, {"Content-Type": "application/json"}
    msg = json.dumps(
        {"v": "1", "device_config": dc, "save": True}, separators=(",", ":")
    )
    ok = await send_json_line_to_ip(wifi_ip, msg)
    if not ok:
        return json.dumps({"error": "Wi-Fi driver not connected"}), 503, {
            "Content-Type": "application/json",
        }
    return json.dumps({"message": "driver-config sent"}), 200, {
        "Content-Type": "application/json",
    }


@controller.post("/<id>/patterns/push")
async def push_patterns_ota(request, id):
    """
    Push all local pattern files directly to a Wi-Fi LED driver over HTTP upload.
    """
    dev = devices.read(id)
    if not dev:
        return json.dumps({"error": "Device not found"}), 404, {
            "Content-Type": "application/json",
        }
    if (dev.get("transport") or "").lower() != "wifi":
        return json.dumps({"error": "Pattern OTA push is only supported for Wi-Fi devices"}), 400, {
            "Content-Type": "application/json",
        }
    wifi_ip = str(dev.get("address") or "").strip()
    if not wifi_ip:
        return json.dumps({"error": "Device has no IP address"}), 400, {
            "Content-Type": "application/json",
        }

    base_dir = driver_patterns_dir()
    try:
        names = sorted(os.listdir(base_dir))
    except OSError as e:
        return json.dumps({"error": str(e)}), 500, {
            "Content-Type": "application/json",
        }

    files = [n for n in names if _safe_pattern_filename(n) and n != "__init__.py"]
    if not files:
        return json.dumps({"error": "No pattern files found"}), 404, {
            "Content-Type": "application/json",
        }

    sent = []
    failed = []
    total = len(files)
    for idx, filename in enumerate(files):
        path = os.path.join(base_dir, filename)
        try:
            with open(path, "r") as f:
                code = f.read()
        except OSError:
            failed.append(filename)
            continue
        reload_patterns = idx == (total - 1)
        ok = _http_post_pattern_source(
            wifi_ip,
            filename,
            code,
            reload_patterns=reload_patterns,
            timeout_s=10.0,
        )
        if ok:
            sent.append(filename)
        else:
            failed.append(filename)

    if not sent:
        return json.dumps({"error": "Wi-Fi driver did not accept pattern uploads", "failed": failed}), 503, {
            "Content-Type": "application/json",
        }

    return json.dumps({
        "message": "Pattern files uploaded",
        "sent_count": len(sent),
        "sent": sent,
        "failed": failed,
    }), 200, {
        "Content-Type": "application/json",
    }