from __future__ import annotations

import dataclasses
import functools
import itertools
import logging
from collections import defaultdict
from copy import deepcopy
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, TypeVar

from auditwheel import json
from auditwheel.elfutils import (
    elf_file_filter,
    elf_find_ucs2_symbols,
    elf_find_versioned_symbols,
    elf_is_python_extension,
    elf_references_pyfpe_jbuf,
)
from auditwheel.error import InvalidLibcError, NonPlatformWheelError
from auditwheel.genericpkgctx import InGenericPkgCtx
from auditwheel.lddtree import DynamicExecutable, DynamicLibrary, ld_paths_from_arg, ldd
from auditwheel.libc import Libc
from auditwheel.policy import ExternalReference, Policy, WheelPolicies

if TYPE_CHECKING:
    from collections.abc import Mapping

    from auditwheel.architecture import Architecture

log = logging.getLogger(__name__)


@dataclass(frozen=True)
class WheelAbIInfo:
    policies: WheelPolicies
    full_external_refs: dict[Path, dict[str, ExternalReference]]
    overall_policy: Policy
    external_refs: dict[str, ExternalReference]
    ref_policy: Policy
    versioned_symbols: dict[str, set[str]]
    sym_policy: Policy
    ucs_policy: Policy
    pyfpe_policy: Policy
    blacklist_policy: Policy
    machine_policy: Policy
    graft_policy: Policy


@dataclass(frozen=True)
class WheelElfData:
    policies: WheelPolicies
    full_elftree: dict[Path, DynamicExecutable]
    full_external_refs: dict[Path, dict[str, ExternalReference]]
    versioned_symbols: dict[str, set[str]]
    uses_ucs2_symbols: bool
    uses_pyfpe_jbuf: bool


def _fixup_elf_trees(
    py_elf_trees: dict[Path, DynamicExecutable],
    nonpy_elftree: dict[Path, DynamicExecutable],
) -> None:
    # Auditwheel assumes that if a dependency in a non-Python ELF tree can be resolved
    # in any other fully resolved ELF tree, then it can be loaded from other
    # non-Python ELF trees as well. Python extension ELF trees are not
    # fixed up; they should always be valid in the first place. We fix up
    # non-Python ELF trees with that assumption in mind.
    resolved_executables = [
        executable
        for executable in itertools.chain(py_elf_trees.values(), nonpy_elftree.values())
        if all(library.realpath is not None for library in executable.libraries.values())
    ]

    def _get_reference(path_to_elf: Path) -> DynamicExecutable | None:
        assert path_to_elf.is_file()  # noqa: S101
        for candidate in resolved_executables:
            for library in candidate.libraries.values():
                assert library.realpath is not None  # noqa: S101
                assert library.realpath.is_file()  # noqa: S101
                if path_to_elf.samefile(library.realpath):
                    return candidate
        return None

    for elf_path, elf_executable in nonpy_elftree.items():
        need_resolution = any(
            elf_library.realpath is None for elf_library in elf_executable.libraries.values()
        )
        if not need_resolution:
            continue
        reference = _get_reference(elf_executable.realpath)
        if reference is None:
            continue
        ref_name = reference.realpath.name
        libraries_new: dict[str, DynamicLibrary] = {}
        soname_all = set(elf_executable.libraries.keys())
        for soname, elf_library in elf_executable.libraries.items():
            if elf_library.realpath is None and soname in reference.libraries:
                libraries_new[soname] = reference.libraries[soname]
                # we need to add newly found dependencies as well
                needed_new = set(libraries_new[soname].needed) - soname_all
                while needed_new:
                    needed_new_copy = needed_new.copy()
                    soname_all |= needed_new_copy
                    for soname_needed in needed_new_copy:
                        if soname_needed in reference.libraries:
                            libraries_new[soname_needed] = reference.libraries[soname_needed]
                            needed_new.update(libraries_new[soname_needed].needed)
                        else:
                            log.warning("%s not found in %s ELF tree", soname_needed, ref_name)
                    needed_new -= soname_all
            else:
                if elf_library.realpath is None:
                    log.warning("%s not found in %s ELF tree", soname, ref_name)
                libraries_new[soname] = elf_library
        nonpy_elftree[elf_path] = dataclasses.replace(elf_executable, libraries=libraries_new)


