import struct, time
from usb import UsbContext, UsbDevice
from devices import Interface, STLinkException, registerInterface

ST_VID, STLINK2_PID = 0x0483, 0x3748

def checkDevice(device):
   return device.VendorId == ST_VID and device.ProductId == STLINK2_PID

DFU_MODE, MASS_MODE, DEBUG_MODE = 0, 1, 2

CORE_RUNNING = 0x80
CORE_HALTED = 0x81

# Commands:
GET_VERSION = 0xf1
DEBUG_COMMAND = 0xf2
DFU_COMMAND = 0xf3
GET_CURRENT_MODE = 0xf5

# dfu commands:
DFU_EXIT = 0x7

# debug commands:
DEBUG_ENTER = 0x20
DEBUG_EXIT = 0x21
DEBUG_ENTER_SWD = 0xa3
DEBUG_GETSTATUS = 0x01
DEBUG_RESETSYS = 0x03
DEBUG_READALLREGS = 0x04
DEBUG_READREG = 0x5
DEBUG_WRITEREG = 0x6
DEBUG_READMEM_32BIT = 0x7
DEBUG_WRITEMEM_32BIT = 0x8
DEBUG_RUNCORE = 0x9
DEBUG_STEPCORE = 0xa

JTAG_WRITEDEBUG_32BIT = 0x35
JTAG_READDEBUG_32BIT = 0x36

# cortex M3
CM3_REG_CPUID = 0xE000ED00

