# ============================================================+
# 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

# ---- Optional OLED ----
try:
    from ssd1306 import SSD1306_I2C
    _i2c = I2C(0, sda=Pin(4), scl=Pin(5))
    _oled = SSD1306_I2C(128, 64, _i2c)
    HAS_OLED = True
except Exception:
    _oled = None
    HAS_OLED = False

# ---- Optional NFC (PN532 I2C) ----
try:
    from pn532 import PN532_I2C
    HAS_NFC = True
except Exception:
    HAS_NFC = False

# ---- Optional sensor (DHT22 on GP15) ----
try:
    import dht
    _dht = dht.DHT22(Pin(15))
    HAS_SENSOR = True
except Exception:
    HAS_SENSOR = False

# ---- Optional relay (GP16) ----
try:
    _relay = Pin(16, Pin.OUT)
    HAS_RELAY = True
except Exception:
    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():
    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', 'identity_tag_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', 'access_control'),
        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 IoT...', 'attempt ' + str(attempt+1) + '/3', '', ''])
            reg = identity.ensure_registered()
            break
        except Exception as e:
            print('[checkin] IoT 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', 'IoT Node ' + reg.get('result', ''), uid_short, '', ''])
    utime.sleep(1)

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

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

    print('[checkin] JWT issued')

    # it_node_checkin su identity_tag con JWT
    _show(['CHECKIN', 'IT Node checkin...', '', '', ''])
    import urequests
    from identity_iot_aes import aes_cbc_encrypt, aes_cbc_decrypt

    payload = ujson.dumps({
        'request': 'it_node_checkin',
        'node_id': str(cfg.get('node_id', '')),
    })
    cipher = aes_cbc_encrypt(payload, key, iv)
    body   = ujson.dumps({'data': cipher})

    for attempt in range(3):
        try:
            resp = urequests.post(
                cfg['identity_node_url'],
                headers={
                    'Content-Type':  'application/json',
                    'Authorization': 'Bearer ' + token,
                },
                data=body,
                timeout=20,
            )
            raw = resp.text
            resp.close()
            outer = ujson.loads(raw)
            clear = aes_cbc_decrypt(outer['data'], key, iv)
            result = ujson.loads(clear)
            if result.get('status') == 'ok':
                print('[checkin] node checkin ok')
                _show_state('CHECKIN OK', cfg.get('node_label', ''))
                utime.sleep(1)
                return {'status': 'ok', 'token': token, 'identity': identity}
            raise Exception('it_node_checkin: ' + str(result.get('code', 'error')))
        except Exception as e:
            print('[checkin] it_node_checkin attempt', attempt + 1, ':', e)
            if attempt == 2:
                raise
            utime.sleep(3)

    raise Exception('it_node_checkin failed')


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

def run():
    state = 'BOOT'

    while True:

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

        # ------------------------------------------------------------------
        elif state == 'CAPABILITY_DETECT':
            caps = detect_capabilities()
            _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'),
                '',
                'Checking config...',
                '',
            ])
            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 = 'WIFI_CONNECT'
            else:
                print('[config] missing or invalid — provisioning mode')
                state = 'PROVISIONING'

        # ------------------------------------------------------------------
        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()
            else:
                from provisioning_server import run_provisioning
                import network as _net
                _wlan = _net.WLAN(_net.STA_IF)
                _ip = _wlan.ifconfig()[0] if _wlan.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)
                # retry
                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 == 'access_control':
                from node_access_control import run_access_control
                run_access_control(cfg, caps)

            elif 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()