"""
DTU WFBLE BLE provisioning script for MicroPython ESP32-S3.
Tries multiple payload formats to provision WiFi credentials.
Writes results to /result.txt for retrieval.
"""
import bluetooth
import time
import struct

SSID = "AhoySmartHome"
PASSWORD = "106granada"
DTU_MAC = bytes([0x24, 0x19, 0x72, 0x61, 0x9D, 0x62])
# Target write handle (value handle 18)
WRITE_VH = 18
NOTIFY_VH1 = 20  # 53300004
NOTIFY_VH2 = 23  # 53300005
CCCD_VH1 = 21
CCCD_VH2 = 24

_IRQ_SCAN_RESULT = 5
_IRQ_SCAN_DONE = 6
_IRQ_PERIPHERAL_CONNECT = 7
_IRQ_PERIPHERAL_DISCONNECT = 8
_IRQ_GATTC_SERVICE_RESULT = 9
_IRQ_GATTC_SERVICE_DONE = 10
_IRQ_GATTC_CHARACTERISTIC_RESULT = 11
_IRQ_GATTC_CHARACTERISTIC_DONE = 12
_IRQ_GATTC_DESCRIPTOR_RESULT = 13
_IRQ_GATTC_DESCRIPTOR_DONE = 14
_IRQ_GATTC_READ_RESULT = 15
_IRQ_GATTC_READ_DONE = 16
_IRQ_GATTC_WRITE_DONE = 17
_IRQ_GATTC_NOTIFY = 18
_IRQ_MTU_EXCHANGED = 21

results = []
log_file = open('/result.txt', 'w')

def log(msg):
    print(msg)
    log_file.write(msg + '\n')
    log_file.flush()

ble = bluetooth.BLE()
ble.active(True)
ble.config(mtu=512)

state = {
    'conn_handle': None,
    'connected': False,
    'notify_data': [],
    'write_done': False,
    'write_status': None,
}

def build_payload_v1():
    """0x28 + ssid_len + ssid + pw_len + pw"""
    ssid_b = SSID.encode()
    pw_b = PASSWORD.encode()
    return bytes([0x28, len(ssid_b)]) + ssid_b + bytes([len(pw_b)]) + pw_b

def build_payload_v2():
    """0x01 + ssid_len + ssid + pw_len + pw"""
    ssid_b = SSID.encode()
    pw_b = PASSWORD.encode()
    return bytes([0x01, len(ssid_b)]) + ssid_b + bytes([len(pw_b)]) + pw_b

def build_payload_v3():
    """Raw HM protobuf-style: field 15 = ssid, field 16 = pw"""
    # protobuf: field 15, wire type 2 (len delimited) = tag 0x7a
    # field 16, wire type 2 = tag 0x82 0x01
    ssid_b = SSID.encode()
    pw_b = PASSWORD.encode()
    proto = bytes([0x7a, len(ssid_b)]) + ssid_b + bytes([0x82, 0x01, len(pw_b)]) + pw_b
    return proto

def build_payload_v4():
    """Full HM-framed SET_CONFIG with protobuf body"""
    # HM frame: 48 4d a3 10 00 01 <len_hi> <len_lo> 00 28 40 01 <protobuf>
    # cmd=0x28, type=0x40, sub=0x01
    ssid_b = SSID.encode()
    pw_b = PASSWORD.encode()
    # protobuf body: field 15 (ssid), field 16 (pw)
    proto = bytes([0x7a, len(ssid_b)]) + ssid_b + bytes([0x82, 0x01, len(pw_b)]) + pw_b
    # HM header: magic 484d, cmd a310, seq 0001, payload_len (big-endian 2 bytes)
    # then: 00 28 40 01 before proto? Let's try the exact known-good hex variant
    # From prior attempt: 484da3100001ae6c002840017a0d...
    # ae6c = checksum? Let me build without checksum assumption
    body = bytes([0x00, 0x28, 0x40, 0x01]) + proto
    length = len(body)
    header = bytes([0x48, 0x4d, 0xa3, 0x10, 0x00, 0x01])
    # Compute simple checksum (sum of body bytes & 0xFFFF)
    csum = sum(body) & 0xFFFF
    frame = header + bytes([csum >> 8, csum & 0xFF]) + body
    return frame

def build_payload_v5():
    """Try command 0x11 (SET_WIFI) with ssid/pw directly"""
    ssid_b = SSID.encode()
    pw_b = PASSWORD.encode()
    return bytes([0x11, len(ssid_b)]) + ssid_b + bytes([len(pw_b)]) + pw_b

def build_payload_v6():
    """Length-prefixed with 2-byte lengths"""
    ssid_b = SSID.encode()
    pw_b = PASSWORD.encode()
    return struct.pack('>BH', 0x28, len(ssid_b)) + ssid_b + struct.pack('>H', len(pw_b)) + pw_b

