"""Deliver driver JSON messages over serial (ESP-NOW) and/or WebSocket (Wi-Fi drivers)."""

import asyncio
import json

from models.device import normalize_mac
from models.wifi_ws_clients import send_json_line_to_ip

# Serial bridge (ESP32): broadcast MAC + this envelope → firmware unicasts ``body`` to each peer.
_SPLIT_MODE = "split"
_BROADCAST_MAC_HEX = "ffffffffffff"


def _split_serial_envelope(inner_json_str, peer_hex_list):
    """One UART frame: broadcast dest + JSON {m:split, peers:[hex,...], body:<object>}."""
    body = json.loads(inner_json_str)
    env = {"m": _SPLIT_MODE, "peers": list(peer_hex_list), "body": body}
    return json.dumps(env, separators=(",", ":"))


def _wifi_message_for_device(msg, device_name):
    """
    For Wi-Fi WebSocket fanout, narrow a v1 select map to a single device name.
    Returns the original message when no narrowing applies.
    """
    if not device_name:
        return msg
    try:
        body = json.loads(msg)
    except Exception:
        return msg
    if not isinstance(body, dict):
        return msg
    select = body.get("select")
    if not isinstance(select, dict):
        return msg
    if device_name not in select:
        return msg
    body["select"] = {device_name: select[device_name]}
    return json.dumps(body, separators=(",", ":"))


def _combine_preset_chunks_for_wifi(chunk_messages):
    """Merge chunked v1 preset messages into one v1 JSON string for Wi-Fi."""
    merged_presets = {}
    save_flag = False
    default_id = None
    for msg in chunk_messages:
        try:
            body = json.loads(msg)
        except Exception:
            continue
        if not isinstance(body, dict):
            continue
        presets = body.get("presets")
        if isinstance(presets, dict):
            merged_presets.update(presets)
        if body.get("save"):
            save_flag = True
        if body.get("default") is not None:
            default_id = body.get("default")
    out = {"v": "1", "presets": merged_presets}
    if save_flag:
        out["save"] = True
    if default_id is not None:
        out["default"] = default_id
    return json.dumps(out, separators=(",", ":"))


async def deliver_preset_broadcast_then_per_device(
    sender,
    chunk_messages,
    target_macs,
    devices_model,
    default_id,
    delay_s=0.1,
):
    """
    Send preset definition chunks: ESP-NOW broadcast once per chunk; same chunk to each
    Wi-Fi driver over WebSocket. If default_id is set, send a per-target default message
    (unicast serial or WebSocket) with targets=[device name] for each registry entry.
    """
    if not chunk_messages:
        return 0

    seen = set()
    ordered = []
    for raw in target_macs:
        m = normalize_mac(str(raw)) if raw else None
        if not m or m in seen:
            continue
        seen.add(m)
        ordered.append(m)

    wifi_ips = []
    for mac in ordered:
        doc = devices_model.read(mac)
        if doc and doc.get("transport") == "wifi" and doc.get("address"):
            wifi_ips.append(str(doc["address"]).strip())

    deliveries = 0
    wifi_combined_msg = _combine_preset_chunks_for_wifi(chunk_messages)
    for msg in chunk_messages:
        tasks = [sender.send(msg, addr=_BROADCAST_MAC_HEX)]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        if results and results[0] is True:
            deliveries += 1
        await asyncio.sleep(delay_s)

    for ip in wifi_ips:
        if not ip:
            continue
        try:
            if await send_json_line_to_ip(ip, wifi_combined_msg):
                deliveries += 1
        except Exception as e:
            print(f"[driver_delivery] Wi-Fi preset send failed: {e!r}")
        await asyncio.sleep(delay_s)

    if default_id:
        did = str(default_id)
        for mac in ordered:
            doc = devices_model.read(mac) or {}
            name = str(doc.get("name") or "").strip() or mac
            body = {"v": "1", "default": did, "save": True, "targets": [name]}
            out = json.dumps(body, separators=(",", ":"))
            if doc.get("transport") == "wifi" and doc.get("address"):
                ip = str(doc["address"]).strip()
                try:
                    if await send_json_line_to_ip(ip, out):
                        deliveries += 1
                except Exception as e:
                    print(f"[driver_delivery] default Wi-Fi send failed: {e!r}")
            else:
                try:
                    await sender.send(out, addr=mac)
                    deliveries += 1
                except Exception as e:
                    print(f"[driver_delivery] default serial failed: {e!r}")
            await asyncio.sleep(delay_s)

    return deliveries


async def deliver_json_messages(sender, messages, target_macs, devices_model, delay_s=0.1):
    """
    Send each message string to the bridge and/or Wi-Fi WebSocket clients.

    If target_macs is None or empty: one serial send per message (default/broadcast address).
    Otherwise: Wi-Fi uses WebSocket in parallel. Multiple ESP-NOW peers are sent in **one** serial
    write to the ESP32 (broadcast + split envelope); the bridge unicasts ``body`` to each
    peer. A single ESP-NOW peer still uses one unicast serial frame. Wi-Fi and serial
    tasks run together in one asyncio.gather.

    Returns (delivery_count, chunk_count) where chunk_count is len(messages).
    """
    if not messages:
        return 0, 0

    if not target_macs:
        deliveries = 0
        for msg in messages:
            await sender.send(msg)
            deliveries += 1
            await asyncio.sleep(delay_s)
        return deliveries, len(messages)

    seen = set()
    ordered_macs = []
    for raw in target_macs:
        m = normalize_mac(str(raw)) if raw else None
        if not m or m in seen:
            continue
        seen.add(m)
        ordered_macs.append(m)

    deliveries = 0
    for msg in messages:
        wifi_tasks = []
        espnow_hex = []
        for mac in ordered_macs:
            doc = devices_model.read(mac)
            if doc and doc.get("transport") == "wifi":
                ip = doc.get("address")
                if ip:
                    name = str(doc.get("name") or "").strip()
                    wifi_msg = _wifi_message_for_device(msg, name)
                    wifi_tasks.append(send_json_line_to_ip(ip, wifi_msg))
            else:
                espnow_hex.append(mac)

        tasks = []
        espnow_peer_count = 0
        if len(espnow_hex) > 1:
            tasks.append(
                sender.send(
                    _split_serial_envelope(msg, espnow_hex),
                    addr=_BROADCAST_MAC_HEX,
                )
            )
            espnow_peer_count = len(espnow_hex)
        elif len(espnow_hex) == 1:
            tasks.append(sender.send(msg, addr=espnow_hex[0]))
            espnow_peer_count = 1

        tasks.extend(wifi_tasks)

        if tasks:
            results = await asyncio.gather(*tasks, return_exceptions=True)
            n_serial = len(tasks) - len(wifi_tasks)
            for i, r in enumerate(results):
                if i < n_serial:
                    if r is True:
                        deliveries += espnow_peer_count
                    elif isinstance(r, Exception):
                        print(f"[driver_delivery] serial delivery failed: {r!r}")
                else:
                    if r is True:
                        deliveries += 1
                    elif isinstance(r, Exception):
                        print(f"[driver_delivery] Wi-Fi delivery failed: {r!r}")

        await asyncio.sleep(delay_s)
    return deliveries, len(messages)