#!/usr/bin/env python3
"""
A scooter that is not there.

The load generators only ever talk: they publish telemetry and never listen. A ride needs the
other half of the loop, because allocation is deliberately asynchronous. booking-service dispatches
UNLOCK and returns PENDING_UNLOCK; the ride only goes ACTIVE when a lock result comes back on
`lock.status` carrying the ride id. With no hardware and nothing listening, every ride sits in
PENDING_UNLOCK until the sweeper expires it.

So this subscribes to the command topic and answers, which is exactly what the lock in the field
does:

    beeb/{id}/down   {"command": "UNLOCK", "correlationId": "<ride id>"}
    beeb/{id}/up     {"lock": "UNLOCKED", "correlationId": "<echoed>"}

The echo matters. iot-bridge does not read the id out of the reply (it claims the one it sent);
what the echo does is mark the result *solicited*, which is what separates "the lock opened because
we asked" from "someone opened it in the street".

    python3 virtual_scooter.py --prefix DEMO --count 6 --mqtt-port 51883

Runs until interrupted, reporting each command it answers.
"""

import argparse
import json
import random
import socket
import struct
import sys
import threading
import time

from mqtt_load import encode_remaining_length, encode_string

# How long a real lock takes to physically open. Answering instantly would hide any race in the
# app's PENDING_UNLOCK handling, which is one of the things this is here to exercise.
UNLOCK_DELAY_SECONDS = 1.2


class MqttDevice:
    """CONNECT, SUBSCRIBE and PUBLISH: enough MQTT 3.1.1 to be a device."""

    def __init__(self, host: str, port: int, client_id: str):
        self.sock = socket.create_connection((host, port), timeout=10)
        self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.lock = threading.Lock()
        self.buffer = b""

        payload = encode_string(client_id)
        variable = encode_string("MQTT") + bytes([4, 0x02]) + struct.pack(">H", 60)
        body = variable + payload
        self.sock.sendall(bytes([0x10]) + encode_remaining_length(len(body)) + body)
        ack = self.sock.recv(4)
        if len(ack) < 4 or ack[3] != 0:
            raise RuntimeError(f"broker refused the connection: {ack!r}")
        # The connect timeout must not survive into the listen loop: a device is idle most of the
        # time, so a read timeout is the normal state, not a dead broker. Keepalive is satisfied by
        # the telemetry this same client publishes.
        self.sock.settimeout(None)

    def subscribe(self, topic_filter: str) -> None:
        body = struct.pack(">H", 1) + encode_string(topic_filter) + bytes([0])   # packet id, QoS 0
        with self.lock:
            self.sock.sendall(bytes([0x82]) + encode_remaining_length(len(body)) + body)

    def publish(self, topic: str, payload: bytes) -> None:
        body = encode_string(topic) + payload
        with self.lock:
            self.sock.sendall(bytes([0x30]) + encode_remaining_length(len(body)) + body)

    def _read(self, n: int) -> bytes:
        while len(self.buffer) < n:
            chunk = self.sock.recv(4096)
            if not chunk:
                raise ConnectionError("broker closed the connection")
            self.buffer += chunk
        out, self.buffer = self.buffer[:n], self.buffer[n:]
        return out

    def read_packet(self):
        """Next inbound packet as (type, payload). Returns None for anything not a PUBLISH."""
        header = self._read(1)[0]
        multiplier, length = 1, 0
        while True:
            digit = self._read(1)[0]
            length += (digit & 0x7F) * multiplier
            if not digit & 0x80:
                break
            multiplier *= 128
        body = self._read(length) if length else b""
        if header & 0xF0 != 0x30:            # not a PUBLISH (SUBACK, PINGRESP, ...)
            return None
        topic_len = struct.unpack(">H", body[:2])[0]
        topic = body[2:2 + topic_len].decode(errors="replace")
        # QoS 0 only, so there is no packet identifier between the topic and the payload.
        return topic, body[2 + topic_len:]

    def close(self) -> None:
        try:
            self.sock.sendall(bytes([0xE0, 0x00]))
        except OSError:
            pass
        finally:
            self.sock.close()