@functools.lru_cache
def get_wheel_elfdata(
    libc: Libc | None,
    architecture: Architecture | None,
    wheel_fn: Path,
    exclude: frozenset[str],
    args_ldpaths: str | None,
) -> WheelElfData:
    full_elftree: dict[Path, DynamicExecutable] = {}
    nonpy_elftree: dict[Path, DynamicExecutable] = {}
    full_external_refs: dict[Path, dict[str, ExternalReference]] = {}
    versioned_symbols: dict[str, set[str]] = defaultdict(set)
    uses_ucs2_symbols = False
    uses_pyfpe_jbuf = False
    policies: WheelPolicies | None = None

    # Android is cross-compiled, so ldpaths should never be loaded from the build machine.
    if libc == Libc.ANDROID and args_ldpaths is None:
        args_ldpaths = ""

    with InGenericPkgCtx(wheel_fn) as ctx:
        shared_libraries_in_purelib = []
        shared_libraries_with_invalid_machine = []

        platform_wheel = False
        for fn, elf in elf_file_filter(ctx.iter_files()):
            # Check for invalid binary wheel format: no shared library should
            # be found in purelib
            so_name = fn.name

            # If this is in purelib, add it to the list of shared libraries in
            # purelib
            if any(p.name == "purelib" for p in fn.parents):
                shared_libraries_in_purelib.append(so_name)

            # If at least one shared library exists in purelib, this is going
            # to fail and there's no need to do further checks
            if not shared_libraries_in_purelib:
                log.debug("processing: %s", fn)
                elftree = ldd(fn, exclude=exclude, ldpaths=ld_paths_from_arg(args_ldpaths))

                try:
                    elf_arch = elftree.platform.baseline_architecture
                except ValueError:
                    shared_libraries_with_invalid_machine.append(so_name)
                    log.warning("ignoring: %s with unknown architecture", so_name)
                    continue
                if architecture is None:
                    log.info("setting architecture to %s", elf_arch.value)
                    architecture = elf_arch
                elif elf_arch != architecture.baseline:
                    shared_libraries_with_invalid_machine.append(so_name)
                    log.warning("ignoring: %s with %s architecture", so_name, elf_arch)
                    continue

                if elftree.libc is not None:
                    if libc is None:
                        log.info("setting libc to %s", elftree.libc)
                        libc = elftree.libc
                    elif libc != elftree.libc:
                        log.warning("ignoring: %s with %s libc", so_name, elftree.libc)
                        continue

                if policies is None and libc is not None and architecture is not None:
                    policies = WheelPolicies(libc=libc, arch=architecture, wheel_fn=wheel_fn)

                platform_wheel = True

                for key, value in elf_find_versioned_symbols(elf):
                    log.debug("key %s, value %s", key, value)
                    versioned_symbols[key].add(value)

                is_py_ext, py_ver = elf_is_python_extension(fn, elf)

                # If the ELF is a Python extention, we definitely need to
                # include its external dependencies.
                if is_py_ext:
                    if policies is None:
                        assert architecture is not None  # noqa: S101
                        assert libc is None  # noqa: S101
                        msg = f"couldn't detect libc for python extension {fn}"
                        raise InvalidLibcError(msg)
                    full_elftree[fn] = elftree
                    uses_pyfpe_jbuf |= elf_references_pyfpe_jbuf(elf)
                    if py_ver == 2:
                        uses_ucs2_symbols |= any(True for _ in elf_find_ucs2_symbols(elf))
                    full_external_refs[fn] = policies.lddtree_external_references(
                        elftree,
                        ctx.path,
                    )
                else:
                    # If the ELF is not a Python extension, it might be
                    # included in the wheel already because auditwheel repair
                    # vendored it, so we will check whether we should include
                    # its internal references later.
                    nonpy_elftree[fn] = elftree

        # If at least one shared library exists in purelib, raise an error
        if shared_libraries_in_purelib:
            libraries = "\n\t".join(shared_libraries_in_purelib)
            msg = (
                "Invalid binary wheel, found the following shared library/libraries "
                f"in purelib folder:\n\t{libraries}\n"
                "The wheel has to be platlib compliant in order to be repaired by "
                "auditwheel."
            )
            raise RuntimeError(msg)

        if not platform_wheel:
            arch = None if architecture is None else architecture.value
            raise NonPlatformWheelError(arch, shared_libraries_with_invalid_machine)

        _fixup_elf_trees(full_elftree, nonpy_elftree)

        # Get a list of all external libraries needed by ELFs in the wheel.
        needed_libs = {
            lib
            for elf in itertools.chain(full_elftree.values(), nonpy_elftree.values())
            for lib in elf.needed
        }

        if policies is None:
            # we have no python extensions, either we have shared libraries with
            # no dependencies on libc (unlikely) or a statically linked executable
            # let's fallback to the host libc
            assert architecture is not None  # noqa: S101
            assert libc is None  # noqa: S101
            libc = Libc.detect()
            log.warning("couldn't detect wheel libc, defaulting to %s", str(libc))
            policies = WheelPolicies(libc=libc, arch=architecture)

        for fn, elf_tree in nonpy_elftree.items():
            # If a non-pyextension ELF file is not needed by something else
            # inside the wheel, then it was not checked by the logic above and
            # we should walk its elftree.
            if fn.name not in needed_libs:
                full_elftree[fn] = elf_tree

            # Even if a non-pyextension ELF file is not needed, we
            # should include it as an external reference, because
            # it might require additional external libraries.
            full_external_refs[fn] = policies.lddtree_external_references(
                elf_tree,
                ctx.path,
            )

    log.debug("full_elftree:\n%s", json.dumps(full_elftree))
    log.debug(
        "full_external_refs (will be repaired):\n%s",
        json.dumps(full_external_refs),
    )

    return WheelElfData(
        policies,
        full_elftree,
        full_external_refs,
        versioned_symbols,
        uses_ucs2_symbols,
        uses_pyfpe_jbuf,
    )


