#!/usr/bin/env python3
"""
STM32 Flash Script using STM32CubeProgrammer CLI
Target: STM32F407xx
"""

import subprocess
import sys
import os
import argparse
from pathlib import Path


SCRIPT_DIR = Path(__file__).parent.resolve()
DEFAULT_ELF = SCRIPT_DIR / "build" / "Debug" / "CloudPlant.elf"
STM_PROGRAMMER_CLI = r"C:\Program Files\STMicroelectronics\STM32Cube\STM32CubeProgrammer\bin\   .exe"


def convert_elf_to_bin(elf_path):
    """Convert ELF to binary"""
    bin_path = str(elf_path).replace(".elf", ".bin")

    try:
        cmd = ["arm-none-eabi-objcopy", "-O", "binary", str(elf_path), bin_path]
        result = subprocess.run(cmd, capture_output=True, text=True)

        if result.returncode != 0:
            print(f"Conversion failed: {result.stderr}")
            return None

        print(f"Converted ELF to BIN: {bin_path}")
        print(f"Binary size: {os.path.getsize(bin_path)} bytes")
        return bin_path
    except FileNotFoundError:
        print("arm-none-eabi-objcopy not found!")
        print("Please install GNU Arm Embedded Toolchain or add to PATH")
        return None


def flash_with_cube(bin_path, port="SWD"):
    """Flash using STM32CubeProgrammer CLI"""
    if not os.path.exists(STM_PROGRAMMER_CLI):
        print(f"STM32CubeProgrammer not found at: {STM_PROGRAMMER_CLI}")
        return False

    bin_path_abs = Path(bin_path).resolve()
    cmd = [
        STM_PROGRAMMER_CLI,
        "-c",
        f"port={port}",
        "-w",
        str(bin_path_abs),
        "0x08000000",
        "-v",
        "-rst",
    ]

    print(f"Flashing {bin_path}...")
    print(f"Port: {port}")
    print(f"Address: 0x08000000")
    print("-" * 50)

    result = subprocess.run(cmd, capture_output=True, text=True)
    print(result.stdout)
    if result.stderr:
        print(result.stderr)

    if result.returncode == 0:
        print("Flash successful!")
        return True
    else:
        print(f"Flash failed (exit code: {result.returncode})")
        return False


def erase_with_cube(port="SWD"):
    """Erase chip using STM32CubeProgrammer CLI"""
    if not os.path.exists(STM_PROGRAMMER_CLI):
        print(f"STM32CubeProgrammer not found at: {STM_PROGRAMMER_CLI}")
        return False

    cmd = [STM_PROGRAMMER_CLI, "-c", f"port={port}", "-e", "all"]

    print("Erasing chip...")
    result = subprocess.run(cmd, capture_output=True, text=True)
    print(result.stdout)

    if result.returncode == 0:
        print("Erase successful!")
        return True
    else:
        print(f"Erase failed")
        return False


def main():
    parser = argparse.ArgumentParser(
        description="STM32 Flash Script using STM32CubeProgrammer CLI"
    )
    parser.add_argument(
        "elf_file",
        nargs="?",
        default=str(DEFAULT_ELF),
        help=f"ELF file to flash (default: {DEFAULT_ELF})",
    )
    parser.add_argument("--port", default="SWD", help="ST-Link port (default: SWD)")
    parser.add_argument(
        "--erase-only", action="store_true", help="Only erase chip, do not flash"
    )

    args = parser.parse_args()

    print("=" * 50)
    print("STM32 Flash Script (STM32CubeProgrammer CLI)")
    print("=" * 50)
    print(f"Target: STM32F407xx")
    print(f"ELF: {args.elf_file}")
    print(f"Programmer: {STM_PROGRAMMER_CLI}")
    print("=" * 50)

    if not os.path.exists(args.elf_file):
        print(f"ELF file not found: {args.elf_file}")
        sys.exit(1)

    if not os.path.exists(STM_PROGRAMMER_CLI):
        print(f"STM32CubeProgrammer not found: {STM_PROGRAMMER_CLI}")
        print("Please install STM32CubeProgrammer")
        sys.exit(1)

    if args.erase_only:
        success = erase_with_cube(args.port)
        sys.exit(0 if success else 1)

    bin_path = convert_elf_to_bin(args.elf_file)
    if not bin_path:
        sys.exit(1)

    success = flash_with_cube(bin_path, args.port)

    try:
        os.remove(bin_path)
    except:
        pass

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()