# ============================================================+
# Identity IoT — Provisioning Server
# (c) Copyright : ae Aeonian Engineering Limited - Hong Kong
# (c) Copyright : WIDE di D. Papa - Naples - Italy
# ============================================================+
# Runs when no config.json found on flash.
# 1. Connects WiFi in AP mode (fallback) or waits on existing LAN
# 2. Broadcasts UDP beacon: "identity-iot|{sn}"
# 3. Exposes HTTP on port 8080:
#      GET  /       → node info (plaintext)
#      POST /provision → AES-encrypted config payload → saves config.json
# ============================================================+

import network
import socket
import ujson
import ubinascii
import machine
import utime

from identity_iot_aes import aes_cbc_encrypt, aes_cbc_decrypt

# Hardcoded provisioning keys — same in Flutter app
# Change only with coordinated firmware + app update
PROG_KEY = b'IdentityIoTProv!Key_WIDE_AEL_32b'   # 32 bytes
PROG_IV  = b'IdentityIoTIV16b'                    # 16 bytes

UDP_PORT  = 5000
HTTP_PORT = 8080
BEACON_INTERVAL_MS = 2000


def _get_sn():
    return ubinascii.hexlify(machine.unique_id()).decode()


def _get_ip():
    wlan = network.WLAN(network.STA_IF)
    if wlan.isconnected():
        return wlan.ifconfig()[0]
    return '0.0.0.0'


def _start_ap():
    """Fallback: start AP so app can connect even without existing WiFi."""
    sn = _get_sn()
    ssid = 'IdentityIoT-' + sn[:8]
    ap = network.WLAN(network.AP_IF)
    ap.active(True)
    ap.config(essid=ssid, password='provision', security=4)   # WPA2
    utime.sleep(2)
    print('[provisioning] AP started:', ssid, '— ip:', ap.ifconfig()[0])
    return ap.ifconfig()[0]


def _beacon_loop(udp_sock, beacon_bytes, stop_flag):
    """Send UDP broadcast beacon. Called between HTTP polls."""
    try:
        udp_sock.sendto(beacon_bytes, ('255.255.255.255', UDP_PORT))
    except Exception as e:
        print('[beacon] error:', e)


def _parse_http_request(conn):
    """Read HTTP request from socket. Returns (method, path, body)."""
    try:
        data = b''
        conn.settimeout(5)
        while True:
            chunk = conn.recv(1024)
            if not chunk:
                break
            data += chunk
            if b'\r\n\r\n' in data:
                # Check if we have full body
                header_end = data.index(b'\r\n\r\n') + 4
                headers_raw = data[:header_end].decode('utf-8', 'ignore')
                body = data[header_end:]
                content_length = 0
                for line in headers_raw.split('\r\n'):
                    if line.lower().startswith('content-length:'):
                        content_length = int(line.split(':')[1].strip())
                while len(body) < content_length:
                    chunk = conn.recv(1024)
                    if not chunk:
                        break
                    body += chunk
                first_line = headers_raw.split('\r\n')[0]
                parts = first_line.split(' ')
                method = parts[0] if len(parts) > 0 else 'GET'
                path   = parts[1] if len(parts) > 1 else '/'
                return method, path, body.decode('utf-8', 'ignore')
    except Exception as e:
        print('[http] parse error:', e)
    return 'GET', '/', ''


def _http_response(conn, status, body):
    resp = (
        'HTTP/1.1 ' + status + '\r\n'
        'Content-Type: application/json\r\n'
        'Content-Length: ' + str(len(body)) + '\r\n'
        'Connection: close\r\n'
        '\r\n' + body
    )
    conn.sendall(resp.encode())


def run_provisioning():
    """
    Block until a valid config is received from the Flutter app.
    Returns config dict ready to save.
    """
    sn = _get_sn()
    print('[provisioning] SN:', sn)

    # Try STA first, fall back to AP
    wlan_sta = network.WLAN(network.STA_IF)
    if not wlan_sta.isconnected():
        ip = _start_ap()
    else:
        ip = wlan_sta.ifconfig()[0]
        print('[provisioning] using existing WiFi, ip:', ip)

    beacon_msg = ujson.dumps({
        'type': 'identity-iot',
        'sn':   sn,
        'ip':   ip,
        'port': HTTP_PORT,
    }).encode()

    # UDP socket for beacon
    udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        udp.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    except Exception:
        pass
    udp.setblocking(False)

    # HTTP server socket
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(('0.0.0.0', HTTP_PORT))
    srv.listen(1)
    srv.setblocking(False)

    print('[provisioning] HTTP listening on', ip + ':' + str(HTTP_PORT))
    print('[provisioning] UDP beacon on port', UDP_PORT)

    last_beacon = 0
    received_cfg = None

    while received_cfg is None:

        # Beacon
        now = utime.ticks_ms()
        if utime.ticks_diff(now, last_beacon) >= BEACON_INTERVAL_MS:
            try:
                udp.sendto(beacon_msg, ('255.255.255.255', UDP_PORT))
            except Exception:
                pass
            last_beacon = now

        # HTTP accept (non-blocking)
        try:
            conn, addr = srv.accept()
            print('[provisioning] connection from', addr)
            method, path, body = _parse_http_request(conn)

            if method == 'GET' and (path == '/' or path == '/info'):
                info = ujson.dumps({
                    'sn':     sn,
                    'ip':     ip,
                    'status': 'awaiting_provision',
                    'type':   'identity-iot',
                    'port':   HTTP_PORT,
                })
                _http_response(conn, '200 OK', info)
                conn.close()

            elif method == 'POST' and path == '/provision':
                try:
                    outer = ujson.loads(body)
                    cipher_b64 = outer.get('data', '')
                    if not cipher_b64:
                        raise Exception('missing data')

                    clear = aes_cbc_decrypt(cipher_b64, PROG_KEY, PROG_IV)
                    cfg   = ujson.loads(clear)

                    # Basic validation
                    required = ['wifi_ssid', 'wifi_pass', 'comm_key', 'comm_iv',
                                'identity_iot_url', 'tenant_id', 'node_type', 'node_id']
                    for k in required:
                        if k not in cfg:
                            raise Exception('missing field: ' + k)

                    _http_response(conn, '200 OK', ujson.dumps({'status': 'ok', 'sn': sn}))
                    conn.close()
                    received_cfg = cfg
                    print('[provisioning] config received OK')

                except Exception as e:
                    print('[provisioning] provision error:', e)
                    _http_response(conn, '400 Bad Request', ujson.dumps({'status': 'error', 'msg': str(e)}))
                    conn.close()

            else:
                _http_response(conn, '404 Not Found', '{}')
                conn.close()

        except OSError:
            # No connection yet — normal in non-blocking mode
            utime.sleep_ms(100)

    udp.close()
    srv.close()
    return received_cfg