def get_external_libs(external_refs: dict[str, ExternalReference]) -> dict[Path, str]:
    """Get external library dependencies for all policies excluding the default
    linux policy
    :param external_refs: external references for all policies
    :return: {realpath: soname} e.g.
    {'/path/to/external_ref.so.1.2.3': 'external_ref.so.1'}
    """
    result: dict[Path, str] = {}
    for external_ref in external_refs.values():
        # linux tag (priority 0) has no white-list, do not analyze it
        if external_ref.policy.priority == 0:
            continue
        # go through all libs, retrieving their soname and realpath
        for libname, realpath in external_ref.libs.items():
            if realpath and realpath not in result:
                result[Path(realpath)] = libname
    return result


def get_versioned_symbols(libs: dict[Path, str]) -> dict[str, dict[str, set[str]]]:
    """Get versioned symbols used in libraries
    :param libs: {realpath: soname} dict to search for versioned symbols e.g.
    {'/path/to/external_ref.so.1.2.3': 'external_ref.so.1'}
    :return: {soname: {depname: set([symbol_version])}} e.g.
    {'external_ref.so.1': {'libc.so.6', set(['GLIBC_2.5','GLIBC_2.12'])}}
    """
    result = {}
    for path, elf in elf_file_filter(libs.keys()):
        # {depname: set(symbol_version)}, e.g. {'libc.so.6', set(['GLIBC_2.5','GLIBC_2.12'])}
        elf_versioned_symbols: dict[str, set[str]] = defaultdict(set)
        for key, value in elf_find_versioned_symbols(elf):
            log.debug("path %s, key %s, value %s", path, key, value)
            elf_versioned_symbols[key].add(value)
        result[libs[path]] = elf_versioned_symbols
    return result