PAYLOADS = [
    ("v1: 0x28+ssid+pw", build_payload_v1),
    ("v2: 0x01+ssid+pw", build_payload_v2),
    ("v3: raw protobuf", build_payload_v3),
    ("v4: HM-framed protobuf", build_payload_v4),
    ("v5: 0x11+ssid+pw", build_payload_v5),
]

def irq_handler(event, data):
    if event == _IRQ_PERIPHERAL_CONNECT:
        conn_handle, addr_type, addr = data
        state['conn_handle'] = conn_handle
        state['connected'] = True
        log(f"Connected: conn_handle={conn_handle}")
    elif event == _IRQ_PERIPHERAL_DISCONNECT:
        conn_handle, addr_type, addr = data
        state['connected'] = False
        state['conn_handle'] = None
        log(f"Disconnected: conn_handle={conn_handle}")
    elif event == _IRQ_MTU_EXCHANGED:
        conn_handle, mtu = data
        log(f"MTU exchanged: {mtu}")
    elif event == _IRQ_GATTC_WRITE_DONE:
        conn_handle, value_handle, status = data
        state['write_done'] = True
        state['write_status'] = status
        log(f"Write done: vh={value_handle} status={status}")
    elif event == _IRQ_GATTC_NOTIFY:
        conn_handle, value_handle, notify_data = data
        nd = bytes(notify_data)
        log(f"NOTIFY vh={value_handle}: {nd.hex()}")
        state['notify_data'].append((value_handle, nd))
    elif event == _IRQ_GATTC_READ_RESULT:
        conn_handle, value_handle, char_data = data
        cd = bytes(char_data)
        log(f"READ_RESULT vh={value_handle}: {cd.hex()} = {cd}")

ble.irq(irq_handler)

def wait_for(condition, timeout=10, interval=0.1):
    deadline = time.time() + timeout
    while time.time() < deadline:
        if condition():
            return True
        time.sleep(interval)
    return False

def connect_to_dtu():
    log("Connecting to DTU...")
    ble.gap_connect(0, DTU_MAC)  # addr_type=0 (public)
    if wait_for(lambda: state['connected'], timeout=15):
        log("Connected!")
        time.sleep(1)  # Let MTU exchange happen
        return True
    log("Connection timeout!")
    return False

def enable_notifications():
    """Enable notifications on both CCCDs"""
    conn = state['conn_handle']
    log(f"Enabling notifications on CCCD vh={CCCD_VH1} and vh={CCCD_VH2}")
    # Write 0x0001 to CCCD to enable notifications
    notify_enable = bytes([0x01, 0x00])
    ble.gattc_write(conn, CCCD_VH1, notify_enable, 1)
    time.sleep(0.5)
    ble.gattc_write(conn, CCCD_VH2, notify_enable, 1)
    time.sleep(0.5)
    log("Notifications enabled")

def try_payload(name, payload):
    log(f"\n--- Trying {name} ---")
    log(f"Payload ({len(payload)} bytes): {payload.hex()}")
    
    state['write_done'] = False
    state['write_status'] = None
    state['notify_data'] = []
    
    conn = state['conn_handle']
    try:
        ble.gattc_write(conn, WRITE_VH, payload, 1)  # with_response=1
    except Exception as e:
        log(f"Write error: {e}")
        return False
    
    # Wait for write done
    if wait_for(lambda: state['write_done'], timeout=10):
        log(f"Write status: {state['write_status']}")
    else:
        log("Write timeout!")
    
    # Wait for notification response
    log("Waiting for notifications (15s)...")
    deadline = time.time() + 15
    while time.time() < deadline:
        if state['notify_data']:
            log(f"Got {len(state['notify_data'])} notification(s)")
            break
        time.sleep(0.2)
    
    if not state['notify_data']:
        log("No notifications received")
    
    return state['write_status'] == 0

# MAIN EXECUTION
log("=== DTU BLE Provisioning ===")
log(f"Target MAC: {':'.join('%02X' % b for b in DTU_MAC)}")
log(f"SSID: {SSID}")

if connect_to_dtu():
    enable_notifications()
    time.sleep(1)
    
    # Read current state of write characteristic
    log("\nReading vh=18 before provisioning...")
    ble.gattc_read(state['conn_handle'], WRITE_VH)
    time.sleep(2)
    
    # Try each payload
    for name, builder in PAYLOADS:
        if not state['connected']:
            log("Disconnected! Reconnecting...")
            if not connect_to_dtu():
                break
            enable_notifications()
            time.sleep(1)
        
        payload = builder()
        try_payload(name, payload)
        time.sleep(3)
    
    log("\n=== Provisioning attempts complete ===")
    if state['connected']:
        log("Disconnecting...")
        ble.gap_disconnect(state['conn_handle'])
        time.sleep(1)
else:
    log("FAILED: Could not connect to DTU")

log("DONE")
log_file.close()