#!/usr/bin/env python3
"""
Connect to EcoFlow River 2 and River 3 via Bluetooth LE and print a telemetry
snapshot.

Supply your device details on the command line, or use --scan to discover them.
After authentication, devices push telemetry packets. River 2 collection waits
for its expected heartbeat types. River 3 prints updates as they arrive and tracks
required measurement and configuration fields across partial reports. Normal
collection finishes when the required information arrives; --timeout is only a
failure limit. Use --watch to print every packet from either model until Ctrl+C,
including raw payloads for investigation. Watch output can contain private data.

Usage:
  uv run --extra ecoflow-bt python tools/ecoflow-bt.py --user-id <id> --scan
  uv run --extra ecoflow-bt python tools/ecoflow-bt.py --email you@example.com --scan
  uv run --extra ecoflow-bt python tools/ecoflow-bt.py --user-id <id> --device <address> --serial-number <serial>
"""

import argparse
import asyncio
import getpass
import sys

try:
    from bmspy.ecoflow_bt import EcoFlowBTManager, EcoFlowBTDevice, EcoFlowPacket, scan_for_devices
    from bmspy.ecoflow_login import get_user_id
    from bmspy.ecoflow_proto import (
        format_payload, format_protobuf, describe_protobuf_field, parse_protobuf_fields,
    )
    from bmspy.ecoflow_river2 import decode_river2_payload
except ImportError as exc:
    sys.exit(f"Missing dependency: {exc}\n"
             "Run: uv sync --extra ecoflow-bt")

# ── Model profiles (no device-specific identifiers) ──────────────────────────

# Baseline from the named fields in observed R651 measurement/configuration
# reports, not every optional field in the protocol schema. Other firmware can
# omit fields: report those as missing rather than claiming a complete report.
# Values are protobuf wire types (0: varint, 5: float32).
RIVER3_REQUIRED_GROUPS = {
    "measurements": {
        **dict.fromkeys((1, 148, 212, 248, 254, 255, 258, 259, 260, 261,
                         268, 269, 275, 281, 282, 355), 0),
        **dict.fromkeys((3, 4, 9, 10, 11, 12, 37, 54, 158, 242, 243,
                         262, 263, 361, 368), 5),
    },
    "configuration": dict.fromkeys(
        (7, 8, 17, 18, 19, 25, 33, 74, 195, 202, 209, 211, 270, 271,
         356, 358, 360, 364, 365, 366, 367, 458), 0,
    ),
}

RIVER2_REQUIRED_HEARTBEATS = {
    "ef_pkt_02_20_02": "PD",
    "ef_pkt_03_20_02": "EMS",
    "ef_pkt_03_20_32": "BMS",
}
RIVER2_OPTIONAL_HEARTBEATS = {
    "ef_pkt_04_20_02": "INV",
    "ef_pkt_05_20_02": "MPPT",
}

MODEL_PROFILES = {
    "R651": {
        "label":            "River 3",
        "enc_type":         7,
        "src_candidates":   [0x02],
        "required_groups":  RIVER3_REQUIRED_GROUPS,
        # DisplayPropertyUpload can arrive as several partial updates.
        # RuntimePropertyUpload (0x16) is not confirmed for River 3.
        "expected_fields": {
            "ef_pkt_02_fe_15", "bms_capacity_charge_ratio",
            "bms_temperature_celsius", "ef_input_power_watts",
            "ef_output_power_watts", "bms_charge_is_charging",
        },
    },
    "R601": {
        "label":            "River 2",
        "enc_type":         7,
        "src_candidates":   [0x02, 0x03, 0x04, 0x05],
        # PD: total power; EMS: time estimates; BMS: battery temperature.
        "required_heartbeats": RIVER2_REQUIRED_HEARTBEATS,
        "expected_fields": set(RIVER2_REQUIRED_HEARTBEATS),
    },
}

MODEL_PROFILES["R603"] = {**MODEL_PROFILES["R601"], "label": "River 2 Max"}
for prefix in ("R653", "R654", "R655"):
    MODEL_PROFILES[prefix] = {**MODEL_PROFILES["R651"]}

# ── Packet decoder ─────────────────────────────────────────────────────────────


