# -*- coding: utf-8 -*-
"""
SPDX-FileCopyrightText: 2023-2025 PeppermintOS Team
(peppermintosteam@proton.me)

SPDX-License-Identifier: GPL-3.0-or-later

This module provides functions to build an initramfs for a Linux system.

It includes functions for:
- Copying the kernel image to the boot directory.
- Generating the initramfs using dracut.
- Retrieving kernel version information from within a chroot environment.

Credits:
  - PeppermintOS Team (peppermintosteam@proton.me) - Development and maintenance of the project.

License:
  This code is distributed under the GNU General Public License version 3 or later (GPL-3.0-or-later).
  For more details, please refer to the LICENSE file included in the project or visit:
  https://www.gnu.org/licenses/gpl-3.0.html
"""

import os
import subprocess
import shutil
import logging
import sys

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, BASE_DIR)

try:
    from builder.core.bootstrap.paths import paths
    from builder.core.xbps_commands import xbps_commands
    from builder.configs import logger_config

    logger = logger_config.setup_logger('initramfs_builder')
except ImportError as e:
    print(f"Error importing necessary modules: {e}. Ensure your environment is set up correctly.")
    sys.exit(1)


def execute_xbps_query_chroot(rootfs_path, xbps_query_args, iso_build_config):
    """Executes xbps-query inside the chroot and returns the output."""
    xbps_arch = iso_build_config.get('target_arch', 'x86_64')
    command_xbps_query = ["chroot", rootfs_path, f"XBPS_ARCH={xbps_arch}", xbps_commands["XBPS_QUERY_CMD"]] + list(xbps_query_args)
    completed_process = subprocess.run(command_xbps_query, capture_output=True, text=True)
    return completed_process.stdout.strip()

def execute_xbps_uhelper_chroot(rootfs_path, xbps_uhelper_args, iso_build_config):
    """Executes xbps-uhelper inside the chroot and returns the output."""
    command_xbps_uhelper = ["chroot", rootfs_path, xbps_commands["XBPS_UHELPER_CMD"]] + list(xbps_uhelper_args)
    completed_process = subprocess.run(command_xbps_uhelper, capture_output=True, text=True)
    return completed_process.stdout.strip()

def copy_kernel_image(rootfs_path, boot_path, iso_build_config):
    """
    Copies the kernel image file (vmlinuz or vmlinux) from the ROOTFS to the ISO's BOOT directory,
    renaming it to 'vmlinuz' (x86_64, i686) or 'vmlinux' (aarch64) and using the correct kernel version.
    """
    logger.info("=> Copying kernel image to the ISO BOOT directory...")

    kernel_versions = get_kernel_version_in_chroot(rootfs_path, iso_build_config)
    if not kernel_versions or 'version_final_part' not in kernel_versions:
        logger.error("=> ERROR: Could not obtain kernel version to copy the image.")
        raise ValueError("Could not obtain kernel version.")
    kernel_file_suffix = kernel_versions['version_final_part']

    target_arch = iso_build_config.get('target_arch', 'x86_64').lower()

    if target_arch.startswith('i686') or target_arch == 'x86_64':
        kernel_filename_in_boot = "vmlinuz"
    elif target_arch == 'aarch64':
        kernel_filename_in_boot = "vmlinux"
    else:
        logger.error(f"=> Unknown target architecture: '{target_arch}'. Cannot copy kernel.")
        raise ValueError(f"Unsupported target architecture: '{target_arch}'")

    boot_dir = os.path.join(rootfs_path, 'boot')
    kernel_filename_in_rootfs = None
    try:
        boot_files = os.listdir(boot_dir)
        for filename in boot_files:
            if target_arch.startswith('i686') or target_arch == 'x86_64':
                if filename.startswith(f'vmlinuz-{kernel_file_suffix}'):
                    kernel_filename_in_rootfs = filename
                    break
            elif target_arch == 'aarch64':
                if filename.startswith(f'vmlinux-{kernel_file_suffix}'):
                    kernel_filename_in_rootfs = filename
                    break
    except FileNotFoundError:
        logger.error(f"=> ERROR: /boot directory not found at: {boot_dir}")
        raise FileNotFoundError(f"/boot directory not found: {boot_dir}")

    if not kernel_filename_in_rootfs:
        logger.error(f"=> ERROR: Kernel file with suffix '{kernel_file_suffix}' not found in: {boot_dir}")
        raise FileNotFoundError(f"Kernel file not found with suffix: '{kernel_file_suffix}'")


    kernel_path_in_rootfs = os.path.join(rootfs_path, "boot", kernel_filename_in_rootfs)
    kernel_path_in_boot = os.path.join(boot_path, kernel_filename_in_boot)

    logger.info(f"=> Copying kernel from: {kernel_path_in_rootfs} to: {kernel_path_in_boot}...")

    shutil.copy2(kernel_path_in_rootfs, kernel_path_in_boot)
    logger.info(f"=> Kernel copied successfully to: {kernel_path_in_boot}")
    return kernel_path_in_boot

