#!/usr/bin/env python3
"""
LSPi Project Flash Script
Support OpenOCD, J-Link, PyOCD and other programming tools
"""

import os
import sys
import subprocess
import shutil
import argparse

# Project configuration
PROJECT_NAME = "LSPi_Project"
ELF_FILE = f"build/{PROJECT_NAME}.elf"
BIN_FILE = f"build/{PROJECT_NAME}.bin"
HEX_FILE = f"build/{PROJECT_NAME}.hex"

# GD32F470 target configuration
TARGET_DEVICE = "gd32f470zg"
TARGET_CPU = "cortex-m4"

# Tool detection
def detect_tools():
    """Detect available programming tools in system"""
    tools = {}
    
    # Detect OpenOCD
    if shutil.which("openocd"):
        tools["openocd"] = "openocd"
    elif shutil.which("openocd.exe"):
        tools["openocd"] = "openocd.exe"
    
    # Detect J-Link
    if shutil.which("jlink"):
        tools["jlink"] = "jlink"
    elif shutil.which("jlink.exe"):
        tools["jlink"] = "jlink.exe"
    elif shutil.which("JLink.exe"):
        tools["jlink"] = "JLink.exe"
    
    # Detect PyOCD
    if shutil.which("pyocd"):
        tools["pyocd"] = "pyocd"
    elif shutil.which("pyocd.exe"):
        tools["pyocd"] = "pyocd.exe"
    
    return tools

def check_build_files(args):
    """Check if build files exist"""
    if not os.path.exists(args.elf):
        print(f"Error: ELF file not found: {args.elf}")
        print("Please run build command first: cmake -B build && cmake --build build")
        return False
    
    print(f"Found ELF file: {args.elf}")
    
    # Check if BIN and HEX files are generated
    if os.path.exists(args.bin):
        print(f"Found BIN file: {args.bin}")
    else:
        print(f"Warning: BIN file not found: {args.bin}")
    
    if os.path.exists(args.hex):
        print(f"Found HEX file: {args.hex}")
    else:
        print(f"Warning: HEX file not found: {args.hex}")
    
    return True

def flash_with_openocd(elf_path):
    """Flash using OpenOCD"""
    print("Flashing with OpenOCD...")
    
    # OpenOCD configuration script
    # Note: need to adjust interface and target config files based on actual hardware
    openocd_script = f"""
# Configuration for CMSIS-DAP debugger
source [find interface/cmsis-dap.cfg]
# Configuration for ST-Link debugger
# source [find interface/stlink.cfg]
# Configuration for J-Link debugger
# source [find interface/jlink.cfg]

# GD32F4xx target configuration (using STM32F4xx config as OpenOCD may not have GD32 specific config)
source [find target/stm32f4x.cfg]

# Reset configuration for GD32
# Try different reset configurations if one doesn't work
# reset_config srst_only
# reset_config srst_nogate
reset_config connect_assert_srst

# Program and verify
program {elf_path} verify reset exit
"""
    
    # Write script to temporary file
    with open("openocd_temp.cfg", "w") as f:
        f.write(openocd_script)
    
    try:
        cmd = ["openocd", "-f", "openocd_temp.cfg"]
        print(f"Executing: {' '.join(cmd)}")
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print("Flash successful!")
        else:
            print(f"Flash failed, return code: {result.returncode}")
            print(f"Stderr: {result.stderr}")
        
        return result.returncode == 0
    except Exception as e:
        print(f"Error executing OpenOCD: {e}")
        return False
    finally:
        # Cleanup temporary file
        if os.path.exists("openocd_temp.cfg"):
            os.remove("openocd_temp.cfg")

def flash_with_jlink(hex_path):
    """Flash using J-Link"""
    print("Flashing with J-Link...")
    
    # J-Link command script
    jlink_script = f"""
device {TARGET_DEVICE}
speed 4000
erase
loadfile {hex_path}
r
g
exit
"""
    
    # Write script to temporary file
    script_file = "jlink_temp.jlink"
    with open(script_file, "w") as f:
        f.write(jlink_script)
    
    try:
        cmd = ["JLink.exe", "-CommanderScript", script_file]
        print(f"Executing: {' '.join(cmd)}")
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print("Flash successful!")
        else:
            print(f"Flash failed, return code: {result.returncode}")
            print(f"Stderr: {result.stderr}")
        
        return result.returncode == 0
    except Exception as e:
        print(f"Error executing J-Link: {e}")
        return False
    finally:
        # Cleanup temporary file
        if os.path.exists(script_file):
            os.remove(script_file)

