#!/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. The script waits for the
expected packet types or the --timeout limit, then prints the snapshot and exits.
River 2 requires valid PD, EMS, and BMS core payloads. Optional heartbeats and
other packets are listed separately. Use --watch to print every decoded/raw
packet until Ctrl+C. Watch output can contain private identifiers.

Usage:
  poetry run python tools/ecoflow-bt.py --user-id <id> --scan
  poetry run python tools/ecoflow-bt.py --email you@example.com --scan
  poetry run 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 decode_payload, decode_protobuf, decode_river2_payload
except ImportError as exc:
    sys.exit(f"Missing dependency: {exc}\n"
             "Run: poetry install -E ecoflow-bt")

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

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],
        # cmd_set=0xFE: DisplayPropertyUpload (0x15) + RuntimePropertyUpload (0x16)
        "expected_fields": {"ef_pkt_02_fe_15", "ef_pkt_02_fe_16"},
    },
    "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"}

# ── 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 = decode_payload(src, cmd_set, cmd_id, payload)
        if lines is not None:
            return ["  " + ln for ln in lines]

    # Fallback: generic protobuf decode
    pb = decode_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]) -> 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}:")
        for line in _decode_snapshot_packet(key, payload_hex, src_candidates):
            print(line)

    print()


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


class PacketCollector:
    """Track valid River 2 heartbeats and optionally print each raw packet."""

    def __init__(self, address: str, profile: dict, watch: bool = False) -> None:
        self.address = address
        self.label = profile["label"]
        self.required_heartbeats = profile.get("required_heartbeats", {})
        self.valid_heartbeats: set[str] = set()
        self.watch = watch
        self.error: Exception | None = None

    def observe_snapshot(self, snap) -> None:
        """Seed receipt tracking for callers that start after reception."""
        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]:
        return {
            f"{label} ({key})" for key, label in self.required_heartbeats.items()
            if key not in self.valid_heartbeats
        }

    def progress(self) -> str:
        return (f"{len(self.valid_heartbeats)}/{len(self.required_heartbeats)} "
                "required heartbeats received")

    def receive(self, packet: EcoFlowPacket) -> None:
        # Re-raise callback/output failures from collect(), which owns cleanup.
        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 = decode_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)


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)


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."""
    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_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 not watch:
                        print(f"  {addr}: waiting for {missing}", flush=True)
                    last_missing[addr] = missing
            else:
                if 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_heartbeats"):
            collector = PacketCollector(addr, d, watch=watch)
            collectors[addr] = collector
            dev.on_packet(collector.receive)
        print(f"Registered {d['label']}  {address}  serial={d['serial_number']}")

    try:
        if watch:
            print("\nWatching 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:
            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 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()