"""Snowflake utility module."""

from __future__ import annotations

import os

from cryptography.hazmat.primitives import serialization


def get_private_key(
    private_key: str | None = None,
    private_key_path: str | None = None,
    private_key_passphrase: str | None = None,
) -> bytes | None:
    """Get private key bytes from either string content or file path.

    Attempts to load private key from string content first, then falls back
    to file path if provided. Returns None if neither is available.

    Args:
        private_key: Private key content as a string.
        private_key_path: Path to private key file.
        private_key_passphrase: Optional passphrase for encrypted keys.

    Returns:
        Private key in DER format, or None if no key source provided.
    """
    if private_key:
        return get_private_key_from_str(private_key, private_key_passphrase)

    if private_key_path:
        return get_private_key_from_file(private_key_path, private_key_passphrase)

    return None


def get_private_key_from_file(key_path: str, passphrase: str | None = None) -> bytes:
    """Parse private key file for Snowflake authentication.

    Args:
        key_path: Path to the private key file (supports tilde expansion).
        passphrase: Optional passphrase for encrypted keys.

    Returns:
        The private key in DER format.
    """
    expanded_path = os.path.expanduser(key_path)
    with open(expanded_path, 'rb') as key:
        password = _encode_passphrase(passphrase)
        return _convert_pem_to_der(key.read(), password)


def get_private_key_from_str(key: str, passphrase: str | None = None) -> bytes:
    """Parse private key for Snowflake authentication.

    Args:
        key: The private key content as a string.
        passphrase: Optional passphrase for encrypted keys.

    Returns:
        The private key in DER format.
    """
    password = _encode_passphrase(passphrase)
    return _convert_pem_to_der(key.encode('utf-8'), password)


def _convert_pem_to_der(key: bytes, password: bytes | None) -> bytes:
    p_key = serialization.load_pem_private_key(key, password=password)

    return p_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )


def _encode_passphrase(passphrase: str | None) -> bytes | None:
    return passphrase.encode('utf-8') if passphrase else None