"""Authentication utilities."""

import base64
import os
from functools import lru_cache
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization

from common.src import logger

logger = logger.new_logger(__name__)


@lru_cache(maxsize=8)
def get_private_key_der(string: str, password: str = None) -> bytes:
    """
    Loads a private key and returns it in DER-encoded PKCS#8 format.

    The input string can be:
    - A PEM-formatted private key string (with BEGIN/END headers)
    - A file path to a PEM or DER private key
    - A base64-encoded PEM (headers removed) or DER key

    Args:
        string (str): The private key input as a string or file path.
        password (str, optional): The password for the encrypted private key, if required.

    Returns:
        bytes: The private key encoded in DER format.

    Raises:
        ValueError: If the key cannot be parsed in either PEM or DER format.
    """
    private_key_content: bytes = _get_private_key_content(string)

    # Try PEM first
    try:
        p_key = serialization.load_pem_private_key(
            private_key_content,
            password=password.encode() if password else None,
            backend=default_backend(),
        )
        logger.debug("Private key loaded successfully from PEM format.")
    except ValueError:
        # If PEM fails, try DER
        p_key = serialization.load_der_private_key(
            private_key_content,
            password=password.encode() if password else None,
            backend=default_backend(),
        )
        logger.debug("Private key loaded successfully from DER format.")

    logger.debug("Converting private key to DER format...")
    return p_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )


def _get_private_key_content(provided_key: str) -> bytes:
    """
    Returns private key content in binary format.

    Accepts:
    - A PEM-formatted key string (with BEGIN/END headers)
    - A path to a PEM or DER key file
    - A base64-encoded PEM (headers removed) or DER string

    Args:
        provided_key (str): A private key string or a file path.

    Returns:
        bytes: Raw key content (decoded from base64 or read from file)

    Raises:
        ValueError: If the base64 decoding fails or format is invalid.
    """
    provided_key = provided_key.strip()

    if provided_key.startswith("-----BEGIN"):
        logger.debug("The provided private key seems to be in PEM format.")
        return provided_key.encode("utf-8")

    if os.path.exists(provided_key):
        logger.debug("The provided private key seems to be a file path.")
        with open(provided_key, "rb") as f:
            return f.read()

    # Handle base64-encoded PEM or DER
    logger.debug("The provided private key seems to be base64-encoded.")
    cleaned = provided_key.replace("\n", "").strip()
    padded = cleaned + "=" * (-len(cleaned) % 4)
    return base64.b64decode(padded)