from __future__ import annotations

import json
import logging
import re
from collections import defaultdict
from dataclasses import dataclass, replace
from pathlib import Path
from typing import TYPE_CHECKING, Any

from auditwheel.architecture import Architecture
from auditwheel.elfutils import filter_undefined_symbols
from auditwheel.error import InvalidLibcError
from auditwheel.lddtree import LIBPYTHON_RE
from auditwheel.libc import Libc
from auditwheel.tools import is_subdir
from auditwheel.wheeltools import android_api_level, get_wheel_platforms

if TYPE_CHECKING:
    from collections.abc import Generator, Iterable

    from auditwheel.lddtree import DynamicExecutable

_HERE = Path(__file__).parent
_MUSL_POLICY_RE = re.compile(r"^musllinux_\d+_\d+$")

logger = logging.getLogger(__name__)

_POLICY_JSON_MAP = {
    Libc.GLIBC: _HERE / "manylinux-policy.json",
    Libc.MUSL: _HERE / "musllinux-policy.json",
    Libc.ANDROID: _HERE / "android-policy.json",
}


@dataclass(frozen=True)
class Policy:
    name: str
    aliases: tuple[str, ...]
    priority: int
    symbol_versions: dict[str, frozenset[str]]
    whitelist: frozenset[str]
    blacklist: dict[str, frozenset[str]]

    def __lt__(self, other: Any) -> bool:  # noqa: ANN401
        if not isinstance(other, Policy):
            raise NotImplementedError
        return self.priority < other.priority


@dataclass()
class ExternalReference:
    libs: dict[str, Path | None]
    blacklist: dict[str, list[str]]
    policy: Policy


