# mBot2 / CyberPi Wi-Fi MQTT listener — ALL-IN-ONE (upload via mBlock, "Upload mode")
#
# Single-file merge of mbot2_listener.py + simple_mqtt.py: upload just THIS file
# to the CyberPi, nothing else. Functionally identical to the two-file setup.
#
# Turns the robot into a permanent Wi-Fi drone: it joins your LAN, connects to
# the Mosquitto broker from the docker-compose stack, reacts to drive / LED /
# mode commands, and streams encoder + color telemetry back. All piloting logic
# lives on the PC.
#
# API notes (see robot/README.md for sources):
#   cyberpi.wifi.connect / is_connect,  mbot2.drive_power,  cyberpi.led.on,
#   mbuild.quad_rgb_sensor.get_red/green/blue.
# The encoder-read and native line-offset function names are auto-detected at
# startup (their exact spelling varies by firmware), each with a fallback.

try:
    import usocket as socket
except ImportError:
    import socket
try:
    import ustruct as struct
except ImportError:
    import struct

import ujson as json
import utime as time

import cyberpi
import mbot2
import mbuild

# ---------------- configuration: EDIT THESE ----------------
WIFI_SSID = "Get yo virus myah"             # /!\ FREQUENCY BAND MUST BE 2.4GHz
WIFI_PASS = "<72260960>"
BROKER_IP = "10.58.116.212"   # IP of the machine running docker compose
BROKER_PORT = 1883

COLOR_PORT = "L1"             # quad-RGB probe used for the color sensor / seek
COLOR_MATCH_TOL = 40           # tolerance for RGB matching (lower = more sensitive, higher = more forgiving)
LINE_L_PORT = "L2"            # left probe for RGB-brightness line steering (outer probe, per hardware layout)
LINE_R_PORT = "R2"            # right probe for RGB-brightness line steering (outer probe, per hardware layout)
LINE_BASE = 22               # forward power while line-following (lowered from 30: turns detected late, needs more reaction time)
LINE_GAIN = 20               # steering strength (lowered from 30: overshooting/oscillating on turns)
LINE_CORR_MAX = 100          # effectively unclamped, matching robot.py's proven reference for this map
ENC_L_SIGN = 1                # flip to -1 if the left wheel reads backwards
ENC_R_SIGN = 1               # flip to -1 if the right wheel reads backwards
# The two motors are mounted mirrored: drive_power(l, r) with equal signs spins
# the robot in place instead of driving straight. Confirmed on hardware via
# /api/drive {"l":40,"r":40}. Flip to 1 if a unit turns out wired the other way.
DRIVE_R_SIGN = -1
# -----------------------------------------------------------

# MQTT contract — keep in sync with backend/app/protocol.py
CMD_DRIVE = "mbot2/cmd/drive"
CMD_LED = "mbot2/cmd/led"
CMD_MODE = "mbot2/cmd/mode"
CMD_STOP = "mbot2/cmd/stop"
TEL_ENC = "mbot2/tel/enc"
TEL_COLOR = "mbot2/tel/color"
TEL_STATUS = "mbot2/tel/status"
TEL_MISSION = "mbot2/tel/mission"

NAMED = {
    "red": (220, 40, 40), "green": (40, 180, 90), "blue": (50, 90, 220),
    "yellow": (235, 210, 60), "black": (20, 20, 20), "white": (235, 235, 235),
}

# mode: manual | color_seek | line_follow ; target: [r,g,b] or None ;
# method: "rgb" (brightness steer) | "offset" (native line API) ;
# distance_ticks/start_l/start_r/last_progress: color_seek mission tracking,
# set when CMD_MODE arrives with a distance_ticks (see TEL_MISSION below).
state = {"mode": "manual", "target": None, "method": "rgb", "pl": 0, "pr": 0,
         "distance_ticks": None, "start_l": 0.0, "start_r": 0.0, "last_progress": None}


# =====================================================================
# Vendored umqtt.simple — minimal MQTT 3.1.1 client for MicroPython.
# From micropython-lib (github.com/micropython/micropython-lib), MIT license.
# =====================================================================
class MQTTException(Exception):
    pass