def _decode_snapshot_packet(
    key: str,
    payload_hex: str,
    src_candidates: list[int],
) -> list[str]:
    parts = key.split("_")
    if len(parts) < 4:
        return [f"    {payload_hex}"]
    try:
        if len(parts) == 5:
            src_candidates = [int(parts[2], 16)]
            cmd_set = int(parts[3], 16)
            cmd_id = int(parts[4], 16)
        else:  # Snapshots from older bmspy versions lack the source address.
            cmd_set = int(parts[2], 16)
            cmd_id = int(parts[3], 16)
    except ValueError:
        return [f"    {payload_hex}"]

    header = f"[cmd_set=0x{cmd_set:02X} cmd_id=0x{cmd_id:02X}]"
    try:
        payload = bytes.fromhex(payload_hex)
    except (ValueError, TypeError):
        return [f"  {header}: invalid hex payload {payload_hex!r}"]
    if not payload:
        return [f"  {header}: (empty payload)"]

    for src in src_candidates:
        lines = format_payload(src, cmd_set, cmd_id, payload)
        if lines is not None:
            return ["  " + ln for ln in lines]

    # Fallback: generic protobuf formatting
    pb = format_protobuf(payload)
    if pb:
        return [f"  {header}"] + ["    " + ln for ln in pb]
    return [f"  {header}: {payload_hex[:120]}{'...' if len(payload_hex) > 120 else ''}"]


# ── Snapshot printer ───────────────────────────────────────────────────────────


def print_snapshot(
    label: str, address: str, snap, src_candidates: list[int],
    decoded_packets: dict[str, list[str]] | None = None,
) -> None:
    print(f"\n{'━' * 60}")
    print(f"  {label}  ({address})")
    print(f"{'━' * 60}")

    pkt_keys = []
    for name, field in snap.items():
        if name.startswith("ef_pkt_"):
            payload_hex = getattr(field, "info", None)
            pkt_keys.append((name, payload_hex if payload_hex is not None else str(field)))
        else:
            val = getattr(field, "value", None)
            if val is None:
                val = getattr(field, "values", None) or getattr(field, "info", str(field))
            unit = getattr(field, "units", "") or ""
            print(f"  {name}: {val}{' ' + unit if unit else ''}")

    for key, payload_hex in pkt_keys:
        print(f"\n  {key}:")
        lines = (decoded_packets or {}).get(key)
        if lines is None:
            lines = _decode_snapshot_packet(key, payload_hex, src_candidates)
        for line in lines:
            print(line)

    print()


# ── Live diagnostics and field-level collection ────────────────────────────────


