"""JWS signature creation and verification using ES256 and EdDSA."""

from __future__ import annotations

import base64
import json
from typing import TYPE_CHECKING

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.asymmetric.utils import (
    decode_dss_signature,
    encode_dss_signature,
)

if TYPE_CHECKING:
    from cryptography.hazmat.primitives.asymmetric.ec import (
        EllipticCurvePrivateKey,
        EllipticCurvePublicKey,
    )


def _b64url_encode(data: bytes) -> str:
    """Base64url encode without padding."""
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64url_decode(data: str) -> bytes:
    """Base64url decode with padding restoration."""
    # Add padding if needed
    padding = 4 - len(data) % 4
    if padding != 4:
        data += "=" * padding
    return base64.urlsafe_b64decode(data)


def create_jws_detached(
    payload: bytes,
    private_key: EllipticCurvePrivateKey | Ed25519PrivateKey,
    key_id: str,
) -> str:
    """
    Create a detached JWS signature over a payload.

    The resulting JWS uses compact serialization with an empty payload:
    base64url(header)..base64url(signature)

    Args:
        payload: The bytes to sign
        private_key: EC P-256 or Ed25519 private key for signing
        key_id: Key identifier to include in header

    Returns:
        Detached JWS compact serialization string
    """
    # Build header
    alg = "ES256"
    if isinstance(private_key, Ed25519PrivateKey):
        alg = "EdDSA"
        
    header = {
        "alg": alg,
        "kid": key_id,
    }
    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
    header_b64 = _b64url_encode(header_bytes)

    # Build signing input (header.payload)
    payload_b64 = _b64url_encode(payload)
    signing_input = f"{header_b64}.{payload_b64}".encode("ascii")

    # Sign with Ed25519
    signature_raw = None
    if isinstance(private_key, Ed25519PrivateKey):
        # raw 64-byte signature
        signature_raw = private_key.sign(signing_input)
    else:
        # Sign with ECDSA P-256
        signature_der = private_key.sign(signing_input, ec.ECDSA(hashes.SHA256()))
        # Convert DER signature to raw R||S format for JWS
        r, s = decode_dss_signature(signature_der)
        # Each component is 32 bytes for P-256
        r_bytes = r.to_bytes(32, byteorder="big")
        s_bytes = s.to_bytes(32, byteorder="big")
        signature_raw = r_bytes + s_bytes
    
    if signature_raw is None:
        raise ValueError("Failed to sign input")

    signature_b64 = _b64url_encode(signature_raw)

    # Return detached JWS (empty payload)
    return f"{header_b64}..{signature_b64}"


def verify_jws_detached(
    jws: str,
    payload: bytes,
    public_key: EllipticCurvePublicKey | Ed25519PublicKey,
) -> bool:
    try:
        parts = jws.split(".")
        if len(parts) != 3:
            return False

        header_b64 = parts[0]
        signature_b64 = parts[2] # parts[1] is empty in detached JWS

        # Decode header
        try:
            header_bytes = _b64url_decode(header_b64)
            header = json.loads(header_bytes.decode("utf-8"))
        except (ValueError, json.JSONDecodeError, UnicodeDecodeError):
            return False

        alg = header.get("alg")

        # Reconstruct signing input: header + "." + b64url(payload)
        payload_b64 = _b64url_encode(payload)
        signing_input = f"{header_b64}.{payload_b64}".encode("ascii")

        # Decode signature
        signature_raw = _b64url_decode(signature_b64)

        if alg == "EdDSA":
            if not isinstance(public_key, Ed25519PublicKey):
                return False
            if len(signature_raw) != 64:
                return False
            
            # Ed25519 verify
            try:
                public_key.verify(signature_raw, signing_input)
                return True
            except InvalidSignature:
                return False

        if alg == "ES256":
            if not isinstance(public_key, ec.EllipticCurvePublicKey):
                return False

            if len(signature_raw) == 64:
                # Raw JWS format: 32-byte R + 32-byte S -> convert to DER
                r = int.from_bytes(signature_raw[:32], byteorder="big")
                s = int.from_bytes(signature_raw[32:], byteorder="big")
                signature_der = encode_dss_signature(r, s)
            else:
                # Assume ASN.1 DER (e.g. WebAuthn/Passkey signatures)
                signature_der = signature_raw

            try:
                public_key.verify(signature_der, signing_input, ec.ECDSA(hashes.SHA256()))
                return True
            except InvalidSignature:
                return False

        # Unsupported algorithm
        return False

    except Exception:
        # Catch-all for unexpected encoding/library errors
        return False
