#!/usr/bin/env python3

import argparse
import asyncio
import aiohttp
import gzip
import hashlib
import os
import re
import shutil
import subprocess
import sys
import tarfile
import tempfile
import zstandard

from collections import deque
from functools import cmp_to_key

async def download_file(session, url, dest_path, max_retries=3, retry_delay=2, timeout=60, checksum=None):
    """Asynchronous file download with retries."""
    attempt = 0
    while attempt < max_retries:
        try:
            async with session.get(url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
                if response.status == 200:
                    with open(dest_path, "wb") as f:
                        content = await response.read()

                        # verify checksum if provided
                        if checksum:
                            sha256 = hashlib.sha256(content).hexdigest()
                            if sha256 != checksum:
                                raise Exception(f"SHA256 mismatch for {url}: expected {checksum}, got {sha256}")

                        f.write(content)
                    print(f"Downloaded {url} at {dest_path}")
                    return
                else:
                    raise Exception(f"Failed to download {url}, Status Code: {response.status}")
        except (asyncio.CancelledError, asyncio.TimeoutError, aiohttp.ClientError) as e:
            print(f"Error downloading {url}: {type(e).__name__} - {e}. Retrying...")

        attempt += 1
        await asyncio.sleep(retry_delay)

    raise Exception(f"Failed to download {url} after {max_retries} attempts.")

async def download_deb_files_parallel(mirror, packages, tmp_dir):
    """Download .deb files in parallel."""
    os.makedirs(tmp_dir, exist_ok=True)

    tasks = []
    timeout = aiohttp.ClientTimeout(total=60)
    async with aiohttp.ClientSession(timeout=timeout) as session:
        for pkg, info in packages.items():
            filename = info.get("Filename")
            if filename:
                url = f"{mirror}/{filename}"
                dest_path = os.path.join(tmp_dir, os.path.basename(filename))
                tasks.append(asyncio.create_task(download_file(session, url, dest_path, checksum=info.get("SHA256"))))

        await asyncio.gather(*tasks)

async def download_package_index_parallel(mirror, arch, suites, check_sig, keyring):
    """Download package index files for specified suites and components entirely in memory."""
    tasks = []
    timeout = aiohttp.ClientTimeout(total=60)

    async with aiohttp.ClientSession(timeout=timeout) as session:
        for suite in suites:
            for component in ["main", "universe"]:
                tasks.append(fetch_and_decompress(session, mirror, arch, suite, component, check_sig, keyring))

        results = await asyncio.gather(*tasks)

    merged_content = ""
    for result in results:
        if isinstance(result, str):
            if merged_content:
                merged_content += "\n\n"
            merged_content += result

    return merged_content

async def fetch_and_decompress(session, mirror, arch, suite, component, check_sig, keyring):
    """Fetch and decompress the Packages.gz file."""

    path = f"{component}/binary-{arch}/Packages.gz"
    url = f"{mirror}/dists/{suite}/{path}"

    async with session.get(url) as response:
        if response.status == 200:
            compressed_data = await response.read()
            decompressed_data = gzip.decompress(compressed_data).decode('utf-8')
            print(f"Downloaded index: {url}")

            if check_sig:
                # Verify the package index against the sha256 recorded in the Release file
                release_file_content = await fetch_release_file(session, mirror, suite, keyring)
                packages_sha = parse_release_file(release_file_content, path)

                sha256 = hashlib.sha256(compressed_data).hexdigest()
                if sha256 != packages_sha:
                    raise Exception(f"SHA256 mismatch for {path}: expected {packages_sha}, got {sha256}")
                print(f"Checksum verified for {path}")

            return decompressed_data
        else:
            print(f"Skipped index: {url} (doesn't exist)")
            return None

async def fetch_release_file(session, mirror, suite, keyring):
    """Fetch Release and Release.gpg files and verify the signature."""

    release_url = f"{mirror}/dists/{suite}/Release"
    release_gpg_url = f"{mirror}/dists/{suite}/Release.gpg"

    with tempfile.NamedTemporaryFile() as release_file, tempfile.NamedTemporaryFile() as release_gpg_file:
        await download_file(session, release_url, release_file.name)
        await download_file(session, release_gpg_url, release_gpg_file.name)

        print("Verifying signature of Release with Release.gpg.")
        # Use gpgv rather than gpg for verification. gpgv verifies a detached
        # signature against a fixed keyring without involving gpg-agent or
        # keyboxd, which makes it robust on hosts running GnuPG 2.4+ (e.g. Azure
        # Linux) where "gpg --keyring" routes through keyboxd and can fail.
        verify_command = ["gpgv"]
        if keyring:
            verify_command += ["--keyring", keyring]
        verify_command += [release_gpg_file.name, release_file.name]
        result = subprocess.run(verify_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        if result.returncode != 0:
            raise Exception(f"Signature verification failed: {result.stderr.decode('utf-8')}")

        print("Signature verified successfully.")

        with open(release_file.name) as f:
            return f.read()

def parse_release_file(content, path):
    """Parses the Release file and returns sha256 checksum of the specified path."""

    # data looks like this:
    # <checksum>  <size>  <path>
    matches = re.findall(r'^ (\S*) +(\S*) +(\S*)$', content, re.MULTILINE)

    for entry in matches:
        # the file has both md5 and sha256 checksums, we want sha256 which has a length of 64
        if entry[2] == path and len(entry[0]) == 64:
            return entry[0]

    raise Exception(f"Could not find checksum for {path} in Release file.")

def parse_debian_version(version):
    """Parse a Debian package version into epoch, upstream version, and revision."""
    match = re.match(r'^(?:(\d+):)?([^-]+)(?:-(.+))?$', version)
    if not match:
        raise ValueError(f"Invalid Debian version format: {version}")
    epoch, upstream, revision = match.groups()
    return int(epoch) if epoch else 0, upstream, revision or ""

def compare_upstream_version(v1, v2):
    """Compare upstream or revision parts using Debian rules."""
    def tokenize(version):
        tokens = re.split(r'([0-9]+|[A-Za-z]+)', version)
        return [int(x) if x.isdigit() else x for x in tokens if x]

    tokens1 = tokenize(v1)
    tokens2 = tokenize(v2)

    for token1, token2 in zip(tokens1, tokens2):
        if type(token1) == type(token2):
            if token1 != token2:
                return (token1 > token2) - (token1 < token2)
        else:
            return -1 if isinstance(token1, str) else 1

    return len(tokens1) - len(tokens2)

def compare_debian_versions(version1, version2):
    """Compare two Debian package versions."""
    epoch1, upstream1, revision1 = parse_debian_version(version1)
    epoch2, upstream2, revision2 = parse_debian_version(version2)

    if epoch1 != epoch2:
        return epoch1 - epoch2

    result = compare_upstream_version(upstream1, upstream2)
    if result != 0:
        return result

    return compare_upstream_version(revision1, revision2)

def resolve_dependencies(packages, aliases, desired_packages):
    """Recursively resolves dependencies for the desired packages."""
    resolved = []
    to_process = deque(desired_packages)

    while to_process:
        current = to_process.popleft()
        resolved_package = current if current in packages else aliases.get(current, [None])[0]

        if not resolved_package:
            print(f"Error: Package '{current}' was not found in the available packages.")
            sys.exit(1)

        if resolved_package not in resolved:
            resolved.append(resolved_package)

            deps = packages.get(resolved_package, {}).get("Depends", "")
            if deps:
                deps = [dep.split(' ')[0] for dep in deps.split(', ') if dep]
                for dep in deps:
                    if dep not in resolved and dep not in to_process and dep in packages:
                        to_process.append(dep)

    return resolved

def parse_package_index(content):
    """Parses the Packages.gz file and returns package information."""
    packages = {}
    aliases = {}
    entries = re.split(r'\n\n+', content)

    for entry in entries:
        fields = dict(re.findall(r'^(\S+): (.+)$', entry, re.MULTILINE))
        if "Package" in fields:
            package_name = fields["Package"]
            version = fields.get("Version")
            filename = fields.get("Filename")
            depends = fields.get("Depends")
            provides = fields.get("Provides", None)
            sha256 = fields.get("SHA256")

            # Only update if package_name is not in packages or if the new version is higher
            if package_name not in packages or compare_debian_versions(version, packages[package_name]["Version"]) > 0:
                packages[package_name] = {
                    "Version": version,
                    "Filename": filename,
                    "Depends": depends,
                    "SHA256": sha256
                }

                # Update aliases if package provides any alternatives
                if provides:
                    provides_list = [x.strip() for x in provides.split(",")]
                    for alias in provides_list:
                        # Strip version specifiers
                        alias_name = re.sub(r'\s*\(=.*\)', '', alias)
                        if alias_name not in aliases:
                            aliases[alias_name] = []
                        if package_name not in aliases[alias_name]:
                            aliases[alias_name].append(package_name)

    return packages, aliases

def install_packages(mirror, packages_info, aliases, tmp_dir, extract_dir, ar_tool, desired_packages):
    """Downloads .deb files and extracts them."""
    resolved_packages = resolve_dependencies(packages_info, aliases, desired_packages)
    print(f"Resolved packages (including dependencies): {resolved_packages}")

    packages_to_download = {}

    for pkg in resolved_packages:
        if pkg in packages_info:
            packages_to_download[pkg] = packages_info[pkg]

        if pkg in aliases:
            for alias in aliases[pkg]:
                if alias in packages_info:
                    packages_to_download[alias] = packages_info[alias]

    asyncio.run(download_deb_files_parallel(mirror, packages_to_download, tmp_dir))

    package_to_deb_file_map = {}
    for pkg in resolved_packages:
        pkg_info = packages_info.get(pkg)
        if pkg_info:
            deb_filename = pkg_info.get("Filename")
            if deb_filename:
                deb_file_path = os.path.join(tmp_dir, os.path.basename(deb_filename))
                package_to_deb_file_map[pkg] = deb_file_path

    for pkg in reversed(resolved_packages):
        deb_file = package_to_deb_file_map.get(pkg)
        if deb_file and os.path.exists(deb_file):
            extract_deb_file(deb_file, tmp_dir, extract_dir, ar_tool)

    print("All done!")

def extract_deb_file(deb_file, tmp_dir, extract_dir, ar_tool):
    """Extract .deb file contents"""

    os.makedirs(extract_dir, exist_ok=True)

    with tempfile.TemporaryDirectory(dir=tmp_dir) as tmp_subdir:
        result = subprocess.run([ar_tool, "t", os.path.abspath(deb_file)], cwd=tmp_subdir, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        tar_filename = None
        for line in result.stdout.decode().splitlines():
            if line.startswith("data.tar"):
                tar_filename = line.strip()
                break

        if not tar_filename:
            raise FileNotFoundError(f"Could not find 'data.tar.*' in {deb_file}.")

        tar_file_path = os.path.join(tmp_subdir, tar_filename)
        print(f"Extracting {tar_filename} from {deb_file}..")

        with open(tar_file_path, "wb") as outfile:
            subprocess.run([ar_tool, "p", os.path.abspath(deb_file), tar_filename], check=True, stdout=outfile, stderr=subprocess.PIPE)

        file_extension = os.path.splitext(tar_file_path)[1].lower()

        if file_extension == ".xz":
            mode = "r:xz"
        elif file_extension == ".gz":
            mode = "r:gz"
        elif file_extension == ".zst":
            # zstd is not supported by standard library yet
            decompressed_tar_path = tar_file_path.replace(".zst", "")
            with open(tar_file_path, "rb") as zst_file, open(decompressed_tar_path, "wb") as decompressed_file:
                dctx = zstandard.ZstdDecompressor()
                dctx.copy_stream(zst_file, decompressed_file)

            tar_file_path = decompressed_tar_path
            mode = "r"
        else:
            raise ValueError(f"Unsupported compression format: {file_extension}")

        with tarfile.open(tar_file_path, mode) as tar:
            tar.extractall(path=extract_dir, filter=_rootfs_extraction_filter)

def _rootfs_extraction_filter(member, dest_path):
    """Tarfile extraction filter based on the 'data' filter that additionally
    rewrites absolute-target symlinks/hardlinks into rootfs-relative paths.
    """
    if (member.issym() or member.islnk()) and os.path.isabs(member.linkname):
        link_dir = os.path.dirname(member.name)
        new_linkname = os.path.relpath(member.linkname.lstrip('/'),
                                       start=link_dir or '.')
        member = member.replace(linkname=new_linkname, deep=False)
    return tarfile.data_filter(member, dest_path)

def finalize_setup(rootfsdir):
    lib_dir = os.path.join(rootfsdir, 'lib')
    usr_lib_dir = os.path.join(rootfsdir, 'usr', 'lib')

    if os.path.exists(lib_dir):
        if os.path.islink(lib_dir):
            os.remove(lib_dir)
        else:
            os.makedirs(usr_lib_dir, exist_ok=True)

            for item in os.listdir(lib_dir):
                src = os.path.join(lib_dir, item)
                dest = os.path.join(usr_lib_dir, item)

                if os.path.isdir(src):
                    shutil.copytree(src, dest, dirs_exist_ok=True)
                else:
                    shutil.copy2(src, dest)

            shutil.rmtree(lib_dir)

    os.symlink(usr_lib_dir, lib_dir)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate rootfs for .NET runtime on Debian-like OS")
    parser.add_argument("--arch", required=True, help="Architecture (e.g., amd64, loong64, etc.)")
    parser.add_argument("--rootfsdir", required=True, help="Destination directory.")
    parser.add_argument('--suite', required=True, action='append', help='Specify one or more repository suites to collect index data.')
    parser.add_argument("--mirror", required=True, help="Mirror (e.g., http://ftp.debian.org/debian-ports etc.)")
    parser.add_argument("--artool", required=False, default="ar", help="ar tool to extract debs (e.g., ar, llvm-ar etc.)")
    parser.add_argument("--force-check-gpg", required=False, action='store_true', help="Verify the packages against signatures in Release file.")
    parser.add_argument("--keyring", required=False, default='', help="Keyring file to check signature of Release file.")
    parser.add_argument("packages", nargs="+", help="List of package names to be installed.")

    args = parser.parse_args()

    DESIRED_PACKAGES = args.packages + [ # base packages
        "dpkg",
        "busybox",
        "libc-bin",
        "base-files",
        "base-passwd",
        "debianutils"
    ]

    print(f"Creating rootfs. rootfsdir: {args.rootfsdir}, arch: {args.arch}, suites: {args.suite}, mirror: {args.mirror}")

    check_sig = args.force_check_gpg
    if check_sig and not args.keyring:
        print("ERROR: --force-check-gpg requires --keyring to specify a keyring file for signature verification.")
        print("Install the appropriate keyring package (e.g., debian-ports-archive-keyring, ubuntu-archive-keyring)")
        print("or pass --skipsigcheck to build-rootfs.sh to disable signature checking.")
        sys.exit(1)

    package_index_content = asyncio.run(download_package_index_parallel(args.mirror, args.arch, args.suite, check_sig, args.keyring))

    packages_info, aliases = parse_package_index(package_index_content)

    with tempfile.TemporaryDirectory() as tmp_dir:
        install_packages(args.mirror, packages_info, aliases, tmp_dir, args.rootfsdir, args.artool, DESIRED_PACKAGES)

    finalize_setup(args.rootfsdir)
