from __future__ import annotations

import base64
import hashlib
import os
from pathlib import Path

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


def _b64url_encode(data: bytes) -> str:
    """Base64url encode without padding (matches Go's RawURLEncoding)."""
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def _b64url_decode(s: str) -> bytes:
    """Base64url decode, adding padding as needed."""
    padding = 4 - len(s) % 4
    if padding != 4:
        s += "=" * padding
    return base64.urlsafe_b64decode(s)


class Keypair:
    """Ed25519 keypair for MRP identity."""

    def __init__(self, private_key: Ed25519PrivateKey):
        self._private_key = private_key
        self._public_key = private_key.public_key()
        # Cache raw bytes
        self._public_bytes = self._public_key.public_bytes(
            serialization.Encoding.Raw, serialization.PublicFormat.Raw
        )
        self._seed = private_key.private_bytes(
            serialization.Encoding.Raw,
            serialization.PrivateFormat.Raw,
            serialization.NoEncryption(),
        )

    @classmethod
    def generate(cls) -> Keypair:
        """Generate a new random Ed25519 keypair."""
        return cls(Ed25519PrivateKey.generate())

    @classmethod
    def from_seed(cls, seed: bytes) -> Keypair:
        """Derive keypair from 32-byte seed."""
        return cls(Ed25519PrivateKey.from_private_bytes(seed))

    @classmethod
    def from_file(cls, path: str, *, create: bool = False) -> Keypair:
        """Load keypair from file (32-byte raw seed).

        If create=True and the file doesn't exist, generate a new keypair and save it.
        """
        p = Path(path)
        if not p.exists():
            if not create:
                raise FileNotFoundError(f"Key file not found: {path}")
            kp = cls.generate()
            kp.save(path)
            return kp
        seed = p.read_bytes()
        if len(seed) != 32:
            raise ValueError(f"Invalid key file: expected 32 bytes, got {len(seed)}")
        return cls.from_seed(seed)

    def save(self, path: str) -> None:
        """Write seed (32 bytes) to file with 0600 permissions."""
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_bytes(self._seed)
        os.chmod(path, 0o600)

    @property
    def public_key(self) -> bytes:
        """Raw 32-byte public key."""
        return self._public_bytes

    @property
    def public_key_b64(self) -> str:
        """Base64url-encoded public key (43 chars, no padding)."""
        return _b64url_encode(self._public_bytes)

    @property
    def seed(self) -> bytes:
        """Raw 32-byte seed."""
        return self._seed

    def sign(self, data: bytes) -> bytes:
        """Ed25519 sign data, return raw 64-byte signature."""
        return self._private_key.sign(data)

    def sign_request(self, method: str, path: str, body: bytes | None, timestamp: str) -> str:
        """Build canonical string, sign, return base64url signature.

        Canonical string: METHOD\\nPATH\\nTIMESTAMP\\nBASE64URL(SHA256(BODY))
        """
        if body is None:
            body = b""
        body_hash = _b64url_encode(hashlib.sha256(body).digest())
        canonical = f"{method}\n{path}\n{timestamp}\n{body_hash}"
        sig = self.sign(canonical.encode())
        return _b64url_encode(sig)
