"""Ed25519 manifests and compact, trust-store-verified runtime tokens."""

from __future__ import annotations

import base64
import json
import os
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any

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

from .canonical import canonical_json, sha256_bytes
from .errors import SignatureError
from .models import SignatureBlock, SignedEnvelope

# Public protocol version tag, not a credential.
TOKEN_PREFIX = "sw1"  # nosec B105


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


def _b64decode(value: str) -> bytes:
    padding = "=" * (-len(value) % 4)
    try:
        return base64.b64decode(value + padding, altchars=b"-_", validate=True)
    except Exception as exc:
        raise SignatureError("invalid base64url value") from exc


def public_fingerprint(public_key: Ed25519PublicKey) -> str:
    raw = public_key.public_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PublicFormat.Raw,
    )
    return sha256_bytes(raw)


def load_private_key(path: str | Path) -> Ed25519PrivateKey:
    try:
        key = serialization.load_pem_private_key(Path(path).read_bytes(), password=None)
    except Exception as exc:
        raise SignatureError("private key is not readable") from exc
    if not isinstance(key, Ed25519PrivateKey):
        raise SignatureError("signing key must be Ed25519")
    return key


def load_public_key(path: str | Path) -> Ed25519PublicKey:
    try:
        key = serialization.load_pem_public_key(Path(path).read_bytes())
    except Exception as exc:
        raise SignatureError("public key is not readable") from exc
    if not isinstance(key, Ed25519PublicKey):
        raise SignatureError("trusted key must be Ed25519")
    return key


def generate_keypair(private_path: str | Path, public_path: str | Path) -> str:
    private = Ed25519PrivateKey.generate()
    private_bytes = private.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )
    public = private.public_key()
    public_bytes = public.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    private_target = Path(private_path)
    public_target = Path(public_path)
    private_target.parent.mkdir(parents=True, exist_ok=True)
    public_target.parent.mkdir(parents=True, exist_ok=True)
    private_target.write_bytes(private_bytes)
    try:
        os.chmod(private_target, 0o600)
    except OSError:
        pass
    public_target.write_bytes(public_bytes)
    return public_fingerprint(public)


@dataclass(frozen=True)
class TrustStore:
    keys: Mapping[str, Ed25519PublicKey]

    @classmethod
    def from_paths(cls, paths: Mapping[str, str | Path]) -> TrustStore:
        return cls({key_id: load_public_key(path) for key_id, path in paths.items()})

    def require(self, key_id: str) -> Ed25519PublicKey:
        key = self.keys.get(key_id)
        if key is None:
            raise SignatureError(f"untrusted signer key id: {key_id}")
        return key


def sign_envelope(
    payload: Mapping[str, Any],
    private_key: Ed25519PrivateKey,
    key_id: str,
) -> SignedEnvelope:
    payload_dict = dict(payload)
    body = canonical_json(payload_dict)
    signature = private_key.sign(body)
    return SignedEnvelope(
        payload=payload_dict,
        signature=SignatureBlock(
            key_id=key_id,
            signer_fingerprint=public_fingerprint(private_key.public_key()),
            payload_sha256=sha256_bytes(body),
            signature_b64=_b64encode(signature),
        ),
    )


def verify_envelope(
    envelope: SignedEnvelope | Mapping[str, Any],
    trust: TrustStore,
) -> dict[str, Any]:
    parsed = (
        envelope
        if isinstance(envelope, SignedEnvelope)
        else SignedEnvelope.model_validate(envelope)
    )
    public = trust.require(parsed.signature.key_id)
    expected_fingerprint = public_fingerprint(public)
    if parsed.signature.signer_fingerprint != expected_fingerprint:
        raise SignatureError("signer fingerprint does not match trusted key")
    body = canonical_json(parsed.payload)
    if sha256_bytes(body) != parsed.signature.payload_sha256:
        raise SignatureError("signed payload digest mismatch")
    try:
        signature = _b64decode(parsed.signature.signature_b64)
        if len(signature) != 64:
            raise SignatureError("invalid Ed25519 signature length")
        public.verify(signature, body)
    except InvalidSignature as exc:
        raise SignatureError("invalid Ed25519 signature") from exc
    return dict(parsed.payload)


def encode_token(
    claims: Mapping[str, Any],
    private_key: Ed25519PrivateKey,
    key_id: str,
) -> str:
    header = {"alg": "EdDSA", "kid": key_id, "typ": "SWT"}
    header_segment = _b64encode(canonical_json(header))
    payload_segment = _b64encode(canonical_json(dict(claims)))
    signed = f"{TOKEN_PREFIX}.{header_segment}.{payload_segment}".encode("ascii")
    signature_segment = _b64encode(private_key.sign(signed))
    return f"{signed.decode('ascii')}.{signature_segment}"


def decode_token(
    token: str,
    trust: TrustStore,
    *,
    expected_kind: str | None = None,
) -> dict[str, Any]:
    if len(token) > 16 * 1024:
        raise SignatureError("ScopeWrit token exceeds the 16 KiB limit")
    parts = token.split(".")
    if len(parts) != 4 or parts[0] != TOKEN_PREFIX:
        raise SignatureError("invalid ScopeWrit token format")
    try:
        header = json.loads(_b64decode(parts[1]))
        claims = json.loads(_b64decode(parts[2]))
    except (json.JSONDecodeError, UnicodeDecodeError) as exc:
        raise SignatureError("invalid ScopeWrit token JSON") from exc
    if header.get("alg") != "EdDSA" or header.get("typ") != "SWT":
        raise SignatureError("unsupported ScopeWrit token header")
    key_id = header.get("kid")
    if not isinstance(key_id, str):
        raise SignatureError("ScopeWrit token is missing a key id")
    public = trust.require(key_id)
    signed = ".".join(parts[:3]).encode("ascii")
    try:
        signature = _b64decode(parts[3])
        if len(signature) != 64:
            raise SignatureError("invalid ScopeWrit token signature length")
        public.verify(signature, signed)
    except InvalidSignature as exc:
        raise SignatureError("invalid ScopeWrit token signature") from exc
    if not isinstance(claims, dict):
        raise SignatureError("ScopeWrit token claims must be an object")
    if expected_kind is not None and claims.get("kind") != expected_kind:
        raise SignatureError("ScopeWrit token kind mismatch")
    return claims
