import pylink
import os

# The serial number of your J-Link probe. 
# This is useful if you have multiple J-Link devices connected.
# You can find the serial number using the 'pylink.jlink_list()' function or the J-Link Commander.
JLINK_SERIAL_NO = None # Or provide your serial number here

# The path to the firmware HEX file you want to flash.
# Example: Using a pre-compiled hex from the nRF Connect SDK.
FIRMWARE_FILE_PATH = os.path.join("c:\Temp", "ModemShellRtt.hex")

# Define the device name. For the Thingy:91's nRF9160 SiP, this is 'nRF9160_xxAA'.
TARGET_DEVICE = 'nRF9160_xxAA'

# The address where the firmware should be flashed. For the nRF9160, this is typically 0.
FLASH_ADDRESS = 0x0

def connect_and_flash():
    """
    Connects to a J-Link probe, flashes firmware to a Thingy:91, and resets the device.
    """
    try:
        # Create a JLink object
        jlink = pylink.JLink()

        # Open a connection to the J-Link. 
        # The serial number is optional but recommended for clarity.
        if JLINK_SERIAL_NO:
            print(f"Connecting to J-Link probe with serial number {JLINK_SERIAL_NO}...")
            jlink.open(serial_no=JLINK_SERIAL_NO)
        else:
            print("Connecting to the first available J-Link probe...")
            jlink.open()
        
        # Connect to the target device (Thingy:91's nRF9160)
        print(f"Connecting to target device: {TARGET_DEVICE}...")
        jlink.connect(TARGET_DEVICE, verbose=True)

        # Flash the firmware file
        print(f"Flashing firmware from {FIRMWARE_FILE_PATH}...")
        jlink.flash(FIRMWARE_FILE_PATH, FLASH_ADDRESS)

        # Reset the device to run the new firmware
        print("Resetting the target device...")
        jlink.reset()

        print("Flashing successful!")

    except pylink.errors.JLinkException as e:
        print(f"An error occurred: {e}")
        if jlink and jlink.connected():
            jlink.close()
    finally:
        # Always close the connection in the 'finally' block to ensure cleanup
        if jlink and jlink.connected():
            jlink.close()
            print("Connection to J-Link closed.")

if __name__ == "__main__":
    if not os.path.exists(FIRMWARE_FILE_PATH):
        print(f"Error: Firmware file not found at {FIRMWARE_FILE_PATH}")
    else:
        connect_and_flash()