"""Ed25519 cryptographic identity for A2A agents."""

from __future__ import annotations

import hashlib
from pathlib import Path

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)

# Base58 alphabet (Bitcoin-style, no ambiguous chars)
_B58_ALPHABET = b"123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"


def _b58encode(data: bytes) -> str:
    """Encode bytes to base58 string."""
    n = int.from_bytes(data, "big")
    result = bytearray()
    while n > 0:
        n, remainder = divmod(n, 58)
        result.append(_B58_ALPHABET[remainder])
    # Preserve leading zeros
    for byte in data:
        if byte == 0:
            result.append(_B58_ALPHABET[0])
        else:
            break
    return bytes(reversed(result)).decode("ascii")


def _b58decode(s: str) -> bytes:
    """Decode base58 string to bytes."""
    n = 0
    for char in s:
        n = n * 58 + _B58_ALPHABET.index(char.encode("ascii"))
    # Count leading '1's (zero bytes)
    leading_zeros = 0
    for char in s:
        if char == "1":
            leading_zeros += 1
        else:
            break
    result = n.to_bytes((n.bit_length() + 7) // 8, "big") if n else b""
    return b"\x00" * leading_zeros + result


def _derive_peer_id(public_key: Ed25519PublicKey) -> str:
    """Derive a PeerID from an Ed25519 public key.

    Format: a2a:<base58(sha256(raw_pubkey)[:20])>
    Produces a ~28-char human-readable identifier.
    """
    raw = public_key.public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw,
    )
    digest = hashlib.sha256(raw).digest()[:20]
    return f"a2a:{_b58encode(digest)}"


class PeerIdentity:
    """An agent's cryptographic identity backed by an Ed25519 keypair."""

    def __init__(
        self,
        private_key: Ed25519PrivateKey,
    ) -> None:
        self._private_key = private_key
        self._public_key = private_key.public_key()
        self._peer_id = _derive_peer_id(self._public_key)

    @classmethod
    def generate(cls) -> PeerIdentity:
        """Generate a new random identity."""
        return cls(Ed25519PrivateKey.generate())

    @classmethod
    def load(cls, directory: Path, passphrase: str = "") -> PeerIdentity:
        """Load identity from PEM files in *directory*.

        Expects ``private.pem`` (and optionally ``public.pem``).
        """
        priv_path = directory / "private.pem"
        enc: serialization.KeySerializationEncryption
        if passphrase:
            enc = serialization.BestAvailableEncryption(passphrase.encode())
        else:
            enc = serialization.NoEncryption()

        priv_bytes = priv_path.read_bytes()
        pwd = passphrase.encode() if passphrase else None
        private_key = serialization.load_pem_private_key(priv_bytes, password=pwd)
        if not isinstance(private_key, Ed25519PrivateKey):
            raise TypeError(f"Expected Ed25519 private key, got {type(private_key).__name__}")
        return cls(private_key)

    def save(self, directory: Path, passphrase: str = "") -> None:
        """Save identity to PEM files in *directory*."""
        directory.mkdir(parents=True, exist_ok=True)

        enc: serialization.KeySerializationEncryption
        if passphrase:
            enc = serialization.BestAvailableEncryption(passphrase.encode())
        else:
            enc = serialization.NoEncryption()

        priv_pem = self._private_key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.PKCS8,
            enc,
        )
        (directory / "private.pem").write_bytes(priv_pem)

        pub_pem = self._public_key.public_bytes(
            serialization.Encoding.PEM,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        (directory / "public.pem").write_bytes(pub_pem)

    @property
    def private_key(self) -> Ed25519PrivateKey:
        return self._private_key

    @property
    def public_key(self) -> Ed25519PublicKey:
        return self._public_key

    @property
    def peer_id(self) -> str:
        return self._peer_id

    def peer_id_short(self) -> str:
        """First 8 chars of the base58 portion (after 'a2a:')."""
        return self._peer_id[:12]  # "a2a:" + 8 chars

    def public_key_pem(self) -> str:
        """Return the public key as a PEM string."""
        return self._public_key.public_bytes(
            serialization.Encoding.PEM,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        ).decode("ascii")

    def public_key_raw(self) -> bytes:
        """Return the raw 32-byte public key."""
        return self._public_key.public_bytes(
            serialization.Encoding.Raw, serialization.PublicFormat.Raw,
        )

    def sign(self, data: bytes) -> bytes:
        """Sign data with the private key."""
        return self._private_key.sign(data)

    @staticmethod
    def verify(
        data: bytes, signature: bytes, public_key: Ed25519PublicKey,
    ) -> bool:
        """Verify a signature against a public key. Returns False on failure."""
        try:
            public_key.verify(signature, data)
            return True
        except Exception:
            return False

    @staticmethod
    def load_public_key_pem(pem: str) -> Ed25519PublicKey:
        """Load a public key from PEM string."""
        key = serialization.load_pem_public_key(pem.encode("ascii"))
        if not isinstance(key, Ed25519PublicKey):
            raise TypeError(f"Expected Ed25519 public key, got {type(key).__name__}")
        return key

    def __repr__(self) -> str:
        return f"PeerIdentity({self._peer_id})"
