"""Ed25519 device identity for OpenClaw gateway authentication.

Manages a persistent Ed25519 keypair used to sign challenge-nonce payloads
during the gateway WebSocket handshake (protocol v3 device identity flow).
Also stores gateway-issued device tokens for reconnection.
"""

import base64
import hashlib
import json
import logging
import os
import time
from dataclasses import dataclass
from pathlib import Path

from nacl.signing import SigningKey

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class DeviceIdentity:
    device_id: str     # SHA-256 hex fingerprint of public key
    public_key: bytes  # Raw 32-byte Ed25519 public key
    private_key: bytes # Raw 32-byte Ed25519 seed


def _b64url_encode(data: bytes) -> str:
    """Base64URL encode without padding (RFC 4648 §5)."""
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64url_decode(s: str) -> bytes:
    """Base64URL decode with padding restoration."""
    padding = 4 - (len(s) % 4)
    if padding != 4:
        s += "=" * padding
    return base64.urlsafe_b64decode(s)


def load_or_create_identity(data_dir: str) -> DeviceIdentity:
    """Load an existing device identity from disk, or generate and persist a new one."""
    path = Path(data_dir) / "device-identity.json"

    if path.exists():
        try:
            data = json.loads(path.read_text())
            if data.get("version") == 1:
                identity = DeviceIdentity(
                    device_id=data["deviceId"],
                    public_key=_b64url_decode(data["publicKey"]),
                    private_key=_b64url_decode(data["privateKey"]),
                )
                logger.info("Loaded device identity %s", identity.device_id[:12])
                return identity
        except (json.JSONDecodeError, KeyError, Exception) as e:
            logger.warning("Corrupt device identity file, regenerating: %s", e)

    signing_key = SigningKey.generate()
    public_key = signing_key.verify_key.encode()
    private_key = signing_key.encode()
    device_id = hashlib.sha256(public_key).hexdigest()

    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps({
        "version": 1,
        "deviceId": device_id,
        "publicKey": _b64url_encode(public_key),
        "privateKey": _b64url_encode(private_key),
        "createdAtMs": int(time.time() * 1000),
    }, indent=2))

    logger.info("Created new device identity %s", device_id[:12])
    return DeviceIdentity(
        device_id=device_id,
        public_key=public_key,
        private_key=private_key,
    )


def build_and_sign_device_payload(
    identity: DeviceIdentity,
    *,
    nonce: str,
    token: str | None,
    client_id: str,
    client_mode: str,
    role: str,
    scopes: list[str],
) -> dict:
    """Build the signed device auth payload for the gateway connect request.

    Returns the ``device`` field dict for the connect params:
    ``{id, publicKey, signature, signedAt, nonce}``.
    """
    signed_at_ms = int(time.time() * 1000)

    # Pipe-delimited v2 payload — must match the gateway's buildDeviceAuthPayload()
    # in src/gateway/device-auth.ts exactly.
    payload_str = "|".join([
        "v2",
        identity.device_id,
        client_id,
        client_mode,
        role,
        ",".join(scopes),
        str(signed_at_ms),
        token or "",
        nonce,
    ])
    payload_bytes = payload_str.encode("utf-8")
    signing_key = SigningKey(identity.private_key)
    signature = signing_key.sign(payload_bytes).signature

    return {
        "id": identity.device_id,
        "publicKey": _b64url_encode(identity.public_key),
        "signature": _b64url_encode(signature),
        "signedAt": signed_at_ms,
        "nonce": nonce,
    }


# ---------------------------------------------------------------------------
# Device token persistence
# ---------------------------------------------------------------------------

def _token_store_path(data_dir: str) -> Path:
    return Path(data_dir) / "device-tokens.json"


def _load_store(path: Path) -> dict:
    if path.exists():
        try:
            return json.loads(path.read_text())
        except (json.JSONDecodeError, OSError):
            pass
    return {}


def _save_store(path: Path, store: dict):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(store, indent=2))


def _token_key(device_id: str, role: str) -> str:
    return f"{device_id}:{role}"


def load_device_token(data_dir: str, device_id: str, role: str) -> str | None:
    """Load a previously stored device token, or None."""
    store = _load_store(_token_store_path(data_dir))
    return store.get(_token_key(device_id, role))


def store_device_token(data_dir: str, device_id: str, role: str, token: str):
    """Persist a gateway-issued device token."""
    path = _token_store_path(data_dir)
    store = _load_store(path)
    store[_token_key(device_id, role)] = token
    _save_store(path, store)
    logger.debug("Stored device token for %s:%s", device_id[:12], role)


def clear_device_token(data_dir: str, device_id: str, role: str):
    """Remove a stored device token (e.g. after revocation)."""
    path = _token_store_path(data_dir)
    store = _load_store(path)
    key = _token_key(device_id, role)
    if key in store:
        del store[key]
        _save_store(path, store)
        logger.debug("Cleared device token for %s:%s", device_id[:12], role)