class WheelPolicies:
    def __init__(
        self,
        *,
        libc: Libc,
        arch: Architecture,
        wheel_fn: Path | None = None,
        musl_policy: str | None = None,
    ) -> None:
        if libc != Libc.MUSL and musl_policy is not None:
            msg = f"'musl_policy' shall be None for libc {libc.name}"
            raise ValueError(msg)
        if libc == Libc.MUSL:
            if musl_policy is None:
                try:
                    musl_version = libc.get_current_version()
                    musl_policy = f"musllinux_{musl_version.major}_{musl_version.minor}"
                except InvalidLibcError:
                    logger.warning(
                        "can't determine musl libc version, latest known version will be used.",
                    )
            elif _MUSL_POLICY_RE.match(musl_policy) is None:
                msg = f"Invalid 'musl_policy': '{musl_policy}'"
                raise ValueError(msg)
        policies = json.loads(_POLICY_JSON_MAP[libc].read_text())
        self._policies: list[Policy] = []
        self._architecture = arch
        self._libc_variant = libc
        self._musl_policy = musl_policy

        base_arch = arch.baseline.value
        _validate_pep600_compliance(policies)
        for policy in policies:
            if self._musl_policy is not None and policy["name"] not in {
                "linux",
                self._musl_policy,
            }:
                continue
            if arch.value in policy["symbol_versions"] or policy["name"] == "linux":
                name = f"{policy['name']}_{base_arch}"
                if policy["name"] != "linux":
                    symbol_versions = {
                        k: frozenset(v) for k, v in policy["symbol_versions"][base_arch].items()
                    }
                else:
                    symbol_versions = {}
                aliases = tuple(f"{alias}_{base_arch}" for alias in policy["aliases"])
                whitelist = _fixup_musl_libc_soname(libc, arch, policy["lib_whitelist"])
                blacklist = {k: frozenset(v) for k, v in policy["blacklist"].items()}
                policy_ = Policy(
                    name=name,
                    aliases=aliases,
                    priority=policy["priority"],
                    symbol_versions=symbol_versions,
                    whitelist=whitelist,
                    blacklist=blacklist,
                )
                self._policies.append(policy_)

        self._policies.sort()

        if self._libc_variant == Libc.MUSL:
            if musl_policy is None:
                self._musl_policy = "_".join(self._policies[1].name.split("_")[0:3])
                self._policies = [self._policies[0], self._policies[1]]
            assert len(self._policies) == 2, self._policies  # noqa: S101

        elif self._libc_variant == Libc.ANDROID:
            # Android wheels always have a platform tag with an API level (OS version). We won't
            # change that tag, we'll only use it to determine which libraries need to be grafted.
            if wheel_fn is None:
                msg = "wheel_fn is required when selecting Android policies"
                raise ValueError(msg)

            platforms = get_wheel_platforms(wheel_fn.name)
            if len(platforms) != 1:
                msg = f"Android wheels must have exactly one platform tag, got {platforms}"
                raise ValueError(msg)
            platform = platforms[0]
            api_level = android_api_level(platform)

            # One or two new API levels are created every year, so the policy file only lists the
            # levels in which the available libraries changed
            # (https://developer.android.com/ndk/guides/stable_apis). We should use the policy with
            # the highest API level that's less than or equal to the wheel's tag.
            for p in self._policies:
                if p.name.startswith("android") and android_api_level(p.name) <= api_level:
                    # Rename the policy to match the existing tag, and remove all other policies.
                    self._policies = [self.linux, replace(p, name=platform)]
                    break
            else:
                # The minimum API level is 21, which was the first version to be officially
                # supported by Python (see PEP 738).
                msg = f"no Android policies match the platform tag {platform!r}"
                raise ValueError(msg)

    def __iter__(self) -> Generator[Policy]:
        yield from self._policies

    @property
    def libc(self) -> Libc:
        return self._libc_variant

    @property
    def architecture(self) -> Architecture:
        return self._architecture

    @property
    def highest(self) -> Policy:
        return self._policies[-1]

    @property
    def lowest(self) -> Policy:
        """Lowest policy that's not linux"""
        return self._policies[1]

    @property
    def linux(self) -> Policy:
        return self._policies[0]

    def get_policy_by_name(self, name: str) -> Policy:
        matches = [p for p in self._policies if p.name == name or name in p.aliases]
        if len(matches) == 0:
            msg = f"no policy named {name!r} found"
            raise LookupError(msg)
        if len(matches) > 1:
            msg = "Internal error. Policies should be unique"
            raise RuntimeError(msg)
        return matches[0]

    def versioned_symbols_policy(
        self,
        versioned_symbols: dict[str, set[str]],
    ) -> Policy:
        def policy_is_satisfied(
            policy_name: str,
            policy_sym_vers: dict[str, set[str]],
        ) -> bool:
            policy_satisfied = True
            for name in set(required_vers) & set(policy_sym_vers):
                if not required_vers[name].issubset(policy_sym_vers[name]):
                    for symbol in required_vers[name] - policy_sym_vers[name]:
                        logger.debug(
                            "Package requires %s, incompatible with policy %s which requires %s",
                            symbol,
                            policy_name,
                            policy_sym_vers[name],
                        )
                    policy_satisfied = False
            return policy_satisfied

        required_vers: dict[str, set[str]] = {}
        for symbols in versioned_symbols.values():
            for symbol in symbols:
                sym_name, _, _ = symbol.partition("_")
                required_vers.setdefault(sym_name, set()).add(symbol)
        for p in self._policies[::-1]:
            policy_sym_vers = {
                sym_name: {sym_name + "_" + version for version in versions}
                for sym_name, versions in p.symbol_versions.items()
            }
            if policy_is_satisfied(p.name, policy_sym_vers):
                return p

        # the base policy (generic linux) should always match
        msg = "Internal error"  # pragma: no cover
        raise RuntimeError(msg)  # pragma: no cover

    def lddtree_external_references(
        self,
        lddtree: DynamicExecutable,
        wheel_path: Path,
    ) -> dict[str, ExternalReference]:
        def filter_libs(
            libs: Iterable[str],
            whitelist: frozenset[str],
        ) -> Generator[str]:
            for lib in libs:
                if "ld-linux" in lib or lib in ["ld64.so.2", "ld64.so.1"]:
                    # always exclude ELF dynamic linker/loader
                    # 'ld64.so.2' on s390x
                    # 'ld64.so.1' on ppc64le
                    # 'ld-linux*' on other platforms
                    continue
                if self.libc == Libc.ANDROID and LIBPYTHON_RE.match(lib):
                    # libpython dependencies are normal on Android.
                    continue
                if lib in whitelist:
                    # exclude any libs in the whitelist
                    continue
                yield lib

        def get_req_external(libs: set[str], whitelist: frozenset[str]) -> set[str]:
            # get all the required external libraries
            libs = libs.copy()
            reqs = set()
            while libs:
                lib = libs.pop()
                reqs.add(lib)
                for dep in filter_libs(lddtree.libraries[lib].needed, whitelist):
                    if dep not in reqs:
                        libs.add(dep)
            return reqs

        ret: dict[str, ExternalReference] = {}
        for p in self._policies:
            needed_external_libs: set[str] = set()
            blacklist = {}

            if p != self.linux:
                # special-case the generic linux platform here, because it
                # doesn't have a whitelist. or, you could say its
                # whitelist is the complete set of all libraries. so nothing
                # is considered "external" that needs to be copied in.
                whitelist = p.whitelist
                blacklist_libs = set(p.blacklist.keys()) & set(lddtree.needed)
                blacklist_reduced = {k: p.blacklist[k] for k in blacklist_libs}
                blacklist = filter_undefined_symbols(
                    lddtree.realpath,
                    blacklist_reduced,
                )
                needed_external_libs = get_req_external(
                    set(filter_libs(lddtree.needed, whitelist)),
                    whitelist,
                )

            pol_ext_deps: dict[str, Path | None] = {}
            for lib in needed_external_libs:
                if is_subdir(lddtree.libraries[lib].realpath, wheel_path):
                    # we didn't filter libs that resolved via RPATH out
                    # earlier because we wanted to make sure to pick up
                    # our elf's indirect dependencies. But now we want to
                    # filter these ones out, since they're not "external".
                    logger.debug("RPATH FTW: %s", lib)
                    continue
                pol_ext_deps[lib] = lddtree.libraries[lib].realpath
            ret[p.name] = ExternalReference(pol_ext_deps, blacklist, p)
        return ret


