# ============================================================+
# Identity IoT — MicroPython Module
# (c) Copyright : ae Aeonian Engineering Limited - Hong Kong
# (c) Copyright : WIDE di D. Papa - Naples - Italy
# ============================================================+
# Ed25519 challenge/response authentication for IoT nodes.
# Compatible with identity_layer/identity_iot PHP backend.
#
# Install via mip:
#   import mip
#   mip.install("https://git.aeonianengineering.net/wide/identity-micropython/raw/branch/main/package.json")
#
# Usage:
#   from identity_iot import IdentityIoT
#   identity = IdentityIoT(base_url="https://api.parta.app/API/V1/identity_layer/identity_iot/")
#   identity.ensure_registered()
#   token = identity.authenticate()
# ============================================================+

import ujson
import ubinascii
import uhashlib
import urandom
import machine
import network
import utime

try:
    import uurequests as requests
except ImportError:
    import urequests as requests

from identity_iot_ed25519 import ed25519_sign, ed25519_keypair_from_seed


class IdentityIoT:
    """
    Identity IoT authentication service for MicroPython.
    Device-bound Ed25519 keypair — private key never leaves the device.
    """

    _F_SEED   = '/identity_seed.bin'
    _F_PK     = '/identity_pk.b64'
    _F_JWT    = '/identity_jwt.txt'
    _F_NODEID = '/identity_node_uid.txt'

    def __init__(
        self,
        base_url:   str,
        app_id:     str   = 'identity_iot',
        node_type:  str   = 'sensor',
        node_label: str   = None,
        tenant_id         = 0,
        node_id:    str   = None,
        comm_key:   bytes = None,
        comm_iv:    bytes = None,
        wifi_ssid:  str   = None,
        wifi_pass:  str   = None,
    ):
        self.base_url   = base_url
        self.app_id     = app_id
        self.node_type  = node_type
        self.node_label = node_label
        self.tenant_id  = tenant_id
        self.node_id    = node_id
        self._comm_key  = comm_key
        self._comm_iv   = comm_iv
        self._wifi_ssid = wifi_ssid
        self._wifi_pass = wifi_pass

    # =========================================================================
    # Public API
    # =========================================================================

    def ensure_registered(self):
        """Register node if not already registered. Idempotent."""
        pk = self._load_pk()
        if pk is None:
            pk = self._generate_keypair()
        node_uid = self._get_node_uid()
        resp = self._post({
            'request':        'register_node',
            'public_key_b64': pk,
            'node_uid':       node_uid,
            'node_type':      self.node_type,
            'node_label':     self.node_label or '',
            'tenant_id':      self.tenant_id,
        })
        if resp.get('status') != 'ok':
            raise Exception('register_node failed: ' + str(resp))
        return resp

    def authenticate(self):
        """Full challenge/response flow. Returns JWT string."""
        jwt = self._load_jwt()
        if jwt and not self._is_jwt_expired(jwt):
            return jwt

        node_uid = self._get_node_uid()

        resp = self._post({'request': 'challenge', 'node_uid': node_uid, 'tenant_id': self.tenant_id})
        if resp.get('status') != 'ok':
            raise Exception('challenge failed: ' + str(resp))

        challenge_id = resp['challenge_id']
        nonce_b64    = resp['nonce_b64']

        nonce_bytes = ubinascii.a2b_base64(_normalize_b64(nonce_b64))
        seed        = self._load_seed()
        signature   = ed25519_sign(seed, nonce_bytes)
        sig_b64     = ubinascii.b2a_base64(signature).decode().strip()

        resp = self._post({
            'request':       'verify_signature',
            'challenge_id':  challenge_id,
            'signature_b64': sig_b64,
            'node_uid':      node_uid,
            'tenant_id':     self.tenant_id,
        })
        if resp.get('status') != 'ok':
            raise Exception('verify_signature failed: ' + str(resp))

        token = resp.get('token', '')
        if token:
            self._save_jwt(token)
        return token

    def get_node_uid(self):
        return self._get_node_uid()

    # =========================================================================
    # Keypair management
    # =========================================================================

    def _generate_keypair(self):
        seed   = bytes(urandom.getrandbits(8) for _ in range(32))
        pk, _  = ed25519_keypair_from_seed(seed)
        pk_b64 = ubinascii.b2a_base64(pk).decode().strip()
        with open(self._F_SEED, 'wb') as f:
            f.write(seed)
        with open(self._F_PK, 'w') as f:
            f.write(pk_b64)
        return pk_b64

    def _load_seed(self):
        try:
            with open(self._F_SEED, 'rb') as f:
                return f.read()
        except OSError:
            raise StateError('No seed found. Call ensure_registered() first.')

    def _load_pk(self):
        try:
            with open(self._F_PK, 'r') as f:
                return f.read().strip()
        except OSError:
            return None

    # =========================================================================
    # Node UID
    # =========================================================================

    def _get_node_uid(self):
        try:
            with open(self._F_NODEID, 'r') as f:
                return f.read().strip()
        except OSError:
            pass
        # Fresh UID — hardware ID + random entropy so each new registration is unique
        raw_uid  = machine.unique_id()
        entropy  = bytes([urandom.getrandbits(8) for _ in range(16)])
        h        = uhashlib.sha256(raw_uid + entropy)
        uid_hex  = ubinascii.hexlify(h.digest()).decode()
        with open(self._F_NODEID, 'w') as f:
            f.write(uid_hex)
        return uid_hex

    # =========================================================================
    # JWT management
    # =========================================================================

    def _load_jwt(self):
        try:
            with open(self._F_JWT, 'r') as f:
                return f.read().strip()
        except OSError:
            return None

    def _save_jwt(self, token):
        with open(self._F_JWT, 'w') as f:
            f.write(token)

    def _is_jwt_expired(self, jwt):
        try:
            parts = jwt.split('.')
            if len(parts) != 3:
                return True
            payload_b64  = _normalize_b64(parts[1])
            payload_json = ubinascii.a2b_base64(payload_b64).decode()
            payload      = ujson.loads(payload_json)
            exp          = payload.get('exp', 0)
            return utime.time() >= exp
        except Exception:
            return True

    # =========================================================================
    # WiFi reconnect
    # =========================================================================

    def _ensure_wifi(self):
        wlan = network.WLAN(network.STA_IF)
        if wlan.isconnected():
            return True
        if not self._wifi_ssid:
            return False
        print('WiFi lost, reconnecting...')
        wlan.active(True)
        wlan.connect(self._wifi_ssid, self._wifi_pass)
        timeout = 15
        while not wlan.isconnected() and timeout > 0:
            utime.sleep(1)
            timeout -= 1
        return wlan.isconnected()

    # =========================================================================
    # HTTP transport
    # =========================================================================

    def _post(self, inner: dict) -> dict:
        if self._comm_key and self._comm_iv:
            return self._post_encrypted(inner)
        return self._post_plain(inner)

    def _post_plain(self, inner: dict) -> dict:
        for attempt in range(3):
            try:
                if not self._ensure_wifi():
                    raise Exception('WiFi not connected')
                body = ujson.dumps({'data': ujson.dumps(inner)})
                resp = requests.post(
                    self.base_url,
                    headers={'Content-Type': 'application/json'},
                    data=body,
                    timeout=20
                )
                outer = ujson.loads(resp.text)
                resp.close()
                if 'data' in outer:
                    try:
                        return ujson.loads(outer['data'])
                    except Exception:
                        return outer
                return outer
            except Exception as e:
                print('plain retry', attempt + 1, ':', e)
                utime.sleep(3)
        raise Exception('POST failed after 3 attempts')

    def _post_encrypted(self, inner: dict) -> dict:
        from identity_iot_aes import aes_cbc_encrypt, aes_cbc_decrypt
        for attempt in range(3):
            try:
                if not self._ensure_wifi():
                    raise Exception('WiFi not connected')
                inner_json = ujson.dumps(inner)
                cipher_b64 = aes_cbc_encrypt(inner_json, self._comm_key, self._comm_iv)
                body       = ujson.dumps({'data': cipher_b64})
                resp       = requests.post(
                    self.base_url,
                    headers={'Content-Type': 'application/json'},
                    data=body,
                    timeout=20
                )
                print('status:', resp.status_code) 
                print('raw:', resp.text[:100])     
                outer = ujson.loads(resp.text)
                resp.close()
                if 'data' not in outer:
                    raise Exception('Missing outer.data')
                clear = aes_cbc_decrypt(outer['data'], self._comm_key, self._comm_iv)
                return ujson.loads(clear)
            except Exception as e:
                print('enc retry', attempt + 1, ':', e)
                utime.sleep(3)
        raise Exception('POST failed after 3 attempts')


# =============================================================================
# Utility
# =============================================================================

def _normalize_b64(s: str) -> str:
    s = s.strip().replace('\n', '').replace('\r', '').replace(' ', '')
    s = s.replace('-', '+').replace('_', '/')
    pad = len(s) % 4
    if pad:
        s += '=' * (4 - pad)
    return s


class StateError(Exception):
    pass