"""Live household power over the local socket.

The app's home page polls ``queryDeviceRealTimeData`` (socket
``DEVICE_DETAIL_INFO``) for its energy-flow screen; per-switch watts come
from ``queryLoadStatusInfo`` (socket ``LOAD_STATUS``); grid/load
current+voltage telemetry comes from ``queryEmsInfoBySocket``.  See
``docs/RE_DART_FINDINGS.md`` §4 (ids) and §16.
"""

from __future__ import annotations

import argparse
import json
import logging
import socket
from typing import Any, Dict, List, Optional, Tuple

from .battery import _walk
from .device_local import LocalDeviceClient

log = logging.getLogger(__name__)

HOST = "10.100.1.1"
PORTS: Tuple[int, ...] = (9000, 18000)

#: Reply keys holding watt readings, in priority order.
POWER_KEYS: Tuple[str, ...] = (
    "totalPower",
    "homeloadPower",
    "gridPower",
    "solarPower",
    "SW1ExpPower",
    "SW2ExpPower",
)


def _to_float(value: Any) -> Optional[float]:
    try:
        if isinstance(value, bool):
            return None
        return float(value)
    except (TypeError, ValueError):
        return None


def extract_power(reply: Dict[str, Any]) -> Dict[str, Any]:
    """Pull watt readings out of a parsed reply envelope.

    Returns ``{"value": float|None, "source": str|None, "detail": {...}}``.
    ``value`` is the first ``POWER_KEYS`` hit (headline: ``totalPower``);
    ``detail`` carries the supporting fields (per-switch watts/currents,
    grid volts/amps) when present.
    """
    data = reply.get("dataArea", reply)
    if isinstance(data, str):
        try:
            data = json.loads(data)
        except (ValueError, TypeError):
            data = {}
    for key in POWER_KEYS:
        for hit in _walk(reply, key):
            num = _to_float(hit)
            if num is not None:
                return {"value": num, "source": key,
                        "detail": _power_detail(data)}
    return {"value": None, "source": None, "detail": _power_detail(data)}


def _power_detail(data: Any) -> Dict[str, Any]:
    detail: Dict[str, Any] = {}
    if not isinstance(data, dict):
        return detail
    for key in (
        "totalPower", "apowerCount", "homeloadPower", "solarPower",
        "gridPower", "SW1ExpPower", "SW2ExpPower", "Sw1Volt", "Sw2Volt",
        "SW1Curr", "gridVol1", "gridVol2", "gridCurr1", "gridCurr2",
        "loadCurr1", "loadCurr2", "gridLineVol", "gridFreq",
    ):
        hits = _walk(data, key)
        if hits:
            detail[key] = hits[0] if len(hits) == 1 else hits
    return detail


def ems_load_watts(reply: Dict[str, Any]) -> Optional[float]:
    """Estimate household draw from an EMS reply, when present.

    ``(loadCurr1 + loadCurr2) * (gridVol1 + gridVol2) / 2`` — per-leg
    current times mean leg voltage.  Returns ``None`` if the fields are
    absent (the EMS reply on the tested unit carries live grid Vol/Curr).
    """
    data = reply.get("dataArea", reply)
    if not isinstance(data, dict):
        return None
    try:
        amps = float(data["loadCurr1"]) + float(data["loadCurr2"])
        volts = (float(data["gridVol1"]) + float(data["gridVol2"])) / 2
    except (KeyError, TypeError, ValueError):
        return None
    return amps * volts


def get_power(
    sn: str,
    host: str = HOST,
    ports: Tuple[int, ...] = PORTS,
    key_radix: int = 16,
    timeout: float = 5.0,
) -> Dict[str, Any]:
    """Read live household wattage over the local TCP socket.

    Tries realtime detail (``DEVICE_DETAIL_INFO`` → ``totalPower``),
    then per-switch load status, then EMS telemetry, probing each port.
    Returns the ``extract_power`` dict plus ``cmd``/``port``/``reply``.
    """
    last_err: Optional[Exception] = None
    for port in ports:
        client = LocalDeviceClient(
            sn=sn, host=host, socket_port=port,
            key_radix=key_radix, timeout=timeout,
        )
        for name, call in (
            ("detail", client.query_detail_info),
            ("load", client.query_load_status),
            ("ems", client.query_ems_info),
        ):
            try:
                reply = call()
            except (ConnectionError, socket.error, TimeoutError) as exc:
                last_err = exc
                log.debug("port %d %s failed: %s", port, name, exc)
                break
            except Exception as exc:
                last_err = exc
                log.debug("port %d %s error: %s", port, name, exc)
                continue
            result = extract_power(reply)
            if result["value"] is None and name == "ems":
                watts = ems_load_watts(reply)
                if watts is not None:
                    result = {"value": watts, "source": "ems(loadCurr*V)",
                              "detail": _power_detail(
                                  reply.get("dataArea", reply))}
            result.update({"cmd": name, "port": port, "reply": reply})
            if result["value"] is not None:
                return result
            last_err = ValueError(
                f"reply on port {port} {name} has no power field")
    raise ConnectionError(f"no power reply via {host}:{list(ports)}") from last_err


def main(argv: Optional[List[str]] = None) -> int:
    parser = argparse.ArgumentParser(
        description="Read live household wattage over the device SoftAP socket.")
    parser.add_argument("--sn", required=True, help="20-char device serial")
    parser.add_argument("--host", default=HOST, help="aGate gateway IP")
    parser.add_argument("--ports", default=",".join(map(str, PORTS)),
                        help="comma-separated socket ports to probe")
    parser.add_argument("--timeout", type=float, default=5.0)
    parser.add_argument("--raw", action="store_true",
                        help="print full reply JSON")
    args = parser.parse_args(argv)

    ports = tuple(int(p) for p in args.ports.split(",") if p.strip())
    try:
        result = get_power(args.sn, host=args.host, ports=ports,
                           timeout=args.timeout)
    except Exception as exc:
        print(f"ERROR: {exc}")
        return 1
    if args.raw:
        print(json.dumps(result["reply"], indent=2, default=str))
    else:
        value = result["value"]
        print(f"{value:g}W ({result['source']})"
              if value is not None else "unknown")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())