import websocket
import struct
import random

# --- Подключение ---
URL = "wss://insert-me.ru/mqtt"
USER = "esp8266-insert-me"
PASS = "insert-me"
CLIENT_ID = f"pyclient_{random.randint(0, 10000)}"
TOPIC = "#"

def encode_length(length: int) -> bytes:
    enc = bytearray()
    while True:
        digit = length & 0x7F
        length >>= 7
        if length > 0:
            digit |= 0x80
        enc.append(digit)
        if length == 0:
            break
    return bytes(enc)

def decode_length(data: bytes, start: int = 0) -> tuple:
    """Декодирует остающуюся длину из data, возвращает (значение, количество_прочитанных_байт)."""
    length = 0
    multiplier = 1
    pos = start
    while True:
        if pos >= len(data):
            raise ValueError("Incomplete length encoding")
        digit = data[pos]
        length += (digit & 0x7F) * multiplier
        multiplier <<= 7
        pos += 1
        if not (digit & 0x80):
            break
    return length, pos - start

def build_connect(client_id: str, username: str = None, password: str = None,
                  keepalive: int = 60, clean_start: bool = True) -> bytes:
    protocol_name = b"MQTT"
    protocol_level = 5
    connect_flags = 0
    if clean_start:
        connect_flags |= 0x02
    if username is not None:
        connect_flags |= 0x80
        if password is not None:
            connect_flags |= 0x40

    properties = b"\x00"  # без свойств

    client_id_bytes = client_id.encode('utf-8')
    payload = struct.pack("!H", len(client_id_bytes)) + client_id_bytes

    if username is not None:
        user_bytes = username.encode('utf-8')
        payload += struct.pack("!H", len(user_bytes)) + user_bytes
    if password is not None:
        pass_bytes = password.encode('utf-8')
        payload += struct.pack("!H", len(pass_bytes)) + pass_bytes

    variable_header = (
        struct.pack("!H", len(protocol_name)) +
        protocol_name +
        bytes([protocol_level]) +
        bytes([connect_flags]) +
        struct.pack("!H", keepalive) +
        properties
    )

    remaining_length = len(variable_header) + len(payload)
    fixed_header = b"\x10" + encode_length(remaining_length)
    return fixed_header + variable_header + payload

def build_subscribe(topic: str, packet_id: int = 1, qos: int = 0) -> bytes:
    fixed_header = b"\x82"
    topic_bytes = topic.encode('utf-8')
    properties = b"\x00"
    payload = struct.pack("!H", packet_id) + properties
    payload += struct.pack("!H", len(topic_bytes)) + topic_bytes
    payload += bytes([qos])
    remaining_length = len(payload)
    return fixed_header + encode_length(remaining_length) + payload

def parse_publish(data: bytes) -> tuple:
    """
    Распарсить PUBLISH пакет MQTT 5.
    Возвращает (topic, payload, qos, packet_id) или (None, None, None, None) если ошибка.
    """
    if not data or data[0] & 0xF0 != 0x30:
        return None, None, None, None

    # Пропускаем фиксированный заголовок (1 байт) и декодируем длину
    pos = 1
    remaining_length, consumed = decode_length(data, pos)
    pos += consumed

    # Читаем топик
    if pos + 2 > len(data):
        return None, None, None, None
    topic_len = struct.unpack("!H", data[pos:pos+2])[0]
    pos += 2
    if pos + topic_len > len(data):
        return None, None, None, None
    topic = data[pos:pos+topic_len].decode('utf-8', errors='replace')
    pos += topic_len

    # QoS и Packet ID
    qos = (data[0] >> 1) & 0x03
    packet_id = None
    if qos > 0:
        if pos + 2 > len(data):
            return None, None, None, None
        packet_id = struct.unpack("!H", data[pos:pos+2])[0]
        pos += 2

    # Свойства: сначала длина свойств (1 байт, т.к. у нас всегда 0)
    if pos >= len(data):
        return None, None, None, None
    prop_len = data[pos]
    pos += 1
    if prop_len > 0:
        # Если есть свойства, пропускаем их (упрощённо)
        pos += prop_len

    # Всё что осталось – payload
    payload = data[pos:].decode('utf-8', errors='replace')
    return topic, payload, qos, packet_id

def main():
    # Создаём WebSocket
    ws = websocket.create_connection(URL, subprotocols=["mqtt"], timeout=10)

    # 1. CONNECT
    connect_pkt = build_connect(CLIENT_ID, username=USER, password=PASS, keepalive=60)
    print("Sending CONNECT:", connect_pkt.hex())
    ws.send_binary(connect_pkt)

    # 2. Читаем CONNACK
    try:
        data = ws.recv()
        print("Received:", data.hex() if data else "(empty)")
        if data and data[0] == 0x20:
            if len(data) >= 4:
                reason_code = data[3]
                if reason_code == 0:
                    print("✅ Connected successfully")
                else:
                    print(f"❌ Connection refused, reason code: {reason_code}")
                    ws.close()
                    return
            else:
                print("Malformed CONNACK")
                ws.close()
                return
        else:
            print("Not a CONNACK received")
            ws.close()
            return
    except Exception as e:
        print(f"Error receiving CONNACK: {e}")
        ws.close()
        return

    # 3. Подписка
    sub_pkt = build_subscribe(TOPIC, packet_id=1, qos=0)
    print("Sending SUBSCRIBE:", sub_pkt.hex())
    ws.send_binary(sub_pkt)

    suback = ws.recv()
    print("SUBACK:", suback.hex())
    if suback and suback[0] == 0x90:
        reason = suback[-1]
        if reason == 0:
            print(f"✅ Subscribed to '{TOPIC}'")
        else:
            print(f"❌ Subscribe failed, reason code: {reason}")
            ws.close()
            return

    # 4. Цикл приёма
    print("Waiting for messages... (Ctrl+C to stop)")
    try:
        ws.settimeout(5)
        while True:
            try:
                msg = ws.recv()
            except websocket.WebSocketTimeoutException:
                ws.send_binary(b"\xC0\x00")   # PINGREQ
                continue
            if not msg:
                break

            # Если PINGRESP – игнорируем
            if msg[0] == 0xD0 and len(msg) == 2 and msg[1] == 0x00:
                continue

            # Парсим PUBLISH
            topic, payload, qos, pid = parse_publish(msg)
            if topic is not None:
                print(f"📨 Topic: {topic}, Payload: {payload}")
                if qos > 0:
                    print(f"   (QoS={qos}, PacketID={pid})")
            else:
                print(f"Other MQTT packet: {msg.hex()}")
    except KeyboardInterrupt:
        pass
    finally:
        ws.close()

if __name__ == "__main__":
    main()