class MQTTClient:
    def __init__(self, client_id, server, port=0, user=None, password=None,
                 keepalive=0, ssl=False, ssl_params={}):
        if port == 0:
            port = 8883 if ssl else 1883
        self.client_id = client_id
        self.sock = None
        self.server = server
        self.port = port
        self.ssl = ssl
        self.ssl_params = ssl_params
        self.pid = 0
        self.cb = None
        self.user = user
        self.pswd = password
        self.keepalive = keepalive
        self.lw_topic = None
        self.lw_msg = None
        self.lw_qos = 0
        self.lw_retain = False

    def _send_str(self, s):
        if isinstance(s, str):
            s = s.encode()
        self.sock.write(struct.pack("!H", len(s)))
        self.sock.write(s)

    def _recv_len(self):
        n = 0
        sh = 0
        while 1:
            b = self.sock.read(1)[0]
            n |= (b & 0x7F) << sh
            if not b & 0x80:
                return n
            sh += 7

    def set_callback(self, f):
        self.cb = f

    def set_last_will(self, topic, msg, retain=False, qos=0):
        assert 0 <= qos <= 2
        assert topic
        self.lw_topic = topic
        self.lw_msg = msg
        self.lw_qos = qos
        self.lw_retain = retain

    def connect(self, clean_session=True):
        self.sock = socket.socket()
        addr = socket.getaddrinfo(self.server, self.port)[0][-1]
        self.sock.connect(addr)
        if self.ssl:
            import ussl
            self.sock = ussl.wrap_socket(self.sock, **self.ssl_params)
        premsg = bytearray(b"\x10\0\0\0\0\0")
        msg = bytearray(b"\x04MQTT\x04\x02\0\0")

        sz = 10 + 2 + len(self.client_id)
        msg[6] = clean_session << 1
        if self.user is not None:
            sz += 2 + len(self.user) + 2 + len(self.pswd)
            msg[6] |= 0xC0
        if self.keepalive:
            assert self.keepalive < 65536
            msg[7] |= self.keepalive >> 8
            msg[8] |= self.keepalive & 0x00FF
        if self.lw_topic:
            sz += 2 + len(self.lw_topic) + 2 + len(self.lw_msg)
            msg[6] |= 0x4 | (self.lw_qos & 0x1) << 3 | (self.lw_qos & 0x2) << 3
            msg[6] |= self.lw_retain << 5

        i = 1
        while sz > 0x7F:
            premsg[i] = (sz & 0x7F) | 0x80
            sz >>= 7
            i += 1
        premsg[i] = sz

        self.sock.write(premsg, i + 2)
        self.sock.write(msg)
        self._send_str(self.client_id)
        if self.lw_topic:
            self._send_str(self.lw_topic)
            self._send_str(self.lw_msg)
        if self.user is not None:
            self._send_str(self.user)
            self._send_str(self.pswd)
        resp = self.sock.read(4)
        assert resp[0] == 0x20 and resp[1] == 0x02
        if resp[3] != 0:
            raise MQTTException(resp[3])
        return resp[2] & 1

    def disconnect(self):
        self.sock.write(b"\xe0\0")
        self.sock.close()

    def ping(self):
        self.sock.write(b"\xc0\0")

    def publish(self, topic, msg, retain=False, qos=0):
        if isinstance(msg, str):
            msg = msg.encode()
        pkt = bytearray(b"\x30\0\0\0")
        pkt[0] |= qos << 1 | retain
        sz = 2 + len(topic) + len(msg)
        if qos > 0:
            sz += 2
        assert sz < 2097152
        i = 1
        while sz > 0x7F:
            pkt[i] = (sz & 0x7F) | 0x80
            sz >>= 7
            i += 1
        pkt[i] = sz
        self.sock.write(pkt, i + 1)
        self._send_str(topic)
        if qos > 0:
            self.pid += 1
            pid = self.pid
            struct.pack_into("!H", pkt, 0, pid)
            self.sock.write(pkt, 2)
        self.sock.write(msg)
        if qos == 1:
            while 1:
                op = self.wait_msg()
                if op == 0x40:
                    sz = self.sock.read(1)
                    assert sz == b"\x02"
                    rcv_pid = self.sock.read(2)
                    rcv_pid = rcv_pid[0] << 8 | rcv_pid[1]
                    if pid == rcv_pid:
                        return
        elif qos == 2:
            assert 0

    def subscribe(self, topic, qos=0):
        assert self.cb is not None, "Subscribe callback is not set"
        pkt = bytearray(b"\x82\0\0\0")
        self.pid += 1
        struct.pack_into("!BH", pkt, 1, 2 + 2 + len(topic) + 1, self.pid)
        self.sock.write(pkt)
        self._send_str(topic)
        self.sock.write(qos.to_bytes(1, "little"))
        while 1:
            op = self.wait_msg()
            if op == 0x90:
                resp = self.sock.read(4)
                assert resp[1] == pkt[2] and resp[2] == pkt[3]
                if resp[3] == 0x80:
                    raise MQTTException(resp[3])
                return

    # Wait for a single incoming MQTT message and process it.
    # Subscribed messages are delivered via the callback.
    def wait_msg(self):
        res = self.sock.read(1)
        self.sock.setblocking(True)
        if res is None:
            return None
        if res == b"":
            raise OSError(-1)
        if res == b"\xd0":  # PINGRESP
            sz = self.sock.read(1)[0]
            assert sz == 0
            return None
        op = res[0]
        if op & 0xF0 != 0x30:
            return op
        sz = self._recv_len()
        topic_len = self.sock.read(2)
        topic_len = (topic_len[0] << 8) | topic_len[1]
        topic = self.sock.read(topic_len)
        sz -= topic_len + 2
        if op & 6:
            pid = self.sock.read(2)
            pid = pid[0] << 8 | pid[1]
            sz -= 2
        msg = self.sock.read(sz)
        self.cb(topic, msg)
        if op & 6 == 2:
            pkt = bytearray(b"\x40\x02\0\0")
            struct.pack_into("!H", pkt, 2, pid)
            self.sock.write(pkt)
        elif op & 6 == 4:
            assert 0
        return op

    # Checks whether a pending message from server is available.
    # If not, returns immediately with None. Otherwise processes it.
    def check_msg(self):
        self.sock.setblocking(False)
        return self.wait_msg()