def flash_with_pyocd(hex_path):
    """Flash using PyOCD"""
    print("Flashing with PyOCD...")
    
    try:
        # Erase chip
        print("Erasing chip...")
        erase_cmd = ["pyocd", "erase", "-c", "-t", TARGET_DEVICE]
        erase_result = subprocess.run(erase_cmd, capture_output=True, text=True)
        
        if erase_result.returncode != 0:
            print(f"Erase failed: {erase_result.stderr}")
        
        # Flash HEX file
        print(f"Flashing HEX file: {hex_path}")
        flash_cmd = ["pyocd", "load", hex_path, "-t", TARGET_DEVICE]
        result = subprocess.run(flash_cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print("Flash successful!")
        else:
            print(f"Flash failed, return code: {result.returncode}")
            print(f"Stderr: {result.stderr}")
        
        return result.returncode == 0
    except Exception as e:
        print(f"Error executing PyOCD: {e}")
        return False

def generate_bin_hex(args):
    """Generate BIN and HEX files"""
    print("Generating BIN and HEX files...")
    
    # Use arm-none-eabi-objcopy tool
    objcopy = "arm-none-eabi-objcopy"
    if shutil.which(objcopy):
        try:
            # Generate HEX file
            hex_cmd = [objcopy, "-O", "ihex", args.elf, args.hex]
            subprocess.run(hex_cmd, check=True)
            print(f"Generated HEX file: {args.hex}")
            
            # Generate BIN file
            bin_cmd = [objcopy, "-O", "binary", "-S", args.elf, args.bin]
            subprocess.run(bin_cmd, check=True)
            print(f"Generated BIN file: {args.bin}")
            
            return True
        except subprocess.CalledProcessError as e:
            print(f"Failed to generate files: {e}")
            return False
    else:
        print(f"Tool not found: {objcopy}, cannot generate BIN/HEX files")
        return False

def main():
    parser = argparse.ArgumentParser(description="LSPi Project Flash Script")
    parser.add_argument("--tool", choices=["openocd", "jlink", "pyocd", "auto"], 
                       default="auto", help="Select programming tool (default: auto)")
    parser.add_argument("--elf", default=ELF_FILE, help="ELF file path")
    parser.add_argument("--hex", default=HEX_FILE, help="HEX file path")
    parser.add_argument("--bin", default=BIN_FILE, help="BIN file path")
    parser.add_argument("--generate", action="store_true", help="Generate BIN/HEX files")
    parser.add_argument("--list-tools", action="store_true", help="List available programming tools")
    
    args = parser.parse_args()
    
    # List available tools
    if args.list_tools:
        tools = detect_tools()
        print("Available programming tools:")
        for tool, path in tools.items():
            print(f"  {tool}: {path}")
        if not tools:
            print("  No programming tools found")
        return
    
    # Check build files
    if not check_build_files(args):
        sys.exit(1)
    
    # Generate BIN/HEX files
    if args.generate or not os.path.exists(HEX_FILE):
        if not generate_bin_hex(args):
            print("Warning: Could not generate BIN/HEX files, continuing...")
    
    # Detect available tools
    available_tools = detect_tools()
    
    if not available_tools:
        print("Error: No programming tools found")
        print("Please install one of the following:")
        print("  - OpenOCD (recommended)")
        print("  - SEGGER J-Link")
        print("  - PyOCD")
        sys.exit(1)
    
    print("Available programming tools:")
    for tool in available_tools:
        print(f"  - {tool}")
    
    # Select programming tool
    tool_to_use = args.tool
    if tool_to_use == "auto":
        # Select by priority
        if "openocd" in available_tools:
            tool_to_use = "openocd"
        elif "jlink" in available_tools:
            tool_to_use = "jlink"
        elif "pyocd" in available_tools:
            tool_to_use = "pyocd"
        else:
            print("Error: No available programming tools")
            sys.exit(1)
    
    print(f"Selected programming tool: {tool_to_use}")
    
    # Execute flash
    success = False
    if tool_to_use == "openocd":
        success = flash_with_openocd(args.elf)
    elif tool_to_use == "jlink":
        if not os.path.exists(args.hex):
            print(f"Error: HEX file not found: {args.hex}")
            print("Please use --generate to generate HEX file")
            sys.exit(1)
        success = flash_with_jlink(args.hex)
    elif tool_to_use == "pyocd":
        if not os.path.exists(args.hex):
            print(f"Error: HEX file not found: {args.hex}")
            print("Please use --generate to generate HEX file")
            sys.exit(1)
        success = flash_with_pyocd(args.hex)
    
    if success:
        print("Flash completed!")
        sys.exit(0)
    else:
        print("Flash failed!")
        sys.exit(1)

if __name__ == "__main__":
    main()