from uuid import uuid4

from pydantic import BaseModel, Field, model_validator


class AccessToken(BaseModel):
    """
    Pydantic model representing a standard OAuth2 Access Token.

    Attributes:
        iss (str): Issuer identifier for the token.
        aud (str): Audience that the token is intended for.
        exp (int): Expiration time (as UNIX timestamp).
        iat (int): Issued at time (as UNIX timestamp).
        client_id (str): Identifier for the client.
        jti (str): Unique identifier for the token (JWT ID). Generated by default.
        sub (str): Subject of the token (usually equals client_id).
        cnf (dict): Confirmation claim (e.g., public key thumbprint).
    """

    iss: str
    aud: str
    exp: int
    iat: int
    client_id: str
    jti: str = Field(default_factory=lambda: str(uuid4()))
    sub: str
    cnf: dict = {}


class RefreshToken(AccessToken):
    """
    Pydantic model representing a Refresh Token, extending the Access Token.

    Adds:
        nbf (int): 'Not Before' claim as a UNIX timestamp. It defaults to the `exp` of the token.

    Notes:
        The `nbf` field ensures the refresh token is not accepted before the expiration of the access token.
    """

    nbf: int = None

    @model_validator(mode="after")
    def set_nbf(self) -> "RefreshToken":
        """
        Validator that sets `nbf` to `exp` if not explicitly provided.
        """
        self.nbf = self.nbf or self.exp
        return self
