# ============================================================+
# Identity IoT — Boot Manager
# (c) Copyright : ae Aeonian Engineering Limited - Hong Kong
# (c) Copyright : WIDE di D. Papa - Naples - Italy
# ============================================================+
# State machine: BOOT → CONFIG_CHECK → PROVISIONING | WIFI_CONNECT
#                → CHECKIN → OPERATIONAL
# ============================================================+

import ujson
import utime
import network
import machine
import ubinascii
from machine import Pin, I2C

# Hardware instances — initialised in detect_capabilities() after config load
_oled   = None
_nfc    = None
_dht    = None
_relay  = None
HAS_OLED    = False
HAS_NFC     = False
HAS_SENSOR  = False
HAS_RELAY   = False

CONFIG_PATH = 'config.json'

STATES = [
    'BOOT',
    'CAPABILITY_DETECT',
    'CONFIG_CHECK',
    'PROVISIONING',
    'WIFI_CONNECT',
    'CHECKIN',
    'OPERATIONAL',
    'ERROR',
]


# =============================================================================
# Display helpers
# =============================================================================

def _show(lines):
    if not HAS_OLED or _oled is None:
        print(' | '.join(str(l) for l in lines))
        return
    _oled.fill(0)
    for i, line in enumerate(lines[:5]):
        _oled.text(str(line)[:16], 0, i * 12)
    _oled.show()


def _show_state(state, detail=''):
    _show(['Identity IoT', 'WIDE / AEL', '', state, detail])


# =============================================================================
# Capability detection
# =============================================================================

def detect_capabilities(cfg=None):
    global _oled, _nfc, _dht, _relay
    global HAS_OLED, HAS_NFC, HAS_SENSOR, HAS_RELAY

    pinout = cfg.get('pinout', {}) if cfg else {}
    i2c0_sda  = pinout.get('i2c0_sda', 4)
    i2c0_scl  = pinout.get('i2c0_scl', 5)
    i2c1_sda  = pinout.get('i2c1_sda', 6)
    i2c1_scl  = pinout.get('i2c1_scl', 7)
    dht_pin   = pinout.get('dht_pin',  15)
    relay_pin = pinout.get('relay_pin', 16)

    # OLED
    try:
        from ssd1306 import SSD1306_I2C
        _i2c0 = I2C(0, sda=Pin(i2c0_sda), scl=Pin(i2c0_scl))
        _oled = SSD1306_I2C(128, 64, _i2c0)
        HAS_OLED = True
    except Exception:
        _oled = None
        HAS_OLED = False

    # NFC (PN532 as provisioning dongle only — not present in normal operation)
    try:
        from pn532 import PN532_I2C
        _i2c1 = I2C(1, sda=Pin(i2c1_sda), scl=Pin(i2c1_scl))
        _nfc = PN532_I2C(_i2c1)
        HAS_NFC = True
    except Exception as e:
        print('[nfc] not detected:', e)
        _nfc = None
        HAS_NFC = False

    # Sensor (DHT22)
    try:
        import dht
        _dht = dht.DHT22(Pin(dht_pin))
        HAS_SENSOR = True
    except Exception:
        _dht = None
        HAS_SENSOR = False

    # Relay
    try:
        _relay = Pin(relay_pin, Pin.OUT)
        HAS_RELAY = True
    except Exception:
        _relay = None
        HAS_RELAY = False

    caps = {
        'has_oled':   HAS_OLED,
        'has_nfc':    HAS_NFC,
        'has_sensor': HAS_SENSOR,
        'has_relay':  HAS_RELAY,
    }
    print('[caps]', caps)
    return caps


# =============================================================================
# Config
# =============================================================================

def load_config():
    try:
        with open(CONFIG_PATH) as f:
            return ujson.load(f)
    except OSError:
        return None


def save_config(cfg):
    with open(CONFIG_PATH, 'w') as f:
        ujson.dump(cfg, f)


def validate_config(cfg):
    required = ['wifi_ssid', 'wifi_pass', 'comm_key', 'comm_iv',
                'tenant_id', 'node_type', 'node_id', 'identity_iot_url']
    for k in required:
        if k not in cfg:
            print('[config] missing key:', k)
            return False
    return True


# =============================================================================
# WiFi
# =============================================================================

def connect_wifi(cfg, timeout=30):
    _show_state('WIFI_CONNECT', cfg['wifi_ssid'])
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    utime.sleep(2)

    if not wlan.isconnected():
        wlan.connect(cfg['wifi_ssid'], cfg['wifi_pass'])
        t = timeout
        while t > 0:
            if wlan.isconnected() and wlan.ifconfig()[0] != '0.0.0.0':
                break
            utime.sleep(1)
            t -= 1

    if not wlan.isconnected() or wlan.ifconfig()[0] == '0.0.0.0':
        _show_state('ERROR', 'WiFi failed')
        raise Exception('WiFi connect failed')

    if cfg.get('wifi_static'):
        wlan.ifconfig((
            cfg['wifi_ip'],
            cfg['wifi_mask'],
            cfg['wifi_gw'],
            cfg['wifi_dns'],
        ))

    ip = wlan.ifconfig()[0]
    print('[wifi] connected:', ip)
    _show_state('WIFI_OK', ip)
    utime.sleep(3)
    return wlan


# =============================================================================
# Checkin
# =============================================================================