def _validate_pep600_compliance(policies: Any) -> None:  # noqa: ANN401
    symbol_versions: dict[str, dict[str, set[str]]] = {}
    lib_whitelist: set[str] = set()
    for policy in sorted(policies, key=lambda x: x["priority"], reverse=True):
        if policy["name"] == "linux":
            continue
        if not lib_whitelist.issubset(set(policy["lib_whitelist"])):
            diff = lib_whitelist - set(policy["lib_whitelist"])
            msg = (
                'Invalid "policy.json" file. Missing whitelist libraries in '
                f'"{policy["name"]}" compared to previous policies: {diff}'
            )
            raise ValueError(msg)
        lib_whitelist.update(policy["lib_whitelist"])
        for arch in policy["symbol_versions"]:
            symbol_versions_arch = symbol_versions.get(arch, defaultdict(set))
            for prefix in policy["symbol_versions"][arch]:
                policy_symbol_versions = set(policy["symbol_versions"][arch][prefix])
                if not symbol_versions_arch[prefix].issubset(policy_symbol_versions):
                    diff = symbol_versions_arch[prefix] - policy_symbol_versions
                    msg = (
                        'Invalid "policy.json" file. Symbol versions missing '
                        f'in "{policy["name"]}_{arch}" for "{prefix}" '
                        f"compared to previous policies: {diff}"
                    )
                    raise ValueError(msg)
                symbol_versions_arch[prefix].update(
                    policy["symbol_versions"][arch][prefix],
                )
            symbol_versions[arch] = symbol_versions_arch


def _fixup_musl_libc_soname(
    libc: Libc,
    arch: Architecture,
    whitelist: list[str],
) -> frozenset[str]:
    if libc != Libc.MUSL:
        return frozenset(whitelist)
    soname_map = {
        "libc.so": {
            Architecture.x86_64: "libc.musl-x86_64.so.1",
            Architecture.i686: "libc.musl-x86.so.1",
            Architecture.aarch64: "libc.musl-aarch64.so.1",
            Architecture.s390x: "libc.musl-s390x.so.1",
            Architecture.ppc64le: "libc.musl-ppc64le.so.1",
            Architecture.armv7l: "libc.musl-armv7.so.1",
            Architecture.riscv64: "libc.musl-riscv64.so.1",
            Architecture.loongarch64: "libc.musl-loongarch64.so.1",
        },
    }
    new_whitelist = []
    for soname in whitelist:
        if soname in soname_map:
            new_soname = soname_map[soname][arch.baseline]
            logger.debug("Replacing whitelisted '%s' by '%s'", soname, new_soname)
            new_whitelist.append(new_soname)
        else:
            new_whitelist.append(soname)
    return frozenset(new_whitelist)


def get_replace_platforms(name: str) -> list[str]:
    """Extract platform tag replacement rules from policy

    >>> get_replace_platforms('linux_x86_64')
    []
    >>> get_replace_platforms('linux_i686')
    []
    >>> get_replace_platforms('manylinux1_x86_64')
    ['linux_x86_64']
    >>> get_replace_platforms('manylinux1_i686')
    ['linux_i686']

    """
    if name.startswith("linux"):
        return []
    if name.startswith("manylinux_"):
        return ["linux_" + "_".join(name.split("_")[3:])]
    if name.startswith("musllinux_"):
        return ["linux_" + "_".join(name.split("_")[3:])]
    return ["linux_" + "_".join(name.split("_")[1:])]


__all__ = [
    "WheelPolicies",
]