# =====================================================================
# Robot logic
# =====================================================================
def _valid_rgb(t):
    try:
        return len(t) == 3 and all(isinstance(v, (int, float)) for v in t)
    except TypeError:
        return False


# ---- encoder reader: auto-detect the firmware's function name + port spelling ----
# The Makeblock library defines mbot2.EM_get_angle(port) (UUID
# fadd7b62ec703bbded9b9b63d464e425); examples never read encoders, so the port
# argument spelling isn't documented — probe a few and prefer the confirmed name.
def _make_encoder_reader():
    port_pairs = (("EM1", "EM2"), ("em1", "em2"), (1, 2), ("1", "2"))
    for name in ("EM_get_angle", "encoder_motor_get_angle",
                 "encoder_motor_get_pos", "encoder_motor_get_distance"):
        fn = getattr(mbot2, name, None)
        if not fn:
            continue
        for pl, pr in port_pairs:
            try:
                fn(pl)
                # Remember which name/port worked so the UI can show whether
                # odometry is real or the open-loop fallback.
                global ENC_SRC, ENC_PORTS
                ENC_SRC = name + ":" + str(pl)
                ENC_PORTS = (pl, pr)
                return lambda a=pl, b=pr, f=fn: (f(a), f(b))
            except Exception:
                pass
    return None


# ---- wheel-speed reader (diagnostics only) ----
# mbot2.EM_get_speed(port) is confirmed in the Makeblock API table (UUID
# 9c703762b00091af8ecdf503fab52745), same undocumented port argument as
# EM_get_angle, so reuse whatever spelling the encoder probe settled on.
def _make_speed_reader():
    fn = getattr(mbot2, "EM_get_speed", None)
    if not fn or not ENC_PORTS:
        return None
    pl, pr = ENC_PORTS
    try:
        fn(pl)
        return lambda a=pl, b=pr, f=fn: (f(a), f(b))
    except Exception:
        return None


