"""Ed25519 signing and JWT minting for SBN service-to-service auth."""

from __future__ import annotations

import base64
import json
import os
import uuid
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Mapping, Sequence

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


def _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _canonical_json(data: Mapping[str, Any]) -> bytes:
    return json.dumps(data, separators=(",", ":"), sort_keys=True).encode("utf-8")


def _now_utc() -> datetime:
    return datetime.now(timezone.utc)


@dataclass(slots=True)
class MintedToken:
    """A freshly minted bearer token and its expiry."""

    token: str
    expires_at: datetime

    @classmethod
    def from_ttl(cls, token: str, ttl_seconds: int) -> MintedToken:
        return cls(token=token, expires_at=_now_utc() + timedelta(seconds=ttl_seconds))

    @property
    def expired(self) -> bool:
        return _now_utc() >= self.expires_at


@dataclass(slots=True)
class SigningKey:
    """Ed25519 signing helper for JWT minting and payload signatures."""

    private_key: Ed25519PrivateKey
    issuer: str
    audience: str
    kid: str = "sbn-ed25519"

    @classmethod
    def from_pem(
        cls,
        source: str | bytes | os.PathLike[str],
        *,
        issuer: str,
        audience: str,
        kid: str = "sbn-ed25519",
    ) -> SigningKey:
        if isinstance(source, (str, os.PathLike)):
            text = str(source)
            path = Path(text)
            pem_data = path.read_bytes() if path.exists() else text.encode("utf-8")
        else:
            pem_data = bytes(source)
        key = serialization.load_pem_private_key(pem_data, password=None)
        if not isinstance(key, Ed25519PrivateKey):
            raise TypeError("Expected an Ed25519 private key")
        return cls(private_key=key, issuer=issuer, audience=audience, kid=kid)

    @classmethod
    def generate(
        cls,
        *,
        issuer: str,
        audience: str,
        kid: str = "sbn-ed25519",
    ) -> SigningKey:
        return cls(
            private_key=Ed25519PrivateKey.generate(),
            issuer=issuer,
            audience=audience,
            kid=kid,
        )

    def issue_token(
        self,
        *,
        subject: str,
        scopes: Sequence[str],
        cdna: str,
        ttl_seconds: int,
        tenant_id: str | None = None,
        role: str = "agent",
    ) -> MintedToken:
        if ttl_seconds <= 0:
            raise ValueError("ttl_seconds must be positive")

        issued_at = _now_utc()
        scope_list = sorted({s.strip() for s in scopes if s})
        header = {"alg": "EdDSA", "typ": "JWT", "kid": self.kid}
        payload: dict[str, Any] = {
            "iss": self.issuer,
            "aud": self.audience,
            "sub": subject,
            "role": role,
            "scope": scope_list,
            "cdna": cdna,
            "iat": int(issued_at.timestamp()),
            "exp": int((issued_at + timedelta(seconds=ttl_seconds)).timestamp()),
            "jti": uuid.uuid4().hex,
        }
        if tenant_id:
            payload["tenant_id"] = tenant_id

        signing_input = f"{_b64url(_canonical_json(header))}.{_b64url(_canonical_json(payload))}"
        signature = self.private_key.sign(signing_input.encode("utf-8"))
        token = f"{signing_input}.{_b64url(signature)}"
        return MintedToken.from_ttl(token, ttl_seconds)

    def sign_payload(self, data: Mapping[str, Any]) -> str:
        """Sign arbitrary canonical JSON and return the base64url signature."""
        canonical = _canonical_json(data)
        signature = self.private_key.sign(canonical)
        return _b64url(signature)

    def private_key_pem(self) -> bytes:
        return self.private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
