from datetime import UTC, datetime, timedelta
import base64
import json
import secrets
from typing import Any

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec, padding, rsa
from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature, encode_dss_signature
from cryptography.hazmat.primitives.hashes import SHA256
import jwt


def now_utc() -> datetime:
    return datetime.now(UTC)


def random_token(bytes_len: int = 32) -> str:
    return secrets.token_urlsafe(bytes_len)


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


def b64url_decode(value: str) -> bytes:
    pad = "=" * (-len(value) % 4)
    return base64.urlsafe_b64decode(value + pad)


def canonical_json(payload: dict[str, Any]) -> bytes:
    return json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def load_public_key(public_key_pem: str):
    return serialization.load_pem_public_key(public_key_pem.encode("utf-8"))


def verify_signature(public_key_pem: str, payload: dict[str, Any], signature_b64: str) -> bool:
    key = load_public_key(public_key_pem)
    signature = b64url_decode(signature_b64)
    message = canonical_json(payload)

    try:
        if isinstance(key, ec.EllipticCurvePublicKey):
            if len(signature) == 64:
                signature = encode_dss_signature(
                    int.from_bytes(signature[:32], "big"),
                    int.from_bytes(signature[32:], "big"),
                )
            key.verify(signature, message, ec.ECDSA(SHA256()))
            return True
        if isinstance(key, rsa.RSAPublicKey):
            key.verify(signature, message, padding.PKCS1v15(), SHA256())
            return True
    except Exception:
        return False
    return False


def generate_oidc_private_key() -> rsa.RSAPrivateKey:
    return rsa.generate_private_key(public_exponent=65537, key_size=2048)


def serialize_private_key(key: rsa.RSAPrivateKey) -> str:
    return key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        serialization.NoEncryption(),
    ).decode("utf-8")


def load_or_create_oidc_key(key_pem: str | None) -> rsa.RSAPrivateKey:
    if key_pem:
        return serialization.load_pem_private_key(key_pem.encode("utf-8"), password=None)
    return generate_oidc_private_key()


def public_jwk(private_key: rsa.RSAPrivateKey, kid: str) -> dict[str, str]:
    numbers = private_key.public_key().public_numbers()
    return {
        "kty": "RSA",
        "use": "sig",
        "kid": kid,
        "alg": "RS256",
        "n": b64url(numbers.n.to_bytes((numbers.n.bit_length() + 7) // 8, "big")),
        "e": b64url(numbers.e.to_bytes((numbers.e.bit_length() + 7) // 8, "big")),
    }


def sign_jwt(claims: dict[str, Any], private_key: rsa.RSAPrivateKey, kid: str) -> str:
    claims = claims.copy()
    claims.setdefault("iat", int(now_utc().timestamp()))
    return jwt.encode(claims, private_key, algorithm="RS256", headers={"kid": kid})


def expires_in(seconds: int) -> datetime:
    return now_utc() + timedelta(seconds=seconds)