# ---- board health probes (all optional; absent APIs simply report None) ----
# Verified in the Makeblock API table: cyberpi.get_battery() (3988e0aa…),
# cyberpi.get_extra_battery() (e8e71442…), cyberpi.get_firmware_version()
# (84130d80…), cyberpi.timer.get() (5535075a…). All take no arguments.
# Note: the API exposes no Wi-Fi RSSI or IP-address call, so signal strength
# cannot be reported — only the boolean wifi.is_connect().
def _probe(fn_path, *args):
    obj = cyberpi
    for part in fn_path.split("."):
        obj = getattr(obj, part, None)
        if obj is None:
            return None
    try:
        return obj(*args)
    except Exception:
        return None


# ---- native line-offset reader ----
# mbuild.quad_rgb_sensor.get_offset_track(index) -> ~ -100..100 (line deviation:
# negative = line left, positive = line right). `index` is the integer mBuild
# port (1), not a probe string. Names vary slightly by firmware, so auto-detect.
def _make_offset_reader():
    for name in ("get_offset_track", "get_line_offset", "get_offset"):
        fn = getattr(mbuild.quad_rgb_sensor, name, None)
        if fn:
            try:
                fn(1)
                global OFFSET_SRC
                OFFSET_SRC = name
                return lambda f=fn: f(1)
            except Exception:
                pass
    return None


ENC_SRC = "open-loop"     # set by _make_encoder_reader if a real encoder answers
ENC_PORTS = None
OFFSET_SRC = None         # set by _make_offset_reader if the native API exists

_read_enc = _make_encoder_reader()
_read_offset = _make_offset_reader()
_read_speed = _make_speed_reader()
FIRMWARE = _probe("get_firmware_version")

_pseudo_l = 0.0
_pseudo_r = 0.0
_last_ms = time.ticks_ms()


def read_encoders():
    global _pseudo_l, _pseudo_r, _last_ms
    if _read_enc:
        l, r = _read_enc()
        return ENC_L_SIGN * l, ENC_R_SIGN * r
    # Open-loop fallback: integrate commanded power into pseudo-encoder degrees
    # (~286.5 deg/s per 100% power for an 80 mm wheel at ~200 mm/s). Apply the
    # same wheel-direction signs as the real-encoder path so odometry matches.
    now = time.ticks_ms()
    dt = time.ticks_diff(now, _last_ms) / 1000.0
    _last_ms = now
    _pseudo_l += state["pl"] / 100.0 * 286.5 * dt
    _pseudo_r += state["pr"] / 100.0 * 286.5 * dt
    return ENC_L_SIGN * _pseudo_l, ENC_R_SIGN * _pseudo_r


def _brightness(port):
    return (int(mbuild.quad_rgb_sensor.get_red(port))
            + int(mbuild.quad_rgb_sensor.get_green(port))
            + int(mbuild.quad_rgb_sensor.get_blue(port)))


def read_color():
    r = int(mbuild.quad_rgb_sensor.get_red(COLOR_PORT))
    g = int(mbuild.quad_rgb_sensor.get_green(COLOR_PORT))
    b = int(mbuild.quad_rgb_sensor.get_blue(COLOR_PORT))
    name = "background"
    best = 90 * 90 * 3
    for n, (cr, cg, cb) in NAMED.items():
        d = (r - cr) ** 2 + (g - cg) ** 2 + (b - cb) ** 2
        if d < best:
            best = d
            name = n
    return r, g, b, name


LINE_DEBUG = {"raw": None, "corr": None, "lb": None, "rb": None}  # surfaced via TEL_COLOR for live tuning