class PacketCollector:
    """Track valid River 2 heartbeats or accumulated River 3 fields."""

    def __init__(self, address: str, profile: dict, watch: bool = False) -> None:
        self.address = address
        self.label = profile["label"]
        self.required_groups = profile.get("required_groups", {})
        self.required_heartbeats = profile.get("required_heartbeats", {})
        self.valid_heartbeats: set[str] = set()
        self.required_types = {
            number: wire_type for fields in self.required_groups.values()
            for number, wire_type in fields.items()
        }
        self.watch = watch
        self.seen: set[tuple[int, int]] = set()
        self.fields: dict[int, bytes] = {}
        self.error: Exception | None = None

    def observe_snapshot(self, snap) -> None:
        """Seed River 2 receipt tracking for callers that start after reception.

        Normal CLI runs register receive() before connecting. Snapshot seeding
        also supports direct collect() callers without packet listeners.
        """
        for key, field in snap.items():
            if key not in self.required_heartbeats or key in self.valid_heartbeats:
                continue
            try:
                payload = bytes.fromhex(field.info)
            except (AttributeError, TypeError, ValueError):
                continue
            self._accept_heartbeat(key, payload)

    def _accept_heartbeat(self, key: str, payload: bytes) -> bool:
        route = tuple(int(part, 16) for part in key.split("_")[2:])
        if decode_river2_payload(*route, payload) is None:
            return False
        self.valid_heartbeats.add(key)
        return True

    def missing(self) -> set[str]:
        if self.required_heartbeats:
            return {
                f"{label} ({key})" for key, label in self.required_heartbeats.items()
                if key not in self.valid_heartbeats
            }
        return {
            f"{group}:{describe_protobuf_field(2, 0xFE, 0x15, number) or number}"
            for group, fields in self.required_groups.items()
            for number, wire_type in fields.items()
            if (number, wire_type) not in self.seen
        }

    def progress(self) -> str:
        if self.required_heartbeats:
            return (f"{len(self.valid_heartbeats)}/{len(self.required_heartbeats)} "
                    "required heartbeats received")
        return ", ".join(
            f"{group} {sum((n, w) in self.seen for n, w in fields.items())}/{len(fields)}"
            for group, fields in self.required_groups.items()
        )

    def receive(self, packet: EcoFlowPacket) -> None:
        # Device callbacks cannot propagate failures through the BLE dispatcher.
        # The collection loop re-raises output errors and closes the connections.
        if self.error is not None:
            return
        try:
            self._receive(packet)
        except Exception as exc:
            self.error = exc

    def _receive(self, packet: EcoFlowPacket) -> None:
        route = (packet.src, packet.cmd_set, packet.cmd_id)
        key = f"ef_pkt_{packet.src:02x}_{packet.cmd_set:02x}_{packet.cmd_id:02x}"
        header = f"\n{self.label} ({self.address}) {key}:"
        if self.watch:
            lines = format_payload(*route, packet.payload)
            print("\n".join([
                header,
                f"  version=0x{packet.version:02x} dst=0x{packet.dst:02x} "
                f"dsrc=0x{packet.dsrc:02x} ddst=0x{packet.ddst:02x} seq={packet.seq.hex()}",
                f"  raw_payload: {packet.payload.hex() or '(empty)'}",
                *("  " + line for line in (lines or ["No decoded fields; retain raw payload for review."])),
            ]), flush=True)

        if key in self.required_heartbeats and not self._accept_heartbeat(key, packet.payload):
            print(f"{header}\n  {self.required_heartbeats[key]} heartbeat rejected by decoder. "
                  "Not counted toward completion", flush=True)
        if not self.required_groups or route != (0x02, 0xFE, 0x15):
            return
        try:
            parsed = parse_protobuf_fields(packet.payload)
        except ValueError as exc:
            print(f"{header}\n  Malformed report, not counted: {exc}", flush=True)
            return
        # Keep repeated/nested values together. A subsequent report replaces a
        # whole top-level field, rather than flattening repeated child fields.
        grouped: dict[int, bytes] = {}
        for field in parsed:
            expected_type = self.required_types.get(field.number)
            if expected_type is not None and field.wire_type != expected_type:
                print(f"{header}\n  Field {field.number}: unexpected wire type "
                      f"{field.wire_type}, not counted", flush=True)
                continue
            self.seen.add((field.number, field.wire_type))
            grouped[field.number] = grouped.get(field.number, b"") + field.encoded
        changed = {n: value for n, value in grouped.items() if self.fields.get(n) != value}
        self.fields.update(grouped)
        if changed and not self.watch:
            lines = format_payload(2, 0xFE, 0x15, b"".join(changed[n] for n in sorted(changed)))
            print("\n".join([header, *("  " + line for line in lines)]), flush=True)

    def decoded_packets(self) -> dict[str, list[str]]:
        if not self.fields:
            return {}
        payload = b"".join(self.fields[n] for n in sorted(self.fields))
        lines = format_payload(2, 0xFE, 0x15, payload)
        return {"ef_pkt_02_fe_15": [
            "  [DisplayPropertyUpload: latest fields collected this run]",
            *("  " + line for line in lines[1:]),
        ]}


def print_river2_packet_summary(snap) -> None:
    """Separate optional heartbeats and other packets from required heartbeats."""
    packets = {key: field for key, field in snap.items() if key.startswith("ef_pkt_")}
    seen = [label for key, label in RIVER2_OPTIONAL_HEARTBEATS.items() if key in packets]
    absent = [label for key, label in RIVER2_OPTIONAL_HEARTBEATS.items() if key not in packets]
    print(f"    Optional heartbeat types seen: {', '.join(seen) or 'none'}. "
          f"Not seen: {', '.join(absent) or 'none'} (not required)", flush=True)
    other = packets.keys() - RIVER2_REQUIRED_HEARTBEATS.keys() - RIVER2_OPTIONAL_HEARTBEATS.keys()
    if other:
        descriptions = [
            key + (" (empty payload)" if getattr(packets[key], "info", None) == "" else "")
            for key in sorted(other)
        ]
        print(f"    Other packet types seen: {', '.join(descriptions)}", flush=True)


