from __future__ import annotations

import base64
import hashlib
import time
from dataclasses import dataclass, field
from pathlib import Path

import jwt
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey


@dataclass
class JWTGenerator:
    """Generate Snowflake JWT assertions from a key-pair file.

    This is used by the Direct Access JWT flow to mint a
    session-login token (`assertion`) which is then exchanged for an ingress token
    scoped to a service endpoint.
    """

    account: str
    user: str
    private_key_path: str
    private_key_passphrase: str | None = None
    lifetime_s: int = 60
    _private_key: RSAPrivateKey | None = field(default=None, init=False, repr=False)
    _fingerprint: str | None = field(default=None, init=False, repr=False)
    _key_file_sig: tuple[int, int] | None = field(default=None, init=False, repr=False)

    def _get_private_key_and_fingerprint(self) -> tuple[RSAPrivateKey, str]:
        key_path = Path(self.private_key_path)
        stat = key_path.stat()
        key_file_sig = (int(stat.st_mtime_ns), int(stat.st_size))
        if (
            self._private_key is not None
            and self._fingerprint is not None
            and self._key_file_sig == key_file_sig
        ):
            return self._private_key, self._fingerprint

        pem = key_path.read_bytes()
        pwd = self.private_key_passphrase.encode("utf-8") if self.private_key_passphrase else None
        private_key = serialization.load_pem_private_key(pem, password=pwd, backend=default_backend())
        if not isinstance(private_key, RSAPrivateKey):
            raise TypeError("jwt requires an RSA private key.")

        public_key_der = private_key.public_key().public_bytes(
            serialization.Encoding.DER,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        digest = hashlib.sha256(public_key_der).digest()
        fingerprint = "SHA256:" + base64.b64encode(digest).decode("utf-8")

        self._private_key = private_key
        self._fingerprint = fingerprint
        self._key_file_sig = key_file_sig
        return private_key, fingerprint

    def generate(self) -> str:
        now = int(time.time())
        exp = now + int(self.lifetime_s)
        account = str(self.account).upper()
        user = str(self.user).upper()
        fully_qualified_user = f"{account}.{user}"
        private_key, fingerprint = self._get_private_key_and_fingerprint()

        payload = {
            "iss": f"{fully_qualified_user}.{fingerprint}",
            "sub": fully_qualified_user,
            "iat": now,
            "exp": exp,
        }
        return jwt.encode(payload, private_key, algorithm="RS256")