class Fleet:
    """Every simulated device on one connection, which is how a gateway would really do it."""

    def __init__(self, host, port, device_ids, lat, lng, radius):
        self.client = MqttDevice(host, port, f"virtual-{int(time.time())}")
        self.devices = {}
        for device in device_ids:
            self.devices[device] = {
                "lat": lat + random.uniform(-1, 1) * radius / 111_320.0,
                "lng": lng + random.uniform(-1, 1) * radius / 90_000.0,
                "battery": random.randint(55, 99),
                "locked": True,
            }
            self.client.subscribe(f"beeb/{device}/down")
        self.running = True
        self.answered = 0

    def listen(self):
        while self.running:
            try:
                packet = self.client.read_packet()
            except (ConnectionError, OSError):
                return
            if packet is None:
                continue
            topic, payload = packet
            device = topic.split("/")[1] if topic.count("/") == 2 else None
            if device not in self.devices:
                continue
            try:
                command = json.loads(payload.decode())
            except ValueError:
                print(f"  {device}: unreadable command {payload!r}")
                continue
            threading.Thread(target=self.act, args=(device, command), daemon=True).start()

    def act(self, device, command):
        name = str(command.get("command", "")).upper()
        correlation = command.get("correlationId", "")
        print(f"  -> {device} received {name} (correlationId={correlation or 'none'})")
        if name not in ("UNLOCK", "LOCK"):
            # An honest device reports what it cannot do rather than staying silent, which would
            # look identical to a dead SIM.
            print(f"  <- {device} ignoring unsupported command {name}")
            return
        time.sleep(UNLOCK_DELAY_SECONDS)
        state = "UNLOCKED" if name == "UNLOCK" else "LOCKED"
        self.devices[device]["locked"] = state == "LOCKED"
        self.client.publish(f"beeb/{device}/up", json.dumps({
            "lock": state,
            # Echoed back: this is what marks the result solicited.
            "correlationId": correlation,
            "errorCode": 0,
        }).encode())
        self.answered += 1
        print(f"  <- {device} reports {state}")

    def telemetry_loop(self, interval):
        while self.running:
            for device, s in self.devices.items():
                # A moving scooter drifts; a parked one does not. Ride tracks are built from these
                # frames, so a stationary fleet would draw a track with no shape to it.
                if not s["locked"]:
                    s["lat"] += random.uniform(-0.00012, 0.00012)
                    s["lng"] += random.uniform(-0.00012, 0.00012)
                    s["battery"] = max(1, s["battery"] - random.choice([0, 0, 1]))
                self.client.publish(f"beeb/{device}/up", json.dumps({
                    "battery": s["battery"],
                    "signal": random.randint(3, 5),
                    "lat": round(s["lat"], 6),
                    "lng": round(s["lng"], 6),
                    "lock": "LOCKED" if s["locked"] else "UNLOCKED",
                }).encode())
            time.sleep(interval)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--mqtt-host", default="localhost")
    ap.add_argument("--mqtt-port", type=int, default=1883)
    ap.add_argument("--prefix", default="DEMO")
    ap.add_argument("--count", type=int, default=6)
    ap.add_argument("--lat", type=float, default=24.7136)
    ap.add_argument("--lng", type=float, default=46.6753)
    ap.add_argument("--radius", type=float, default=400)
    ap.add_argument("--interval", type=float, default=3.0, help="telemetry period, seconds")
    args = ap.parse_args()

    ids = [f"{args.prefix}{i + 1:03d}" for i in range(args.count)]
    fleet = Fleet(args.mqtt_host, args.mqtt_port, ids, args.lat, args.lng, args.radius)
    print(f"{len(ids)} virtual scooters listening on beeb/+/down, reporting every "
          f"{args.interval:g}s: {', '.join(ids)}")

    threading.Thread(target=fleet.telemetry_loop, args=(args.interval,), daemon=True).start()
    try:
        fleet.listen()
    except KeyboardInterrupt:
        pass
    finally:
        fleet.running = False
        fleet.client.close()
        print(f"answered {fleet.answered} lock commands")
    return 0


if __name__ == "__main__":
    sys.exit(main())