# ── Collection logic ───────────────────────────────────────────────────────────


async def collect(
    devices: list[tuple[str, EcoFlowBTDevice, set[str]]],
    max_wait: float,
    collectors: dict[str, PacketCollector] | None = None,
    watch: bool = False,
) -> None:
    """Finish on required information, or keep watching until interrupted.

    The deadline is a failure limit, not evidence that a report is complete.
    """
    collectors = dict(collectors or {})
    for addr, dev, expected in devices:
        if addr not in collectors and expected and expected <= RIVER2_REQUIRED_HEARTBEATS.keys():
            collectors[addr] = PacketCollector(addr, {
                "label": "River 2",
                "required_heartbeats": {
                    key: label for key, label in RIVER2_REQUIRED_HEARTBEATS.items() if key in expected
                },
            })
    done: set[str] = set()
    last_missing: dict[str, set[str]] = {}
    deadline = asyncio.get_event_loop().time() + max_wait

    def missing_fields(addr, dev, expected):
        collector = collectors.get(addr)
        if collector is not None and (collector.required_groups or collector.required_heartbeats):
            if collector.required_heartbeats:
                collector.observe_snapshot(dev.snapshot)
            return collector.missing()
        return expected - {name for name, _ in dev.snapshot.items()}

    while True:
        for addr, dev, expected in devices:
            collector = collectors.get(addr)
            if collector is not None and collector.error is not None:
                raise collector.error
            if watch and not dev.is_authenticated:
                raise ConnectionError(f"{addr}: watch stopped because the device is not authenticated")
            if addr in done:
                continue
            received = {name for name, _ in dev.snapshot.items()}
            missing = missing_fields(addr, dev, expected)
            if missing:
                if last_missing.get(addr) != missing:
                    if collector is not None and collector.required_groups:
                        print(f"  {addr}: {collector.progress()}; missing {sorted(missing)}", flush=True)
                    elif not watch:
                        print(f"  {addr}: waiting for {missing}", flush=True)
                    last_missing[addr] = missing
            else:
                if collector is not None and collector.required_groups:
                    print(f"  {addr}: all required information received ({collector.progress()})", flush=True)
                elif collector is not None and collector.required_heartbeats:
                    labels = ", ".join(collector.required_heartbeats.values())
                    print(f"  {addr}: {collector.progress()} ({labels})", flush=True)
                    print_river2_packet_summary(dev.snapshot)
                else:
                    packet_count = sum(name.startswith("ef_pkt_") for name in received)
                    print(f"  {addr}: complete ({packet_count} packet type(s))", flush=True)
                if watch:
                    print(f"  {addr}: continuing watch for optional or unknown information", flush=True)
                done.add(addr)

        if len(done) == len(devices) and not watch:
            return
        remaining = deadline - asyncio.get_event_loop().time()
        if not watch and remaining <= 0:
            for addr, dev, expected in devices:
                if addr not in done:
                    missing = missing_fields(addr, dev, expected)
                    print(f"  {addr}: timed out, incomplete; missing {sorted(missing)}", flush=True)
            return
        await asyncio.sleep(0.5 if watch else min(0.5, remaining))


# ── Main ───────────────────────────────────────────────────────────────────────


