from dataclasses import dataclass
from typing import Optional, Dict


@dataclass(frozen=True)
class JwtPublicKey:
    kid: str
    public_key_pem: str


@dataclass(frozen=True)
class JwtPublicKeySet:
    active: JwtPublicKey
    previous: Optional[JwtPublicKey] = None

    def public_keys_by_kid(self) -> Dict[str, str]:
        out = {self.active.kid: self.active.public_key_pem}
        if self.previous:
            out[self.previous.kid] = self.previous.public_key_pem
        return out