def get_kernel_version_in_chroot(rootfs_path, iso_build_config):
    logger.info("=> DEBUG: Entering get_kernel_version_in_chroot...") # Debug line

    # >>> DEBUG: Verificar target_arch recebido <<<
    target_arch = iso_build_config.get('rootfs', {}).get('target_arch', 'x86_64').lower()
    logger.info(f"=> DEBUG: get_kernel_version_in_chroot using target_arch: '{target_arch}'") # Debug line

    boot_dir = os.path.join(rootfs_path, 'boot')
    logger.info(f"=> DEBUG: get_kernel_version_in_chroot checking directory: '{boot_dir}'") # Debug line

    try:
        boot_files = os.listdir(boot_dir)
        logger.info(f"=> DEBUG: get_kernel_version_in_chroot found files: {boot_files}") # Debug line
    except FileNotFoundError:
        logger.error(f"=> ERROR: /boot directory not found at: {boot_dir}")
        return None

    kernel_version = None
    for filename in boot_files:
        logger.info(f"=> DEBUG: Checking filename: '{filename}'") # Debug line
        if target_arch.startswith('i686') or target_arch == 'x86_64':
            if filename.startswith('vmlinuz-'):
                logger.info(f"=> DEBUG: Matched vmlinuz pattern for x86/i686: '{filename}'") # Debug line
                kernel_version = filename.replace('vmlinuz-', '')
                break
        elif target_arch == 'aarch64':
            if filename.startswith('vmlinux-'):
                logger.info(f"=> DEBUG: Matched vmlinux pattern for aarch64: '{filename}'") # Debug line
                kernel_version = filename.replace('vmlinux-', '')
                break # Correct break for inner if

    # >>> DEBUG: Verificar valor de kernel_version após o loop <<<
    logger.info(f"=> DEBUG: kernel_version after loop: '{kernel_version}'") # Debug line

    if not kernel_version:
        logger.error(f"=> ERROR: Kernel image not found in: {boot_dir}")
        return None

    logger.info(f"=> Kernel version found in /boot: {kernel_version}")

    kernel_versions_dict = {
        'version_final_part': kernel_version,
        'version_major_minor': ".".join(kernel_version.split('.')[:2])
    }
    logger.info(f"=> DEBUG: kernel_versions_dict: {kernel_versions_dict}") # Debug line
    return kernel_versions_dict

def get_kernel_module_dir_in_chroot(rootfs_path, kernel_version):
    """
    Gets the path of the kernel modules directory in the chroot using the provided version.
    """
    kernel_modules_dir = f"/lib/modules/{kernel_version}"
    return kernel_modules_dir

def create_initramfs(rootfs_path, boot_path, xbps_commands, iso_build_config):
    logger.info("=> Starting initramfs creation...")
    logger.info("=> Creating initramfs using dracut...")

    kernel_versions = get_kernel_version_in_chroot(rootfs_path, iso_build_config)
    if not kernel_versions or 'version_final_part' not in kernel_versions:
        logger.error("=> ERROR: Could not obtain kernel version to create the initramfs.")
        raise ValueError("Could not obtain kernel version.")
    kernel_version_final_part = kernel_versions["version_final_part"]

    kernel_modules_dir = get_kernel_module_dir_in_chroot(rootfs_path, kernel_version_final_part)

    KERNELVERSION = execute_xbps_uhelper_chroot(rootfs_path, ["getpkgversion", kernel_version_final_part], iso_build_config)

    dracut_command = [
        "chroot",
        rootfs_path,
        xbps_commands["DRACUT_CMD"],
        "-N",
        "--force",
        f"--kmoddir={kernel_modules_dir}",
        "--add-drivers", "ahci",
        "--force-add", "vmklive",
        "--omit", "systemd",
        "/boot/initrd",
        kernel_version_final_part,
    ]
    logger.info("=> *** EXTRA CHECK: EXECUTING CORRECTED ISO_BUILDER.PY SCRIPT VERSION (BEFORE DRACUT!) ***")
    try:
        subprocess.run(dracut_command, check=True)
    except subprocess.CalledProcessError as e:
        logger.error(f"=> Error executing dracut: Exit code: {e.returncode}")
        logger.error(f"=> Failed command: {e.cmd}")
        logger.error(f"=> Stderr: {e.stderr.decode('utf-8', errors='ignore') if e.stderr else '<empty>'}")
        raise
    else:
        logger.info("=> Initramfs created successfully using dracut!")
        initrd_path_in_rootfs = os.path.join(rootfs_path, "boot", "initrd")
        initrd_path_in_boot = os.path.join(boot_path, "initrd")
        logger.info(f"=> Moving initramfs from: {initrd_path_in_rootfs} to: {initrd_path_in_boot}...")
        shutil.move(initrd_path_in_rootfs, initrd_path_in_boot)
        logger.info(f"=> Initramfs moved successfully to: {initrd_path_in_boot}")
        return initrd_path_in_boot, kernel_versions