import os
import sys
import tempfile
import shutil
from pathlib import Path

import image_tools
from genimage_py import GenImageTool

def convert_firmware_to_rtapp_image(firmware_file: Path, output_dir: Path) -> Path:
    """
    Converts a raw firmware file (ELF/BIN) into a K230 RT-App image container
    and places the final gzipped image in the specified output directory.

    Args:
        firmware_file: Path to the raw input firmware file.
        output_dir: Directory where the final 'rtapp.elf.gz' should be placed.

    Returns:
        Path to the final 'rtapp.elf.gz' file.
    """
    print(f"Converting firmware file {firmware_file.name} to K230 RT-App image...")

    output_dir.mkdir(parents=True, exist_ok=True)
    
    # a. Gzip the firmware file
    gzip_tool = image_tools.K230PrivGzip()
    rtapp_gzipped_file = gzip_tool.compress_file(str(firmware_file))

    # b. Create the K230 image container (MkImage)
    rtapp_image_temp = image_tools.generate_temp_file_path("rt_app_", "_img")

    # Arguments simplified for this example; use actual args from MkImage documentation
    mkimg = image_tools.MkImage()
    mkimg.create_image(
        [rtapp_gzipped_file], rtapp_image_temp, "riscv", "u-boot", "firmware", "gzip", "0x00000000", "0x00000000", "rtapp",verbose = True
    )

    # c. Copy the final K230 RT-App image to the variant folder
    final_rtapp_name = "rtapp.elf.gz" # Common name for the generated RT-App image
    rtapp_output_path = output_dir / final_rtapp_name

    if not image_tools.generate_k230_image(rtapp_image_temp, str(rtapp_output_path)):
         raise RuntimeError("RTApp K230 final image generation failed.")

    # Cleanup temporary files
    os.remove(rtapp_gzipped_file)
    os.remove(rtapp_image_temp)

    print(f"   -> RT-App image placed in variant directory: {rtapp_output_path}")
    return rtapp_output_path

def find_generated_flash_image(output_dir: Path, expected_name: str = "sysimage_sdcard.kdimg") -> Path:
    expected_image = output_dir / expected_name
    if expected_image.exists():
        return expected_image

    candidates = sorted(
        path for path in output_dir.iterdir()
        if path.is_file() and path.suffix in (".kdimg", ".img")
    )
    if len(candidates) == 1:
        return candidates[0]

    if candidates:
        names = ", ".join(path.name for path in candidates)
        raise RuntimeError(f"GenImageTool produced multiple image files: {names}")

    raise RuntimeError("GenImageTool failed to produce a final image file (.kdimg or .img)")

def generate_rtapp(variant_path: str, firmware_path: str) -> str:
    """
    Orchestrates the image generation process:
    1. Converts the raw firmware to a K230 RT-App image.
    2. Calls the genimage utility using the variant's configuration.

    Args:
        variant_path: The root path for the selected platform/variant configuration.
        firmware_path: Path to the raw firmware file (the ELF/BIN to be flashed).

    Returns:
        The path to the final generated flash image file.
    """
    print(f"\n--- Starting Image Generation for Variant: {variant_path} ---")
    
    # 1. Setup Paths
    variant_root = Path(variant_path).resolve()
    firmware_file = Path(firmware_path).resolve()
    
    # Define the output location for the temporary RT-App component
    rtapp_output_dir = variant_root / "rtapp"
    
    # Define the genimage config file location
    genimage_config_file = variant_root / "genimage.cfg"
    
    if not genimage_config_file.exists():
        print(f"Error: Genimage config file not found at: {genimage_config_file}")
        raise FileNotFoundError(f"Missing genimage.cfg in variant folder: {variant_root}")

    # 2. Convert Firmware to RT-App Image (Now in its own function)
    convert_firmware_to_rtapp_image(firmware_file, rtapp_output_dir)

    # 3. Call GenImageTool
    print("Running GenImageTool to assemble final flash image...")

    temp_output_path = Path(tempfile.gettempdir()) / "genimage_output"
    
    if temp_output_path.exists():
        shutil.rmtree(temp_output_path)

    # rootpath is the variant path, config is the genimage.cfg in the variant path
    tool = GenImageTool(
        rootpath=str(variant_root), 
        outputpath=str(temp_output_path), 
        config_file=str(genimage_config_file)
    )
    tool.run()

    # 4. Find the resulting image file.
    final_image_files = find_generated_flash_image(temp_output_path)

    print(f"3. Final flash image generated at: {final_image_files}")
    return str(final_image_files)