def line_powers():
    """Return (left, right) power for line following per the selected method."""
    if state["method"] == "offset" and _read_offset:
        off = _read_offset()          # expected ~ -100..100, 0 = centered
        # Sign flipped: the "negative = line left" assumption in get_offset_track's
        # docs doesn't hold on this unit — confirmed by hardware turning the wrong
        # way on a real corner.
        corr = -LINE_GAIN * (off / 100.0)
        corr = max(-LINE_CORR_MAX, min(LINE_CORR_MAX, corr))
        LINE_DEBUG["raw"], LINE_DEBUG["corr"] = off, corr
        return LINE_BASE + corr, LINE_BASE - corr
    # RGB-brightness steer: the darker probe is over the line.
    lb, rb = _brightness(LINE_L_PORT), _brightness(LINE_R_PORT)
    err = (rb - lb) / 765.0           # >0 => left darker => line on the left
    corr = LINE_GAIN * err
    LINE_DEBUG["lb"], LINE_DEBUG["rb"] = lb, rb
    corr = max(-LINE_CORR_MAX, min(LINE_CORR_MAX, corr))
    LINE_DEBUG["raw"], LINE_DEBUG["corr"] = round(err, 3), corr
    return LINE_BASE - corr, LINE_BASE + corr


def drive(l, r):
    l = max(-100, min(100, int(l)))
    r = max(-100, min(100, int(r)))
    state["pl"], state["pr"] = l, r
    mbot2.drive_power(l, DRIVE_R_SIGN * r)


def publish_status(client):
    """Status + board diagnostics for the web UI's debug panel (1 Hz).

    Every health field is best-effort: a firmware that lacks a call reports null
    rather than crashing the control loop.
    """
    spd = None
    if _read_speed:
        try:
            sl, sr = _read_speed()
            spd = [sl, sr]
        except Exception:
            spd = None
    client.publish(TEL_STATUS, json.dumps({
        "mode": state["mode"], "method": state["method"], "online": True,
        "battery": _probe("get_battery"),
        "extra_battery": _probe("get_extra_battery"),
        "fw": FIRMWARE,
        "uptime": _probe("timer.get"),
        "wifi": _probe("wifi.is_connect"),
        "enc_src": ENC_SRC,
        "offset_src": OFFSET_SRC,
        "spd": spd,
    }))


def on_message(topic, payload):
    try:
        p = json.loads(payload) if payload else {}
    except Exception:
        return
    t = topic.decode() if isinstance(topic, bytes) else topic
    if t == CMD_DRIVE:
        if state["mode"] == "manual":
            drive(p.get("l", 0), p.get("r", 0))
    elif t == CMD_STOP:
        state["mode"] = "manual"
        state["target"] = None
        drive(0, 0)
    elif t == CMD_LED:
        r, g, b = int(p.get("r", 0)), int(p.get("g", 0)), int(p.get("b", 0))
        idx = int(p.get("index", 0))
        if idx == 0:
            cyberpi.led.on(r, g, b)            # all LEDs
        else:
            cyberpi.led.on(r, g, b, id=idx)    # single LED 1..5
    elif t == CMD_MODE:
        state["mode"] = p.get("mode", "manual")
        tgt = p.get("target")
        state["target"] = tgt if _valid_rgb(tgt) else None
        if p.get("method") in ("rgb", "offset"):
            state["method"] = p["method"]
        dist = p.get("distance_ticks")
        if state["mode"] in ("color_seek", "line_follow") and isinstance(dist, (int, float)) and dist > 0:
            state["distance_ticks"] = dist
            state["start_l"], state["start_r"] = read_encoders()
            state["last_progress"] = None
        else:
            state["distance_ticks"] = None
        if state["mode"] == "manual":
            drive(0, 0)


def ensure_wifi():
    if not cyberpi.wifi.is_connect():
        cyberpi.wifi.connect(WIFI_SSID, WIFI_PASS)
        while not cyberpi.wifi.is_connect():
            cyberpi.console.println("wifi...")
            time.sleep(0.5)
    cyberpi.console.println("wifi ok")


def connect_mqtt(client):
    ensure_wifi()
    client.connect()
    for topic in (CMD_DRIVE, CMD_LED, CMD_MODE, CMD_STOP):
        client.subscribe(topic)
    publish_status(client)
    cyberpi.console.println("mqtt ok")


