# SPDX-FileCopyrightText: 2018 Helder Eijs <helderijs@gmail.com>
# SPDX-License-Identifier: BSD-2-Clause

from __future__ import annotations

import os
from typing import Callable

from Crypto.Util._raw_api import load_pycryptodome_raw_lib

_raw_cpuid_lib = load_pycryptodome_raw_lib(
    "Crypto.Util._cpuid_c",
    """
                                           int have_aes_ni(void);
                                           int have_clmul(void);
                                           int have_avx2(void);
                                           int have_bmi1(void);
                                           int have_bmi2(void);
                                           int have_sha_ni(void);
                                           int have_adx(void);
                                           int have_aes_armv8(void);
                                           int have_pmull(void);
                                           int have_sha256_armv8(void);
                                           int have_sha512_armv8(void);
                                           int have_sha3_armv8(void);
                                           """,
)


def have_aes_ni() -> int:
    return _raw_cpuid_lib.have_aes_ni()


def have_clmul() -> int:
    return _raw_cpuid_lib.have_clmul()


def have_avx2() -> int:
    return _raw_cpuid_lib.have_avx2()


def have_bmi1() -> int:
    return _raw_cpuid_lib.have_bmi1()


def have_bmi2() -> int:
    return _raw_cpuid_lib.have_bmi2()


def have_sha_ni() -> int:
    return _raw_cpuid_lib.have_sha_ni()


def have_adx() -> int:
    return _raw_cpuid_lib.have_adx()


def have_aes_armv8() -> int:
    return _raw_cpuid_lib.have_aes_armv8()


def have_pmull() -> int:
    return _raw_cpuid_lib.have_pmull()


def have_sha256_armv8() -> int:
    return _raw_cpuid_lib.have_sha256_armv8()


def have_sha512_armv8() -> int:
    return _raw_cpuid_lib.have_sha512_armv8()


def have_sha3_armv8() -> int:
    return _raw_cpuid_lib.have_sha3_armv8()


def load_optional_lib(name: str, cdecl: str, cpu_supported: Callable[[], int]):
    """Load the functions in ``cdecl`` from the module ``name``,
    compiled for instructions that the CPU may not have.

    Return None if ``cpu_supported()`` is false (the CPU does not have them),
    or if the module was not compiled in.
    """

    if not cpu_supported():
        return None
    try:
        return load_pycryptodome_raw_lib(name, cdecl)
    except OSError:
        return None


def load_avx2_bmi2_lib(name: str, cdecl: str):
    """Load the functions in ``cdecl`` from the module ``name``,
    compiled for AVX2, BMI1 and BMI2.

    Return None if the CPU does not support all three instruction sets,
    or if the module was not compiled in.
    """

    if not (have_avx2() and have_bmi1() and have_bmi2()):
        return None
    try:
        return load_pycryptodome_raw_lib(name, cdecl)
    except OSError:
        return None


def available_cores() -> int:
    """Return the number of CPU cores this process can run on."""

    count: int | None
    # Python 3.13+: it takes into account CPU affinity and -X cpu_count
    if hasattr(os, "process_cpu_count"):
        count = os.process_cpu_count()
    elif hasattr(os, "sched_getaffinity"):
        count = len(os.sched_getaffinity(0))
    elif hasattr(os, "cpu_count"):
        count = os.cpu_count()
    else:
        import multiprocessing

        try:
            count = multiprocessing.cpu_count()
        except NotImplementedError:
            count = None
    return count or 1