#!/bin/env python3

import argparse
import fnmatch
import os
import sys
from contextlib import contextmanager, closing

from msys2_devtools.exttarfile import tarfile
from msys2_devtools.utils import vercmp
from fastprogress.fastprogress import progress_bar
from concurrent.futures import ThreadPoolExecutor
import pefile


def find_dbs(target_dir: str) -> list[str]:
    db_paths = set()
    target_dir = os.path.realpath(target_dir)
    for root, dirs, files in os.walk(target_dir):
        for name in files:
            if fnmatch.fnmatch(name, "*.db"):
                db_paths.add(os.path.realpath(os.path.join(root, name)))

    return sorted(db_paths)


def get_package_paths(root_path: str) -> list[str]:
    dbs = find_dbs(root_path)
    paths = set()
    for db_path in dbs:

        with tarfile.TarFile.open(db_path, mode="r") as tar:
            for info in tar.getmembers():
                file_name = info.name.rsplit("/", 1)[-1]
                if file_name == "desc":
                    infodata = tar.extractfile(info).read().decode()
                    lines = infodata.splitlines()
                    filename = lines[lines.index("%FILENAME%") + 1]
                    paths.add(os.path.join(os.path.dirname(db_path), filename))
    return sorted(paths)


def parse_buildinfo(buildinfo: str) -> dict[str, list[str]]:
    res: dict[str, list[str]] = {}
    for line in buildinfo.splitlines():
        line = line.strip()
        if not line:
            continue

        key, value = line.split(" =", 1)
        value = value.strip()
        values = [value] if value else []
        res.setdefault(key, []).extend(values)

    return res


@contextmanager
def pefile_set_max_import_symbols(max_symbols: int):
    old = pefile.MAX_IMPORT_SYMBOLS
    pefile.MAX_IMPORT_SYMBOLS = max_symbols
    try:
        yield
    finally:
        pefile.MAX_IMPORT_SYMBOLS = old


# python uses .pyd for extensions, ruby uses .so, and octave uses .oct
PE_FILE_EXTENSIONS = (".dll", ".exe", ".pyd", ".so", ".oct")


def get_imported_dlls(data: bytes, name: str) -> set[str]:
    try:
        with pefile_set_max_import_symbols(0x10000), \
            closing(pefile.PE(data=data, fast_load=True)) as pe:
            pe.parse_data_directories(directories=[
                pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT']
            ])

            if len(pe.get_warnings()) > 0:
                print(f"Warnings for PE ({name}):")
                pe.show_warnings()

            if not hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
                return set()
            return {entry.dll.decode().lower() for entry in pe.DIRECTORY_ENTRY_IMPORT}
    except pefile.PEFormatError:
        print(f"PEFormatError while parsing PE file ({name}), skipping imports")
        return set()


def get_buildinfo(path: str, needs_files: bool, needs_imports: bool) -> tuple[dict, set[str], set[str]]:
    buildinfo = None
    files = set()
    imports = set()
    with tarfile.TarFile.open(path, mode="r") as tar:
        for info in tar.getmembers():
            if info.name == ".BUILDINFO":
                buildinfo = tar.extractfile(info).read().decode()
            if needs_files and info.isfile() and not info.name.startswith("."):
                files.add(info.name)
            if needs_imports and info.isfile() and info.name.lower().endswith(PE_FILE_EXTENSIONS):
                data = tar.extractfile(info).read()
                imports.update(get_imported_dlls(data, info.name))

    if buildinfo is None:
        raise RuntimeError(f"Cannot find .BUILDINFO in {path}")
    return parse_buildinfo(buildinfo), files, imports


def parse_vercmp(comparison_expression) -> tuple[str, str | None, str | None]:
    operators = ['>=', '<=', '=', '>', '<']

    package_name = comparison_expression
    operator = None
    version = None

    for op in operators:
        if op in comparison_expression:
            parts = comparison_expression.split(op, 1)
            package_name = parts[0].strip()
            version = parts[1].strip()
            operator = op
            break

    return package_name, operator, version


def parse_full_package_name(full_name: str) -> tuple[str, str]:
    name = full_name.rsplit("-", 3)[0]
    version = full_name[len(name) + 1:].rsplit("-", 1)[0]
    return name, version


def package_matches(full_name: str, expression: str) -> bool:
    package_name, package_version = parse_full_package_name(full_name)
    ex_name, ex_operator, ex_version = parse_vercmp(expression)
    if ex_name != package_name:
        return False

    if ex_operator is None or ex_version is None:
        return True

    if ex_operator == "=":
        return ex_version == package_version
    elif ex_operator == ">":
        return vercmp(package_version, ex_version) > 0
    elif ex_operator == "<":
        return vercmp(package_version, ex_version) < 0
    elif ex_operator == ">=":
        return vercmp(package_version, ex_version) >= 0
    elif ex_operator == "<=":
        return vercmp(package_version, ex_version) <= 0
    else:
        raise ValueError(f"Unknown operator: {ex_operator}")


def main(argv):
    parser = argparse.ArgumentParser(
        description="List things about the package contents", allow_abbrev=False
    )
    parser.add_argument("root", help="path to root dir")
    parser.add_argument("--built-with-package",
                        help="filter packages that had this package installed during build time "
                             "(optionally with a version constraint, e.g. 'foo>=1.0')")
    parser.add_argument("--contains-file",
                        help="filter packages that contain files matching this glob pattern (e.g. '*.a')")
    parser.add_argument("--imports-dll",
                        action="append",
                        default=[],
                        metavar="DLL",
                        help="filter packages that import this DLL (e.g. 'libfoo.dll'); "
                             "can be given multiple times")

    args = parser.parse_args(argv[1:])

    installed_filter = args.built_with_package
    file_pattern = args.contains_file
    dll_filters = args.imports_dll

    found = set()
    paths = get_package_paths(args.root)

    def process_path(p):
        buildinfo, files, imports = get_buildinfo(p, needs_files=file_pattern is not None, needs_imports=bool(dll_filters))
        return buildinfo, files, imports

    with ThreadPoolExecutor(4) as executor:
        for buildinfo, files, imports in progress_bar(executor.map(process_path, paths), total=len(paths), leave=False):
            if installed_filter is not None:
                for installed_full_name in buildinfo["installed"]:
                    if package_matches(installed_full_name, installed_filter):
                        break
                else:
                    continue

            if file_pattern is not None:
                for file in files:
                    if fnmatch.fnmatch(file, file_pattern):
                        break
                else:
                    continue

            if dll_filters:
                if not any(dll in imports for dll in dll_filters):
                    continue

            found.update(buildinfo["pkgbase"])

    for f in sorted(found):
        print(f)


if __name__ == "__main__":
    main(sys.argv)