@registerInterface((ST_VID, STLINK2_PID))
class STLink2(Interface):
   """ STlink2 interface implementation. """
   def __init__(self, stlink2=None):
      self.devHandle = None
      if not stlink2:
         context = UsbContext()
         stlink2s = list(filter(checkDevice, context.DeviceList))
         if not stlink2s:
            raise STLinkException('Could not find an ST link 2 interface')
         if len(stlink2s) > 1:
            print('More then one stlink2 found, picking first one')
         stlink2 = stlink2s[0]
      assert isinstance(stlink2, UsbDevice) # Nifty type checking
      assert checkDevice(stlink2)
      self.stlink2 = stlink2
   def __del__(self):
      if self.IsOpen:
         if self.CurrentMode == DEBUG_MODE:
            self.exitDebugMode()
         self.close()
   def __str__(self):
      if self.IsOpen:
         return 'STlink2 device version {0}'.format(self.Version)
      else:
         return 'STlink2 device'
   def open(self):
      if self.IsOpen:
         return
      self.devHandle = self.stlink2.open()
      if self.devHandle.Configuration != 1:
         self.devHandle.Configuration = 1
      self.devHandle.claimInterface(0)

      # First initialization:
      if self.CurrentMode == DFU_MODE:
         self.exitDfuMode()
      if self.CurrentMode != DEBUG_MODE:
         self.enterSwdMode()
      #self.reset()
   def close(self):
      if self.IsOpen:
         self.devHandle.close()
         self.devHandle = None
   @property
   def IsOpen(self):
      return self.devHandle != None
   # modes:
   def getCurrentMode(self):
      cmd = bytearray(16)
      cmd[0] = GET_CURRENT_MODE
      reply = self.send_recv(cmd, 2) # Expect 2 bytes back
      return reply[0]
   CurrentMode = property(getCurrentMode)
   @property
   def CurrentModeString(self):
      modes = {DFU_MODE: 'dfu', MASS_MODE: 'massmode', DEBUG_MODE:'debug'}
      return modes[self.CurrentMode]
   def exitDfuMode(self):
      cmd = bytearray(16)
      cmd[0:2] = DFU_COMMAND, DFU_EXIT
      self.send_recv(cmd)
   def enterSwdMode(self):
      cmd = bytearray(16)
      cmd[0:3] = DEBUG_COMMAND, DEBUG_ENTER, DEBUG_ENTER_SWD
      self.send_recv(cmd)
   def exitDebugMode(self):
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, DEBUG_EXIT
      self.send_recv(cmd)
      
   def getVersion(self):
      cmd = bytearray(16)
      cmd[0] = GET_VERSION
      data = self.send_recv(cmd, 6) # Expect 6 bytes back
      # Parse 6 bytes into various versions:
      b0, b1, b2, b3, b4, b5 = data
      stlink_v = b0 >> 4
      jtag_v = ((b0 & 0xf) << 2) | (b1 >> 6)
      swim_v = b1 & 0x3f
      vid = (b3 << 8) | b2
      pid = (b5 << 8) | b4
      return 'stlink={0} jtag={1} swim={2} vid:pid={3:04X}:{4:04X}'.format(\
         stlink_v, jtag_v, swim_v, vid, pid)
   Version = property(getVersion)
   
   @property
   def ChipId(self):
      return self.read_debug32(0xE0042000)
   @property
   def CpuId(self):
      u32 = self.read_debug32(CM3_REG_CPUID)
      implementer_id = (u32 >> 24) & 0x7f
      variant = (u32 >> 20) & 0xf
      part = (u32 >> 4) & 0xfff
      revision = u32 & 0xf
      return implementer_id, variant, part, revision

   def getStatus(self):
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, DEBUG_GETSTATUS
      reply = self.send_recv(cmd, 2)
      return reply[0]
   Status = property(getStatus)
   @property
   def StatusString(self):
      s = self.Status
      statii = {CORE_RUNNING: 'CORE RUNNING', CORE_HALTED: 'CORE HALTED'}
      if s in statii:
         return statii[s]
      return 'Unknown status'

   def reset(self):
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, DEBUG_RESETSYS
      self.send_recv(cmd, 2)

   # debug commands:
   def step(self):
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, DEBUG_STEPCORE
      self.send_recv(cmd, 2)
   def run(self):
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, DEBUG_RUNCORE
      self.send_recv(cmd, 2)

   # Helper 1 functions:
   def write_debug32(self, address, value):
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, JTAG_WRITEDEBUG_32BIT
      cmd[2:10] = struct.pack('<II', address, value)
      self.send_recv(cmd, 2)
   def read_debug32(self, address):
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, JTAG_READDEBUG_32BIT
      cmd[2:6] = struct.pack('<I', address) # pack into u32 little endian
      reply = self.send_recv(cmd, 8)
      return struct.unpack('<I', reply[4:8])[0]
   def write_reg(self, reg, value):
      cmd = bytearray(16)
      cmd[0:3] = DEBUG_COMMAND, DEBUG_WRITEREG, reg
      cmd[3:7] = struct.pack('<I', value)
      r = self.send_recv(cmd, 2)
   def read_reg(self, reg):
      cmd = bytearray(16)
      cmd[0:3] = DEBUG_COMMAND, DEBUG_READREG, reg
      reply = self.send_recv(cmd, 4)
      return struct.unpack('<I', reply)[0]
   def read_all_regs(self):
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, DEBUG_READALLREGS
      reply = self.send_recv(cmd, 84)
      fmt = '<' + 'I' * 21 # unpack 21 register values
      return list(struct.unpack(fmt, reply))
   def write_mem32(self, address, content):
      assert len(content) % 4 == 0
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, DEBUG_WRITEMEM_32BIT
      cmd[2:8] = struct.pack('<IH', address, len(content))
      self.send_recv(cmd)
      self.send_recv(content)
   def read_mem32(self, address, length):
      assert length % 4 == 0
      cmd = bytearray(16)
      cmd[0:2] = DEBUG_COMMAND, DEBUG_READMEM_32BIT
      cmd[2:8] = struct.pack('<IH', address, length)
      reply = self.send_recv(cmd, length) # expect memory back!
      return reply

   # Helper 2 functions:
   def send_recv(self, tx, rxsize=0):
      """ Helper function that transmits and receives data in bulk mode. """
      # TODO: we could use here the non-blocking libusb api.
      tx = bytes(tx)
      #assert len(tx) == 16
      self.devHandle.bulkWrite(2, tx) # write to endpoint 2
      if rxsize > 0:
         return self.devHandle.bulkRead(1, rxsize) # read from endpoint 1

if __name__ == '__main__':
   # Test program
   sl = STLink2()
   sl.open()
   sl.reset()
   print('version:', sl.Version)
   print('mode before doing anything:', sl.CurrentModeString)
   if sl.CurrentMode == DFU_MODE:
      sl.exitDfuMode()
   sl.enterSwdMode()
   print('mode after entering swd mode:', sl.CurrentModeString)

   i = sl.ChipId
   print('chip id: 0x{0:X}'.format(i))
   print('cpu: {0}'.format(sl.CpuId))

   print('status: {0}'.format(sl.StatusString))

   # test registers:
   sl.write_reg(0, 0xdeadbeef)
   sl.write_reg(1, 0xcafebabe)
   sl.write_reg(2, 0xc0ffee)
   sl.write_reg(3, 0x1337)
   sl.write_reg(5, 0x1332)
   sl.write_reg(6, 0x12345)
   assert sl.read_reg(3) == 0x1337
   assert sl.read_reg(5) == 0x1332
   assert sl.read_reg(6) == 0x12345
   regs = sl.read_all_regs()
   for i in range(len(regs)):
      print('R{0}=0x{1:X}'.format(i, regs[i]))

   sl.exitDebugMode()
   print('mode at end:', sl.CurrentModeString)

   sl.close()
   print('Test succes!')