async def run(user_id: str, args: argparse.Namespace) -> None:
    found = []
    if args.serial_number and not args.device:
        sys.exit("--serial-number requires --device.")
    if not args.device and not args.scan:
        sys.exit("Supply --device and --serial-number, or use --scan.")
    if args.scan:
        print("Scanning for EcoFlow devices (10 s)…", flush=True)
        found = await scan_for_devices(timeout=10.0, debug=args.debug)
        if found:
            for d in found:
                print(f"  {d.address}  serial={d.serial_number}  "
                      f"enc_type={d.encrypt_type}  rssi={d.rssi}")
        else:
            print("  (none found)")
        print()

    discovered = {d.address.upper(): d for d in found}
    addresses = [args.device.upper()] if args.device else list(discovered)
    if not addresses:
        sys.exit("No EcoFlow devices found.")
    selected = {}
    for address in addresses:
        detected = discovered.get(address)
        serial = args.serial_number or (detected.serial_number if detected else None)
        if not serial:
            sys.exit("Supply --serial-number, or use --scan to discover the selected device.")
        profile = MODEL_PROFILES.get(serial[:4])
        if profile is None:
            if args.device:
                sys.exit("Unsupported serial prefix. Supported prefixes: " + ", ".join(MODEL_PROFILES))
            print(f"Skipping unsupported model prefix: {serial[:4]}")
            continue
        encrypt_type = detected.encrypt_type if detected else profile["enc_type"]
        selected[address] = {
            **profile, "serial_number": serial, "enc_type": encrypt_type,
        }
    if not selected:
        sys.exit("No supported EcoFlow devices found.")

    manager = EcoFlowBTManager(user_id=user_id, debug=args.debug)
    devices     = []
    labels      = {}
    src_by_addr = {}
    collectors = {}
    watch = getattr(args, "watch", False)

    for address, d in selected.items():
        addr = address.upper()
        labels[addr]      = d["label"]
        src_by_addr[addr] = d["src_candidates"]
        dev = manager.add_device(
            address=address,
            serial_number=d["serial_number"],
            encrypt_type=d["enc_type"],
        )
        devices.append((addr, dev, d["expected_fields"]))
        if watch or d.get("required_groups") or d.get("required_heartbeats"):
            collector = PacketCollector(addr, d, watch=watch)
            collectors[addr] = collector
            # Register before connecting so the first notification is not lost.
            dev.on_packet(collector.receive)
        print(f"Registered {d['label']}  {address}  serial={d['serial_number']}")

    try:
        if watch:
            print("\nWatching both models until Ctrl+C; telemetry timeout is ignored.\n"
                  "Watch output can contain private identifiers. Redact it before sharing.", flush=True)
        print("\nConnecting…", flush=True)
        await manager.connect_all(auth_timeout=30.0)

        if not watch:
            print(f"\nCollecting required information (failure limit={args.timeout}s)…", flush=True)
        await collect(devices, max_wait=args.timeout, collectors=collectors, watch=watch)

        print("\nSnapshot:")
        for addr, dev, _ in devices:
            if addr in collectors and collectors[addr].required_groups:
                print_snapshot(labels[addr], addr, dev.snapshot, src_by_addr[addr],
                               decoded_packets=collectors[addr].decoded_packets())
            else:
                print_snapshot(labels[addr], addr, dev.snapshot, src_by_addr[addr])
    finally:
        try:
            print("Disconnecting…", flush=True)
        finally:
            await manager.disconnect_all()
    print("Done.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Retrieve a telemetry snapshot from EcoFlow River 2 and River 3",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    creds = parser.add_mutually_exclusive_group(required=True)
    creds.add_argument("--user-id", metavar="ID",
                       help="EcoFlow user ID (skip cloud login)")
    creds.add_argument("--email", metavar="EMAIL",
                       help="EcoFlow account email (fetches user ID via API)")
    parser.add_argument("--password", metavar="PASS")
    parser.add_argument("--region", default="default", choices=["default", "EU", "US"])
    parser.add_argument("--device", default=None, metavar="BT_ADDRESS",
                        help="BT address to connect to (default: all discovered devices)")
    parser.add_argument("--serial-number", metavar="SERIAL",
                        help="Device serial number (required without --scan)")
    parser.add_argument("--scan", action="store_true",
                        help="Scan for nearby EcoFlow devices before connecting")
    parser.add_argument("--timeout", type=float, default=120.0, metavar="SEC",
                        help="Failure limit for required information (default: 120 s; ignored by --watch)")
    parser.add_argument("--watch", action="store_true",
                        help="Print decoded and raw packets from either model until Ctrl+C, including unknown information")
    parser.add_argument("--debug", type=int, default=0, metavar="N")
    args = parser.parse_args()

    if args.user_id:
        user_id = args.user_id
    else:
        password = args.password or getpass.getpass(f"Password for {args.email}: ")
        print("Fetching user ID…", flush=True)
        user_id = get_user_id(args.email, password, args.region)

    try:
        asyncio.run(run(user_id, args))
    except KeyboardInterrupt:
        print("\nInterrupted.")


if __name__ == "__main__":
    main()