def main():
    client = MQTTClient("mbot2-" + str(time.ticks_ms()), BROKER_IP,
                        port=BROKER_PORT, keepalive=30)
    client.set_callback(on_message)
    # Last-Will: broker publishes this if the robot drops, so the UI clears.
    client.set_last_will(TEL_STATUS, json.dumps({"mode": "offline", "online": False}))
    connect_mqtt(client)

    last_enc = last_col = last_status = last_ping = time.ticks_ms()
    while True:
        try:
            client.check_msg()   # non-blocking; fires on_message for commands
            now = time.ticks_ms()
            mode = state["mode"]

            # Read the color sensor only when a mode needs it or telemetry is due.
            col_due = time.ticks_diff(now, last_col) >= 200
            color = read_color() if (mode != "manual" or col_due) else None

            if mode == "color_seek" and state["target"]:
                drive(25, 25)
                if state["distance_ticks"]:
                    el, er = read_encoders()
                    traveled = (abs(el - state["start_l"]) + abs(er - state["start_r"])) / 2.0
                    pct = int(min(100, traveled / state["distance_ticks"] * 100))
                    if pct != state["last_progress"]:
                        state["last_progress"] = pct
                        client.publish(TEL_MISSION, json.dumps(
                            {"event": "progress", "progress": pct}))
                tr, tg, tb = state["target"]
                if color and abs(color[0] - tr) < COLOR_MATCH_TOL and abs(color[1] - tg) < COLOR_MATCH_TOL \
                        and abs(color[2] - tb) < COLOR_MATCH_TOL:
                    drive(0, 0)
                    state["mode"] = "manual"
                    was_mission = state["distance_ticks"] is not None
                    state["distance_ticks"] = None
                    publish_status(client)     # report the stop immediately
                    if was_mission:
                        client.publish(TEL_MISSION, json.dumps({"event": "arrival"}))
            elif mode == "line_follow":
                l, r = line_powers()
                drive(l, r)
                if state["distance_ticks"]:
                    el, er = read_encoders()
                    traveled = (abs(el - state["start_l"]) + abs(er - state["start_r"])) / 2.0
                    pct = int(min(100, traveled / state["distance_ticks"] * 100))
                    if pct != state["last_progress"]:
                        state["last_progress"] = pct
                        client.publish(TEL_MISSION, json.dumps(
                            {"event": "progress", "progress": pct}))
                if state["target"] and color:  # optional stop-on-color
                    tr, tg, tb = state["target"]
                    if abs(color[0] - tr) < COLOR_MATCH_TOL and abs(color[1] - tg) < COLOR_MATCH_TOL \
                            and abs(color[2] - tb) < COLOR_MATCH_TOL:
                        drive(0, 0)
                        state["mode"] = "manual"
                        was_mission = state["distance_ticks"] is not None
                        state["distance_ticks"] = None
                        publish_status(client)
                        if was_mission:
                            client.publish(TEL_MISSION, json.dumps({"event": "arrival"}))

            if time.ticks_diff(now, last_enc) >= 50:      # 20 Hz encoders
                el, er = read_encoders()
                client.publish(TEL_ENC, json.dumps({"l": el, "r": er, "t": now}))
                last_enc = now
            if col_due:                                    # 5 Hz color
                if color is None:
                    color = read_color()
                client.publish(TEL_COLOR, json.dumps(
                    {"r": color[0], "g": color[1], "b": color[2], "name": color[3],
                     "line_raw": LINE_DEBUG["raw"], "line_corr": LINE_DEBUG["corr"],
                     "line_lb": LINE_DEBUG["lb"], "line_rb": LINE_DEBUG["rb"]}))
                last_col = now
            if time.ticks_diff(now, last_status) >= 1000:
                publish_status(client)
                last_status = now
            if time.ticks_diff(now, last_ping) >= 15000:   # keep session alive
                client.ping()
                last_ping = now
            time.sleep(0.01)
        except Exception as e:
            # Safety first: stop the motors, then rebuild the Wi-Fi/MQTT link so a
            # transient drop can't leave the robot driving at the last power.
            try:
                mbot2.drive_power(0, 0)
            except Exception:
                pass
            cyberpi.console.println("link lost, reconnecting")
            try:
                client.disconnect()
            except Exception:
                pass
            while True:
                try:
                    connect_mqtt(client)
                    break
                except Exception:
                    time.sleep(2)
            last_enc = last_col = last_status = last_ping = time.ticks_ms()


main()