"""Mesh identity key management (X25519)."""
from __future__ import annotations

import base64
import hashlib
import os
import socket
from dataclasses import dataclass

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.x25519 import (
    X25519PrivateKey,
    X25519PublicKey,
)

from persistence import MUGI_DIR


MESH_DIR = os.path.join(MUGI_DIR, "mesh")
IDENTITY_KEY_PATH = os.path.join(MESH_DIR, "identity.key")
IDENTITY_PUB_PATH = os.path.join(MESH_DIR, "identity.pub")


@dataclass(frozen=True)
class MeshIdentity:
    private_key: X25519PrivateKey
    public_key_b64: str
    fingerprint: str
    hostname: str


def _b64(data: bytes) -> str:
    return base64.b64encode(data).decode("ascii")


def _fingerprint_from_pub_bytes(pub_raw: bytes) -> str:
    digest = hashlib.sha256(pub_raw).digest()
    return "SHA256:" + base64.b64encode(digest).decode("ascii")


def _write_mode_600(path: str, data: bytes) -> None:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    try:
        os.write(fd, data)
    finally:
        os.close(fd)
    try:
        os.chmod(path, 0o600)
    except OSError:
        pass


def _load_private_key(path: str) -> X25519PrivateKey:
    with open(path, "rb") as f:
        raw = f.read()
    key = serialization.load_pem_private_key(raw, password=None)
    if not isinstance(key, X25519PrivateKey):
        raise ValueError("identity.key is not an X25519 private key")
    return key


def _public_b64(private_key: X25519PrivateKey) -> str:
    pub_raw = private_key.public_key().public_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PublicFormat.Raw,
    )
    return _b64(pub_raw)


def _hostname() -> str:
    try:
        host = (socket.gethostname() or "").strip()
        return host or "mugi-host"
    except Exception:
        return "mugi-host"


def load_or_create_identity() -> MeshIdentity:
    os.makedirs(MESH_DIR, exist_ok=True)
    if not os.path.isfile(IDENTITY_KEY_PATH):
        private_key = X25519PrivateKey.generate()
        pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        _write_mode_600(IDENTITY_KEY_PATH, pem)
        _write_mode_600(IDENTITY_PUB_PATH, (_public_b64(private_key) + "\n").encode("utf-8"))

    private_key = _load_private_key(IDENTITY_KEY_PATH)
    pub_raw = private_key.public_key().public_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PublicFormat.Raw,
    )
    public_key_b64 = _b64(pub_raw)
    fingerprint = _fingerprint_from_pub_bytes(pub_raw)
    if not os.path.isfile(IDENTITY_PUB_PATH):
        _write_mode_600(IDENTITY_PUB_PATH, (public_key_b64 + "\n").encode("utf-8"))
    return MeshIdentity(
        private_key=private_key,
        public_key_b64=public_key_b64,
        fingerprint=fingerprint,
        hostname=_hostname(),
    )


def public_key_b64_to_fingerprint(public_key_b64: str) -> str:
    raw = base64.b64decode(public_key_b64.encode("ascii"))
    return _fingerprint_from_pub_bytes(raw)


def decode_public_key(public_key_b64: str) -> X25519PublicKey:
    raw = base64.b64decode(public_key_b64.encode("ascii"))
    return X25519PublicKey.from_public_bytes(raw)