def get_symbol_policies(
    policies: WheelPolicies,
    versioned_symbols: dict[str, set[str]],
    external_versioned_symbols: dict[str, dict[str, set[str]]],
    external_refs: dict[str, ExternalReference],
) -> list[tuple[Policy, dict[str, set[str]]]]:
    """Get symbol policies
    Since white-list is different per policy, this function inspects
    versioned_symbol per policy when including external refs
    :param versioned_symbols: versioned symbols for the current wheel
    :param external_versioned_symbols: versioned symbols for external libs
    :param external_refs: external references for all policies
    :return: list of tuples of the form (policy, versioned_symbols),
    e.g. [(<Policy: manylinux...>, {'libc.so.6', set(['GLIBC_2.5'])})]
    """
    result = []
    for external_ref in external_refs.values():
        # skip the linux policy
        if external_ref.policy.priority == 0:
            continue
        policy_symbols = deepcopy(versioned_symbols)
        for soname in external_ref.libs:
            if soname not in external_versioned_symbols:
                continue
            ext_symbols = external_versioned_symbols[soname]
            for k in iter(ext_symbols):
                policy_symbols[k].update(ext_symbols[k])
        # if the white-list policy changed, we don't want to allow highest priority policy
        # than the current one, that is, only restrict to a lower priority policy
        found_policy = min(
            external_ref.policy,
            policies.versioned_symbols_policy(policy_symbols),
        )
        result.append((found_policy, policy_symbols))
    return result


def _get_machine_policy(
    policies: WheelPolicies,
    elftree_by_fn: dict[Path, DynamicExecutable],
    external_so_names: frozenset[str],
) -> Policy:
    """
    Determine the most appropriate machine policy for the wheel based on ELF
    architectures and whitelisted external dependencies.
    The function inspects the extended architecture of each top-level ELF file
    in ``elftree_by_fn`` and of any library dependencies whose SONAME appears in
    ``external_so_names``. If an ELF file requires instructions that are not
    covered by ``policies.architecture``, the overall policy is downgraded to
    reflect that requirement.
    If the offending ELF file corresponds to an external shared object
    (i.e. its SONAME is in ``external_so_names``), the function looks for a
    policy whose whitelist contains that SONAME and uses the highest-priority
    such policy instead of unconditionally falling back to the baseline
    ``policies.linux`` policy. Informational or warning logs are emitted to
    record any such downgrades.
    :param policies: Collection of available wheel policies, including the
        target architecture, the baseline ``linux`` policy, and the highest
        (most permissive) default policy.
    :param elftree_by_fn: Mapping from ELF file paths to their corresponding
        :class:`DynamicExecutable` objects, used to inspect platforms and
        dependency graphs.
    :param external_so_names: Set of SONAMEs that are treated as external
        references and may be whitelisted by individual policies.
    :return: The selected :class:`Policy` after considering all relevant ELF
        files and their dependencies.
    """
    result = policies.highest
    machine_to_check: dict[Path, tuple[str | None, Architecture | None]] = {}
    for fn, dynamic_executable in elftree_by_fn.items():
        if fn in machine_to_check:
            continue
        machine_to_check[fn] = (None, dynamic_executable.platform.extended_architecture)
        for dependency in dynamic_executable.libraries.values():
            if dependency.soname not in external_so_names:
                continue
            if dependency.realpath is None:
                continue
            assert dependency.platform is not None  # noqa: S101
            if dependency.realpath in machine_to_check:
                continue
            machine_to_check[dependency.realpath] = (
                dependency.soname,
                dependency.platform.extended_architecture,
            )

    for fn, (soname, extended_architecture) in machine_to_check.items():
        if extended_architecture is None:
            continue
        if policies.architecture.is_superset(extended_architecture):
            continue
        if soname is not None:
            found_policy = policies.linux
            for policy in policies:
                if soname in policy.whitelist:
                    found_policy = max(found_policy, policy)
            if policies.linux.priority < found_policy.priority:
                log.info(
                    "ELF file %r requires %r instruction set, not in %r, whitelisted in %r",
                    fn,
                    extended_architecture.value,
                    policies.architecture.value,
                    found_policy.name,
                )
                result = min(result, found_policy)
                continue
        log.warning(
            "ELF file %r requires %r instruction set, not in %r",
            fn,
            extended_architecture.value,
            policies.architecture.value,
        )
        result = policies.linux

    return result


