#!/usr/bin/env python3
import sys
import os
import time
import argparse
from pathlib import Path

# Add current directory to path to import local modules
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from k230_flash import K230FlashTool

def main():
    parser = argparse.ArgumentParser(
        description="K230 Arduino Flash Tool - Cross-platform firmware flashing",
        formatter_class=argparse.RawTextHelpFormatter 
    )

    generate_group = parser.add_argument_group(
        "Image Generation Options",
        description="Options used when generating the flash image structure before burning (useful for custom builds)."
    )

    generate_group.add_argument(
        "firmware_file",
        nargs="?",
        help="Path to the firmware file to flash"
    )

    generate_group.add_argument(
        "-p",
        "--platform-path",
        type=str,
        default=None,
        help="Path to the arduino platform path. Default: None"
    )

    generate_group.add_argument(
        "-v",
        "--variant",
        type=str,
        help="Specifies the target build variant (e.g., debug, release, test)"
    )

    burn_group = parser.add_argument_group(
        "Burn Firmware Options",
        description="Options related to the actual flashing process: specifying file, medium, and timing."
    )

    burn_group.add_argument(
        "-t",
        "--timeout",
        type=int,
        default=10,
        help="Timeout in seconds for bootloader detection (default: 10)"
    )

    burn_group.add_argument(
        "-m",
        "--medium-type",
        choices=["EMMC", "SDCARD", "SPI_NAND", "SPI_NOR", "OTP"],
        default="SDCARD",
        help="Target medium type (default: SDCARD)"
    )

    args = parser.parse_args()

    # Check firmware and variant folder
    if not args.firmware_file:
        print("Error: Firmware file is required.")
        parser.print_help()
        sys.exit(1)

    firmware_path = Path(args.firmware_file)
    if not firmware_path.exists():
        print(f"Error: Firmware file not found: {args.firmware_file}")
        sys.exit(1)

    variant_folder = Path(args.platform_path) / "variants" / args.variant
    if not variant_folder.exists():
        print(f"Error: Variant folder not found: {variant_folder}")
        sys.exit(1)

    try:
        from common import ensure_device_in_bootloader
        from genrtapp import generate_rtapp

        print("=== Generating Flash Image ===")
        image_path = generate_rtapp(variant_folder, firmware_path)

        # Resolve the uploader before entering the short-lived USB boot ROM window.
        flash_tool = K230FlashTool()

        print("=== Ensuring Device is in Bootloader Mode ===")
        ensure_device_in_bootloader(args.timeout)

        # Flash the firmware
        print(f"=== Flashing Firmware ===")
        print(f"Firmware: {firmware_path}")
        print(f"Image: {image_path}")
        print(f"Medium: {args.medium_type}")
        print()

        flash_tool.flash(
            image_path=str(image_path),
            medium_type=args.medium_type,
            auto_reboot=True
        )

        os.remove(image_path)

        print("\n=== Flash Operation Completed Successfully ===")

    except KeyboardInterrupt:
        print("\nOperation cancelled by user.")
        sys.exit(1)
    except Exception as e:
        print(f"\nError: {e}")
        sys.exit(1)

if __name__ == "__main__":
    print(sys.argv[1:])
    main()
