"""Snowflake utility module."""

import os

from abacus_common_logic.constants.math import Bytes
from cryptography.hazmat.primitives import serialization


def get_max_query_bytes() -> int:
    """Return the hard-coded query text size limit.

    Returns:
        int: The size in bytes.
    """
    return 16 * Bytes.MB  # type: ignore


def get_max_query_values() -> int:
    """Return the hard-coded query value limit.

    Returns:
        int: The size in bytes.
    """
    return 30_000


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