def analyze_wheel_abi(
    libc: Libc | None,
    architecture: Architecture | None,
    wheel_fn: Path,
    exclude: frozenset[str],
    *,
    disable_isa_ext_check: bool,
    allow_graft: bool,
    requested_policy_base_name: str | None = None,
    args_ldpaths: str | None = None,
) -> WheelAbIInfo:
    data = get_wheel_elfdata(libc, architecture, wheel_fn, exclude, args_ldpaths)
    policies = data.policies
    elftree_by_fn = data.full_elftree
    external_refs_by_fn = data.full_external_refs
    versioned_symbols = data.versioned_symbols

    external_refs: dict[str, ExternalReference] = {
        p.name: ExternalReference({}, {}, p) for p in policies
    }

    for fn in elftree_by_fn:
        update(external_refs, external_refs_by_fn[fn])

    log.debug("external reference info")
    log.debug(json.dumps(external_refs))

    external_libs = get_external_libs(external_refs)
    external_versioned_symbols = get_versioned_symbols(external_libs)
    symbol_policies = get_symbol_policies(
        policies,
        versioned_symbols,
        external_versioned_symbols,
        external_refs,
    )
    symbol_policy = policies.versioned_symbols_policy(versioned_symbols)

    # let's keep the highest priority policy and
    # corresponding versioned_symbols
    symbol_policy, versioned_symbols = max(
        symbol_policies,
        key=lambda x: x[0],
        default=(symbol_policy, versioned_symbols),
    )

    ref_policy = max(
        (e.policy for e in external_refs.values() if len(e.libs) == 0),
        default=policies.linux,
    )

    blacklist_policy = max(
        (e.policy for e in external_refs.values() if len(e.blacklist) == 0),
        default=policies.linux,
    )

    # don't allow highest priority policies with more grafted libraries
    # than the requested policy
    if requested_policy_base_name is None or requested_policy_base_name == "auto":
        graft_lib_count = min(
            (len(e.libs) for e in external_refs.values() if e.policy != policies.linux),
            default=0,
        )
    else:
        graft_lib_count = min(
            (
                len(e.libs)
                for e in external_refs.values()
                if e.policy.name.startswith(f"{requested_policy_base_name}_")
            ),
            default=0,
        )
    graft_policy = max(
        (e.policy for e in external_refs.values() if len(e.libs) <= graft_lib_count),
        default=policies.linux,
    )

    if disable_isa_ext_check:
        machine_policy = policies.highest
    else:
        machine_policy = _get_machine_policy(
            policies,
            elftree_by_fn,
            frozenset(external_libs.values()),
        )

    ucs_policy = policies.linux if data.uses_ucs2_symbols else policies.highest
    pyfpe_policy = policies.linux if data.uses_pyfpe_jbuf else policies.highest

    overall_policy = min(
        symbol_policy,
        ucs_policy,
        pyfpe_policy,
        blacklist_policy,
        machine_policy,
    )

    if requested_policy_base_name is not None:
        overall_policy = min(overall_policy, graft_policy)

    if not allow_graft:
        overall_policy = min(overall_policy, ref_policy)

    return WheelAbIInfo(
        policies,
        external_refs_by_fn,
        overall_policy,
        external_refs,
        ref_policy,
        versioned_symbols,
        symbol_policy,
        ucs_policy,
        pyfpe_policy,
        blacklist_policy,
        machine_policy,
        graft_policy,
    )


_T = TypeVar("_T", ExternalReference, Path | None)


def update(d: dict[str, _T], u: Mapping[str, _T]) -> None:
    for k, v in u.items():
        if isinstance(v, ExternalReference):
            assert k in d  # noqa: S101
            assert isinstance(d[k], ExternalReference)  # noqa: S101
            assert d[k].policy == v.policy  # noqa: S101
            # update blacklist
            for lib, symbols in v.blacklist.items():
                if lib not in d[k].blacklist:
                    d[k].blacklist[lib] = sorted(symbols)
                else:
                    d[k].blacklist[lib] = sorted(
                        set(d[k].blacklist[lib]) | set(symbols),
                    )
            # libs
            update(d[k].libs, v.libs)
        elif isinstance(v, (Path, type(None))):
            d[k] = v
        else:
            msg = f"can't update {d} with {k}:{v}"
            raise TypeError(msg)
