# ============================================================+
# 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://ap.parta.app/API/V1/identity_layer/identity_iot/index.php")
#   await identity.ensure_registered(tenant_id=1)
#   token = await identity.authenticate()
# ============================================================+

import ujson
import ubinascii
import uhashlib
import urandom
import machine
import network
import uasyncio as asyncio

try:
    import uurequests as requests
    _ASYNC = False
except ImportError:
    import urequests as requests
    _ASYNC = False

# Pure-Python Ed25519 — replaced by natmod in production
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.
    Communicates with identity_layer/identity_iot PHP backend.
    """

    # Storage keys (on flash filesystem)
    _F_SEED    = '/identity_seed.bin'     # 32 bytes raw seed
    _F_PK      = '/identity_pk.b64'       # public key b64
    _F_JWT     = '/identity_jwt.txt'      # current JWT
    _F_NODEID  = '/identity_node_uid.txt' # hashed node UID

    def __init__(
        self,
        base_url: str,
        app_id: str = 'identity_iot',
        node_type: str = 'sensor',
        node_label: str = None,
        tenant_id: int = 0,
        comm_key: bytes = None,
        comm_iv: bytes = 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
        # Bootstrap AES-CBC keys — replaced after init_key if used
        self._comm_key  = comm_key
        self._comm_iv   = comm_iv

    # =========================================================================
    # 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.
        Stores JWT on flash for next session.
        """
        # Try existing JWT first
        jwt = self._load_jwt()
        if jwt and not self._is_jwt_expired(jwt):
            return jwt

        node_uid = self._get_node_uid()

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

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

        # Sign nonce
        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()

        # Verify
        resp = self._post({
            'request':       'verify_signature',
            'challenge_id':  challenge_id,
            'signature_b64': sig_b64,
            'node_uid':      node_uid,
        })
        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 the hashed node UID (safe to share)."""
        return self._get_node_uid()

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

    def _generate_keypair(self):
        """Generate Ed25519 keypair from secure random seed, persist to flash."""
        seed = bytes(urandom.getrandbits(8) for _ in range(32))
        pk, _ = ed25519_keypair_from_seed(seed)
        pk_b64 = ubinascii.b2a_base64(pk).decode().strip()
        # Persist
        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 — device fingerprint, hashed before sending
    # =========================================================================

    def _get_node_uid(self):
        """
        Returns SHA-256 hex of device UID.
        Cached on flash after first call.
        """
        try:
            with open(self._F_NODEID, 'r') as f:
                return f.read().strip()
        except OSError:
            pass

        # Pico W: machine.unique_id() returns 8 bytes
        raw_uid = machine.unique_id()
        h = uhashlib.sha256(raw_uid)
        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)
            # MicroPython: use time.time() — requires NTP sync
            import utime
            return utime.time() >= exp
        except Exception:
            return True  # assume expired on any error

    # =========================================================================
    # HTTP transport — ITEMS envelope (AES-CBC) or plain JSON
    # =========================================================================

    def _post(self, inner: dict) -> dict:
        """
        POST to backend.
        If comm_key/iv provided: ITEMS AES-CBC envelope.
        Otherwise: plain JSON (for PoC / testing).
        """
        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:
        """Plain JSON POST — for PoC and testing."""
        body = ujson.dumps({'data': ujson.dumps(inner)})
        resp = requests.post(
            self.base_url,
            headers={'Content-Type': 'application/json'},
            data=body,
        )
        outer = ujson.loads(resp.text)
        resp.close()
        # Backend returns encrypted data — decode if possible, else return raw
        if 'data' in outer:
            try:
                return ujson.loads(outer['data'])
            except Exception:
                return outer
        return outer

    def _post_encrypted(self, inner: dict) -> dict:
        """ITEMS AES-CBC envelope POST."""
        from identity_iot_aes import aes_cbc_encrypt, aes_cbc_decrypt
        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,
        )
        outer = ujson.loads(resp.text)
        resp.close()
        if 'data' not in outer:
            raise Exception('Missing outer.data in response')
        clear = aes_cbc_decrypt(outer['data'], self._comm_key, self._comm_iv)
        return ujson.loads(clear)


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

def _normalize_b64(s: str) -> str:
    """Normalize base64 string for ubinascii."""
    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