def do_checkin(cfg):
    _show_state('CHECKIN', '')
    from identity_iot import IdentityIoT

    key = cfg['comm_key'].encode() if isinstance(cfg['comm_key'], str) else cfg['comm_key']
    iv  = cfg['comm_iv'].encode()  if isinstance(cfg['comm_iv'],  str) else cfg['comm_iv']

    identity = IdentityIoT(
        base_url   = cfg['identity_iot_url'],
        node_type  = cfg.get('node_type', 'sensor'),
        node_label = cfg.get('node_label', ''),
        tenant_id  = cfg['tenant_id'],
        node_id    = cfg['node_id'],
        comm_key   = key,
        comm_iv    = iv,
        wifi_ssid  = cfg['wifi_ssid'],
        wifi_pass  = cfg['wifi_pass'],
    )

    # Register — retry loop
    for attempt in range(3):
        try:
            _show(['CHECKIN', 'Registering...', 'attempt ' + str(attempt+1) + '/3', '', ''])
            reg = identity.ensure_registered()
            break
        except Exception as e:
            print('[checkin] register attempt', attempt + 1, ':', e)
            if attempt == 2:
                _show_state('ERROR', 'Check network')
                raise
            utime.sleep(3)

    uid_short = reg.get('node_uid', '')[:12] + '...'
    _show(['CHECKIN', 'Node ' + reg.get('result', ''), uid_short, '', ''])
    utime.sleep(1)

    # Authenticate → JWT
    _show(['CHECKIN', 'Challenge...', '', '', ''])
    utime.sleep(0.5)
    _show(['CHECKIN', 'Signing...', '', 'Please wait', ''])
    token = identity.authenticate()

    if not token:
        _show_state('ERROR', 'Auth failed')
        raise Exception('authenticate returned no token')

    print('[checkin] JWT issued')
    _show_state('CHECKIN OK', cfg.get('node_label', ''))
    utime.sleep(1)

    return {'status': 'ok', 'token': token, 'identity': identity}


# =============================================================================
# Boot Manager — main entry point
# =============================================================================

def run():
    state = 'BOOT'
    cfg   = None
    caps  = {'has_oled': False, 'has_nfc': False, 'has_sensor': False, 'has_relay': False}

    while True:

        # ------------------------------------------------------------------
        if state == 'BOOT':
            _show(['Identity IoT', 'WIDE / AEL', '', 'Booting...', ''])
            utime.sleep(1)
            state = 'CONFIG_CHECK'

        # ------------------------------------------------------------------
        elif state == 'CONFIG_CHECK':
            cfg = load_config()
            if cfg and validate_config(cfg):
                print('[config] found and valid')
                state = 'CAPABILITY_DETECT'
            else:
                print('[config] missing or invalid — provisioning mode')
                caps = detect_capabilities()
                state = 'PROVISIONING'

        # ------------------------------------------------------------------
        elif state == 'CAPABILITY_DETECT':
            caps = detect_capabilities(cfg)
            _show([
                'Identity IoT',
                'NFC:'    + ('Y' if caps['has_nfc']    else 'N') +
                ' SNS:'   + ('Y' if caps['has_sensor'] else 'N') +
                ' RLY:'   + ('Y' if caps['has_relay']  else 'N'),
                '',
                'READY',
                '',
            ])
            utime.sleep(1)
            state = 'WIFI_CONNECT'

        # ------------------------------------------------------------------
        elif state == 'PROVISIONING':
            if caps.get('has_nfc'):
                from provisioning_nfc import run_provisioning_nfc
                _show(['PROVISIONING', 'NFC mode', '', 'Tap writer badge', ''])
                cfg = run_provisioning_nfc(_nfc)
            else:
                from provisioning_server import run_provisioning
                wlan_sta = network.WLAN(network.STA_IF)
                _ip = wlan_sta.ifconfig()[0] if wlan_sta.isconnected() else '0.0.0.0'
                _show(['PROVISIONING', 'IP: ' + _ip, 'HTTP:8080', 'UDP :5000', 'Waiting app...'])
                cfg = run_provisioning()
            save_config(cfg)
            print('[provisioning] config saved — rebooting')
            _show_state('CONFIG SAVED', 'Rebooting...')
            utime.sleep(2)
            machine.reset()

        # ------------------------------------------------------------------
        elif state == 'WIFI_CONNECT':
            try:
                connect_wifi(cfg)
                state = 'CHECKIN'
            except Exception as e:
                print('[wifi] error:', e)
                _show_state('ERROR', 'WiFi')
                utime.sleep(10)
                state = 'WIFI_CONNECT'

        # ------------------------------------------------------------------
        elif state == 'CHECKIN':
            try:
                import gc
                gc.collect()
                checkin_result = do_checkin(cfg)
                identity = checkin_result.get('identity')
                state = 'OPERATIONAL'
            except Exception as e:
                print('[checkin] error:', e)
                _show_state('ERROR', 'Checkin')
                utime.sleep(15)
                state = 'CHECKIN'

        # ------------------------------------------------------------------
        elif state == 'OPERATIONAL':
            node_type = cfg.get('node_type', 'sensor')
            _show_state('OPERATIONAL', node_type)

            if node_type == 'sensor':
                from node_sensor import run_sensor
                run_sensor(cfg, caps)

            elif node_type == 'relay':
                from node_relay import run_relay
                run_relay(cfg, caps)

            else:
                print('[operational] unknown node_type:', node_type)
                _show_state('ERROR', 'node_type?')
                utime.sleep(30)

        # ------------------------------------------------------------------
        elif state == 'ERROR':
            _show_state('ERROR', 'halted')
            utime.sleep(60)
            machine.reset()