#!/usr/bin/env lua
rs232 = require("luars232")
zlib = require("zlib")
bit = require("bit")

DEBUG = false

-- Per-platform configuration
PLATFORMS = {
  ['olimex_evb']  = {  ['port'] = '/dev/ttyS1',    ['baud'] = rs232.RS232_BAUD_460800, ['zlib'] = 'lzlib',    ['reset'] = {'gpio', 13, 11},              ['backdoor'] = {11, 0x0}},
  ['prototype_v1']  = {['port'] = '/dev/ttyS0',    ['baud'] = rs232.RS232_BAUD_460800, ['zlib'] = 'lzlib',    ['reset'] = {'gpio', 21, 'node_backdoor'}, ['backdoor'] = {8, 0x0}},
  ['prototype_v2']  = {['port'] = '/dev/ttyS1',    ['baud'] = rs232.RS232_BAUD_460800, ['zlib'] = 'lzlib',    ['reset'] = {'gpio', 21, 'node_backdoor'}, ['backdoor'] = {11, 0x0}},
  ['prototype_v4']  = {['port'] = '/dev/ttyS1',    ['baud'] = rs232.RS232_BAUD_460800, ['zlib'] = 'lzlib',    ['reset'] = {'gpio', 21, 'node_backdoor'}, ['backdoor'] = {11, 0x0}},
  ['prototype_v5']  = {['port'] = '/dev/ttyS1',    ['baud'] = rs232.RS232_BAUD_460800, ['zlib'] = 'lzlib',    ['reset'] = {'gpio', 21, 'node_backdoor'}, ['backdoor'] = {11, 0x0}},
  ['pc_xds100v3'] = {  ['port'] = '/dev/ttyUSB1',  ['baud'] = rs232.RS232_BAUD_460800, ['zlib'] = 'lua-zlib', ['reset'] = {'manual'},                    ['backdoor'] = {11, 0x0}}
}
PLATFORMS['pc'] = PLATFORMS['pc_xds100v3']
PLATFORMS['olimex'] = PLATFORMS['olimex_evb']
PLATFORMS['v1'] = PLATFORMS['prototype_v1']
PLATFORMS['v2'] = PLATFORMS['prototype_v2']
PLATFORMS['v4'] = PLATFORMS['prototype_v4']
PLATFORMS['v5'] = PLATFORMS['prototype_v5']
PLATFORMS['cherry'] = PLATFORMS['prototype_v5']

local out = io.stderr

function hexdump(s)
  if s == nil then return "" end
  res = ""
  for c in s:gmatch"." do
    res = res .. string.format('%02x ', c:byte())
  end
  return res
end

function bytes_to_string(bs)
  if bs == nil then return '' end
  local s = {}
  for _,b in ipairs(bs) do table.insert(s, string.char(b)) end
  return table.concat(s)
end

function string_to_bytes(s)
  if s == nil then return {} end
  bs = {}
  for c in s:gmatch"." do table.insert(bs, c:byte()) end
  return bs
end

function debug(msg)
  if DEBUG then print(msg) end
end

function sleep(ms)
  local ntime = os.clock() + ms/1000
  repeat until os.clock() > ntime
end

function parse_int(data)
  -- data should be array of length 4
  return bit.band(0xffffffff, bit.bor(
    bit.lshift(data[1], 24),
    bit.lshift(data[2], 16),
    bit.lshift(data[3], 08),
    bit.lshift(data[4], 00)
  ));
end

function parse_int_le(data)
  -- data should be array of length 4
  return bit.band(0xffffffff, bit.bor(
    bit.lshift(data[1], 00),
    bit.lshift(data[2], 08),
    bit.lshift(data[3], 16),
    bit.lshift(data[4], 24)
  ));
end

table.sum = function(t)
  function sumf(a, ...) return a and a + sumf(...) or 0 end
  return sumf(unpack(t))
end

table.size = function(t)
  i=0
  for k,v in pairs(t) do
    i = i+1
  end
  return i
end

table.merge = function(...)
  res = {}
  for i, t in ipairs(arg) do
    for _, v in pairs(t) do table.insert(res, v) end
  end
  return res
end

table.slice = function(t, s, e)
  assert(s <= #t and e <= #t)
  res = {}
  for i=s,e do
    table.insert(res, t[i])
  end
  return res
end

string.trim = function(s)
  return (s:gsub("^%s*(.-)%s*$", "%1"))
end

gpio = {}

gpio.export = function(number, direction)
  res = os.execute(string.format([[
        cd /sys/class/gpio/ &&
        ([ -d gpio%d ] || echo %d > export) &&
        echo '%s' > gpio%d/direction
       ]], number, number, direction, number))
  if res ~= 0 then return nil, string.format("%d", res) end
  return string.format("gpio%d", number)
end

gpio.unexport = function(name)
  number = string.match(name, "^gpio(%d+)$")
  if number == nil then return false end
  res = os.execute(string.format([[
        echo %d > /sys/class/gpio/unexport
       ]], number))
  if res ~= 0 then return false, string.format("%d", res) end
  return true
end

gpio.set = function(name, value)
  res = os.execute(string.format([[
        echo %d > /sys/class/gpio/%s/value
       ]], value, name))
  if res ~= 0 then return false, string.format("%d", res) end
  return true
end

local function Bootloader(platform, port_override)
  local self = {}

  -- Timeout waiting for first ACK
  START_TIMEOUT = 5000 -- [ms]
  -- Timeout for simple commands
  SHORT_TIMEOUT = 100 -- [ms]
  -- Timeout for commands accessing whole flash
  LONG_TIMEOUT = 5000 -- [ms]
  -- Sleep time for GPIO manipulation
  GPIO_TIMEOUT = 100 -- [ms]
  -- ARM needs some time to get ready to receive after sending.
  -- This causes random failures when programming cpu (i.e. PC) is too fast.
  SEND_AFTER_RECV_TIMEOUT = 10 -- [ms]

  ACK = {0x00, 0xcc}

  -- public

  function self.open_uart()
    self.settings = PLATFORMS[platform]
    if port_override ~= nil then
      self.settings.port = port_override
    end
    if not self.settings then return false, string.format("Unsupported platform: %s", platform) end

    res, msg = save_uart()
    if not res then return res, msg end

    e, self.rs = rs232.open(self.settings.port)
    local res, msg = check_uart(e)
    if not res then return res, msg end

    local res, msg = check_uart(self.rs:set_baud_rate(self.settings.baud))
    if not res then return res, msg end
    local res, msg = check_uart(self.rs:set_data_bits(rs232.RS232_DATA_8))
    if not res then return res, msg end
    local res, msg = check_uart(self.rs:set_parity(rs232.RS232_PARITY_NONE))
    if not res then return res, msg end
    local res, msg = check_uart(self.rs:set_stop_bits(rs232.RS232_STOP_1))
    if not res then return res, msg end
    local res, msg = check_uart(self.rs:set_flow_control(rs232.RS232_FLOW_OFF))
    if not res then return res, msg end

    return true, string.format("UART %s opened.", self.settings.port)
  end

  function self.reset_node()
    mode = self.settings.reset[1]

    if mode == "gpio" then
      if #self.settings.reset ~= 3 then
        return false, "Invalid reset settings"
      end
      r = self.settings.reset[2]
      s = self.settings.reset[3]

      if(type(r) == "number") then
        r, msg = gpio.export(r, 'out')
        if not r then return false, string.format("Error exporting GPIO (reset): %s", msg) end
        unexport_r = true
      end
      if(type(s) == "number") then
        s, msg = gpio.export(s, 'out')
        if not s then return false, string.format("Error exporting GPIO (backdoor): %s", msg) end
        unexport_s = true
      end

      bd_act = self.settings.backdoor[2]

      res = gpio.set(r, 0)
      if not res then return false, "Error manipulating GPIOs [1]." end
      res = gpio.set(s, bd_act)
      if not res then return false, "Error manipulating GPIOs [2]." end
      sleep(GPIO_TIMEOUT)

      res = gpio.set(r, 1)
      if not res then return false, "Error manipulating GPIOs [3]." end

      sleep(GPIO_TIMEOUT)
      res = gpio.set(s, 0x1 - bd_act)
      if not res then return false, "Error manipulating GPIOs [4]." end

      if unexport_s then
        res = gpio.unexport(s)
        if not res then print("Warning: Unexporting gpio for backdoor failed.") end
      end
      if unexport_r then
        res = gpio.unexport(r)
        if not res then print("Warning: Unexporting gpio for reset failed.") end
      end

      return true, "Blinked GPIOs to reset node."
    elseif mode == "manual" then
      out:write("Please reset node to bootloader (RESET+SELECT; release RESET; release SELECT) and press enter.")
      pcall(io.read())
      return true, "Hopefully user reset the node"
    elseif mode == nil then
      return false, "Unspecified reset mode."
    else
      return false, string.format("Unsupported reset mode: %s", mode)
    end
  end

  function self.sync()

    local res, msg = send_bytes({0x55, 0x55})
    if not res then return res, msg end

    -- ensure we wait at least SHORT_TIMEOUT for the answer
    -- we need to flush any garbage collected due to line hazzard
    -- and then the ACK, after which the line is pulled up
    repeat
       _, _, size = rs_read(1024, SHORT_TIMEOUT)
    until(size == 0)

    -- ping node and check for ACK
    local res, msg = send_bytes({0x03,0x20,0x20})
    if not res then return res, msg end
    local res, msg = wait_for_ack(START_TIMEOUT)
    if not res then return res, msg end
    return true, "Synced with bootloader."
  end

  function self.check_status()
    res, msg, data = execute_command({0x23}, SHORT_TIMEOUT, true)
    if not res then return res, msg end

    if #data ~= 1 or data[1] ~= 0x40 then
      return false, string.format("Unexpected status: %s", hexdump(bytes_to_string(data)))
    end
    return true, "Status: success"
  end

  function self.identify()
    res, msg, data1 = execute_command({0x28}, SHORT_TIMEOUT, true)
    if not res then return res, msg end

    FLASH_SIZE_CFG = 0x4003002C
    ACCESS_32BIT = 0x1
    SECTOR_SIZE = 4096
    res, msg, data2 = execute_command(table.merge({0x2A}, encode_address(FLASH_SIZE_CFG), {ACCESS_32BIT, 1}), SHORT_TIMEOUT, true)
    if not res then return res, msg end

    self.flash_size = data2[1] * SECTOR_SIZE

    res, msg = self.check_status()
    if not res then return res, msg end

    return true, string.format("Device identified as: %s, flash size: 0x%x", hexdump(bytes_to_string(data1)), self.flash_size)
  end

  function self.cleanup()
    if self.rs ~= nil then
      local res, msg = check_uart(self.rs:close())
      local res2, msg2 = restore_uart()
      if not res2 then print(string.format("Warning: %s", msg2)) end
      if not res then return res, msg end
    end
    return true, "Closed bootloader connection."
  end

  function self.erase()
    res, msg, _ = execute_command({0x2c}, LONG_TIMEOUT, false)
    if not res then return res, msg end
    res, msg = self.check_status()
    if not res then return res, msg end
    return true, "Erased flash."
  end

  function self.write_image(segments)
    for addr,data in pairs(segments) do
      self.write_segment(addr, data)
      if not res then return res, msg end
    end
    return true, string.format("Wrote all %d segments.", table.size(segments))
  end

  function self.write_segment(addr, data)
    MAX_DATA_PER_PACKET = 252
    MAX_RETRIES = 5
    data = string_to_bytes(data)

    res, msg, _ = execute_command(table.merge(
      {0x21},
      encode_address(addr),
      encode_address(#data)
    ), LONG_TIMEOUT, false)
    if not res then return res, msg end

    self.check_status()
    if not res then return res, msg end

    for i=1,#data,MAX_DATA_PER_PACKET do
      packet = table.slice(data, i, math.min(#data, i+MAX_DATA_PER_PACKET-1))
      for retry = 1,MAX_RETRIES do
        res, msg, _ = execute_command(table.merge(
          {0x24},
          packet
        ), LONG_TIMEOUT, false)
        if res then break end
      end
      if not res then return res, msg end
    end

    self.check_status()
    if not res then return res, msg end

    return true, string.format("Written %d bytes of image to %08x", #data, addr)
  end

  function self.check_crc(segments)
    for addr,data in pairs(segments) do
      self.check_segment(addr, data)
      if not res then return res, msg end
    end
    return true, string.format("CRC correct for all %d segments.", table.size(segments))
  end

  function self.check_segment(addr, data)
    if platform.zlib == 'lzlib' then
      crc = bit.band(math.floor(zlib.crc32(0, data)), 0xffffffff)
    elseif platform.zlib == 'lua-zlib' then
      crc = bit.band(math.floor(zlib.crc32()(data)), 0xffffffff)
    else
      return false, string.format("Unknown zlib library")
    end
    data = string_to_bytes(data)
    len = #data
    res, msg, data = execute_command(table.merge(
      {0x27},
      encode_address(addr),
      encode_address(len),
      encode_address(0)
    ), LONG_TIMEOUT, true)
    if not res then return res, msg end

    self.check_status()
    if not res then return res, msg end

    got_crc = parse_int(data)
    if crc ~= got_crc then
      return false, string.format("CRC mismatch for segment at %08x, len %x: data %08x, flash %08x", addr, len, crc, got_crc)
    end
    return true, string.format("CRC correct for %d byte segment at %08x: %08x", len, addr, crc)
  end

  function self.launch()
    res, msg, _ = execute_command({0x25}, LONG_TIMEOUT, false)
    if not res then return res, msg end
    return true, string.format("Launched new image.")
  end

  -- private

  function execute_command(data, timeout, receive)
    res, msg = send_command(data)
    if not res then return res, msg end
    local res, msg = wait_for_ack(timeout)
    if not res then return res, msg end
    if receive then
      res, msg, data = receive_packet(timeout)
      if not res then return res, msg end
      res, msg = send_ack()
      if not res then return res, msg end
    end
    return true, "Executed command", data
  end

  function encode_address(addr)
    return {
      bit.band(bit.rshift(addr, 24), 0xff),
      bit.band(bit.rshift(addr, 16), 0xff),
      bit.band(bit.rshift(addr, 08), 0xff),
      bit.band(bit.rshift(addr, 00), 0xff),
    }
  end

  function send_command(bs)
    len = #bs + 2
    if len > 0xff then
      return false, string.format("Data too long: %d", len)
    end
    chk = table.sum(bs) % 256

    local res, msg = send_bytes(table.merge({len, chk}, bs))
    if not res then return res, msg end
    return true, string.format("Sent cmd %02x of total length %d", bs[1], len)
  end

  function send_ack()
    local res, msg = send_bytes({0x00, 0xcc})
    if not res then return res, msg end
    return true, "Sent ACK."
  end

  function receive_packet(timeout)
    local HDR_LEN = 2
    local e, data, size = rs_read(HDR_LEN, timeout)
    if size ~= HDR_LEN then
      return false, string.format("Timed out waiting for packet header, got only '%s'", hexdump(data))
    end
    local res, msg = check_uart(e)
    if not res then return res, msg end
    local len, chk = unpack(string_to_bytes(data))
    len = len - HDR_LEN
    local e, data, size = rs_read(len, timeout)
    if size ~= len then
      return false, string.format("Timed out waiting for packet body, got only '%s'", hexdump(data))
    end
    data = string_to_bytes(data)
    local data_chk = table.sum(data) % 256
    if chk ~= data_chk then
      return false, string.format("Checksum mismatch, expected %02x, got %02x on data: %s", chk, data_chk, hexdump(bytes_to_string(data)))
    end

    return true, "Received packet", data
  end

  function rs_write(data)
    local e, size = self.rs:write(data)
    debug(string.format("UART write [%d/%d]: %s", size, data:len(), hexdump(data)))
    return e, size
  end

  function rs_read(max_size, timeout)
    local e, data, size = self.rs:read(max_size, timeout)
    debug(string.format("UART read [%d/%d]: %s", size, max_size, hexdump(data)))
    return e, data, size
  end

  function check_uart(e)
    if e ~= rs232.RS232_ERR_NOERROR then
      return false, string.format("UART %s error: %s", self.settings.port, rs232.error_tostring(e))
    end
    return true, "UART No error"
  end

  function send_bytes(bs)
    sleep(SEND_AFTER_RECV_TIMEOUT)
    local e, written = rs_write(bytes_to_string(bs))
    local len = table.getn(bs)
    local res, msg = check_uart(e)
    if not res then return res, msg end
    if written ~= len then
      return false, string.format("Short write, tried %d, written %d", len, written)
    end
    return true, string.format("Sent %d bytes.", len)
  end

  function wait_for_ack(timeout)
    sleep(SEND_AFTER_RECV_TIMEOUT)
    local e, data, size = rs_read(table.getn(ACK), timeout)
    if size ~= table.getn(ACK) then
      return false, string.format("Timed out waiting for ACK, got only '%s'", hexdump(data))
    end
    local res, msg = check_uart(e)
    if not res then return res, msg end
    if data ~= bytes_to_string(ACK) then
      return false, string.format("Expected ACK, got '%s'", hexdump(data))
    end
    return true, "Got ACK"
  end

  function save_uart()
    f = io.popen(string.format("bash -c 'stty -F %s -g 2>&1 && echo CMD_OK'", self.settings.port), 'r')
    if not f then return false, "Save UART: Popen failed" end
    s = f:read('*a')
    f:close()
    if not s then return false, "Save UART: Read failed" end
    s = string.trim(s)
    uart = string.match(s, '(.*)\nCMD_OK')
    if not uart then
      return false, string.format("Saving UART settings failed: %s", s)
    end
    self.saved_uart = uart
    return true, "Saved UART settings"
  end

  function restore_uart()
    res = os.execute(string.format("stty -F %s %s", self.settings.port, self.saved_uart))
    if res ~= 0 then
      return false, "UART restore failed."
    end
    return true, "UART restored"
  end

  return self
end

function Binary(file, platform)
  local self = {}
  self.file = file
  self.settings = PLATFORMS[platform]
  self.segments = {} -- pairs address, data

  function self.load_ihex()
    -- Load Intel HEX file

    local f, msg = io.open(self.file, 'r')
    if f == nil then
      return false, string.format("Cannot open file %s: %s", self.file, msg)
    end

    upper_addr = 0

    while true do
      line = f:read("*line")
      if line == nil then break end
      err, addr, kind, data = parseIHex(line)
      if err ~= nil then
        return false, string.format("Cannot parse line (%s): %s", err, line)
      end
      if     kind == 0 then                             -- real data
        self.segments[upper_addr + addr] = data
      elseif kind == 1 then break                       -- end of file
      elseif kind == 2 then                             -- extended segment address
        a, b = string.byte(data,1,2)
        -- "The segment address from the most recent 02 record is multiplied by 16 and added to each subsequent data record address" - https://en.wikipedia.org/wiki/Intel_HEX
        upper_addr = (bit.lshift(a, 8) + b) * 16
      elseif kind == 3 then do end                      -- 20bit EIP, ignore
      elseif kind == 4 then                             -- extended linear address
        a, b = string.byte(data,1,2)
        -- "The two encoded, big endian data bytes specify the upper 16 bits of the 32 bit absolute address for all subsequent type 00 records" - https://en.wikipedia.org/wiki/Intel_HEX
        upper_addr = bit.lshift(a, 24) + bit.lshift(b, 16)
      elseif kind == 5 then do end                      -- 32bit EIP, ignore
      else
       return false, string.format("Unsupported record type: %02d", kind)
      end
    end

    for a, d in pairs(self.segments) do
      if d ~= nil then
        while self.segments[a + #d] ~= nil do
          a2, d2 = a + #d, self.segments[a + #d]
          self.segments[a] = d .. d2
          self.segments[a2] = nil
          d = self.segments[a]
        end
      end
    end

    return true, "Opened image file (Intex HEX)"
  end

  function self.load_bin()
    -- Load raw binary file
    local f, msg = io.open(self.file, 'r')
    if f == nil then
      return false, string.format("Cannot open file %s: %s", self.file, msg)
    end
    addr = 0
    data = f:read("*all")
    f:close()
    if data == nil then
      return false, "Failed to read"
    end
    self.segments[addr] = data
    return true, "Opened image file (raw binary)"
  end

  function self.load()
    loader = nil
    if string.find(self.file, '%.bin$') then
      loader = self.load_bin
    end
    if string.find(self.file, '%.hex$') then
      loader = self.load_ihex
    end
    if string.find(self.file, '%.ihex$') then
      loader = self.load_ihex
    end

    if loader == nil then
      return false, string.format("Cannot detect file type by extension")
    end

    res, msg = loader()
    if not res then return res, msg end
    return true, string.format("Opened image file with %d segments.", table.size(self.segments))
  end

  function self.find_segment(addr)
    for k,v in pairs(self.segments) do
      if(k <= addr and addr+3 <= k+#v) then
        return k, v
      end
    end
  end

  function self.check_value(offset, value, mask, msg)
    addr = self.ccfg_base + offset
    debug(string.format("Checking at %08x", addr))
    seg_base, seg = self.find_segment(addr)
    if not seg then return true, "Address not written" end
    write_val = table.slice(string_to_bytes(seg), addr-seg_base+1, addr-seg_base+4)
    write_val = parse_int_le(write_val)
    debug(string.format("Writing %s, expected %s, mask %s", bit.tohex(write_val), bit.tohex(value), bit.tohex(mask)))
    if bit.band(mask, bit.bxor(value, write_val)) ~= 0 then
      return false, string.format("Bad value at offset 0x%x, expected %s, got %s (mask %s) - %s", offset, bit.tohex(value), bit.tohex(write_val), bit.tohex(mask), msg)
    end
    return true, "Value as expected"
  end

  function self.sanity_check(flash_size)
    if(flash_size ~= 32*1024 and flash_size ~= 64*1024 and flash_size ~= 128*1024) then
      return false, string.format("Unexpected flash size: %d (0x%08x)", flash_size, flash_size)
    end

    max_addr = 0
    for k,v in pairs(self.segments) do
      max_addr = math.max(max_addr, k + #v)
    end
    if max_addr > flash_size then
      return false, string.format("Image too big - largest address: %d, flash size: %d", max_addr, flash_size)
    end

    self.ccfg_base = flash_size - 0x100
    -- See 9.1.1 in TI's swcu117f.pdf

    res, msg = self.check_value(0xd8, 0xc5ff00c5, 0xfffe00ff, "This image may disable bootloader!")
    if not res then return res, msg end

    res, msg = self.check_value(0xec, 0x00000000, 0xffffffff, "Incorrect IMAGE_VALID field value")
    if not res then
      print(string.format("Warning: %s", msg))
      print("Continuing, as this will not prevent further reprogramming")
    else
      -- Only check bootloader backdoor settings if image is valid, as
      -- otherwise uC will enter bootloader unconditionally anyway
      pin = self.settings.backdoor[1]
      active = self.settings.backdoor[2]
      value = bit.bor(0x00fe0000, bit.bor(bit.lshift(active, 16), bit.lshift(pin, 8)))
      res, msg = self.check_value(0xd8, value, 0x0001ff00, "This image sets different backdoor pin, may not be programmable on this platform again!")
      if not res then return res, msg end
    end

    res, msg = self.check_value(0xdc, 0x00000101, 0x00000101, "This image may disable mass erase")
    if not res then return res, msg end

    res, msg = self.check_value(0xe0, 0x000000c5, 0x000000ff, "TI failure analysis disabled")
    if not res then return res, msg end
    res, msg = self.check_value(0xe4, 0x00c5c5c5, 0x00ffffff, "Some JTAG TAPs disabled in DAP_0")
    if not res then return res, msg end
    res, msg = self.check_value(0xe8, 0x00c5c5c5, 0x00ffffff, "Some JTAG TAPs disabled in DAP_1")
    if not res then return res, msg end

    res, msg = self.check_value(0xf0, 0xffffffff, 0xffffffff, "Memory protection enabled")
    if not res then return res, msg end
    res, msg = self.check_value(0xf4, 0xffffffff, 0xffffffff, "Memory protection enabled")
    if not res then return res, msg end
    res, msg = self.check_value(0xf8, 0xffffffff, 0xffffffff, "Memory protection enabled")
    if not res then return res, msg end
    res, msg = self.check_value(0xfc, 0xffffffff, 0xffffffff, "Memory protection enabled")
    if not res then return res, msg end

    return true, "Sanity checks passed"
  end

  function hexdump_to_string(h)
    s = ""
    for i = 1,#h,2 do
      s = s .. string.char(tonumber(string.sub(h,i,i+1), 16))
    end
    return s
  end

  function parseIHex(line)
    len, addr, kind, data, checksum = string.match(string.trim(line), "^:(%x%x)(%x%x%x%x)(%x%x)(%x*)(%x%x)$")
    if len == nil then return "Does not match regexp" end
    len = tonumber(len, 16)
    addr = tonumber(addr, 16)
    kind = tonumber(kind, 16)
    checksum = tonumber(checksum, 16)
    if len == nil or addr == nil or kind == nil or checksum == nil then return "Cannot convert header fields to numbers" end
    data = hexdump_to_string(data)
    if len ~= #data then return "Invalid length" end
    checksum = checksum + len + kind + bit.band(addr, 0xff) + bit.rshift(addr, 8)
    for c in data:gmatch"." do
      checksum = (checksum + string.byte(c)) % 256
    end
    if bit.band(checksum, 0xff) ~= 0 then return string.format("Bad checksum (%s)", bit.tohex(checksum, 2)) end
    return nil, addr, kind, data
  end

  return self
end

function StateMachine(platform, file_name, port_override, retry)
  local state = "INIT"
  local b = Bootloader(platform, port_override)
  local bin = Binary(file_name, platform)

  function set_state(state_ok, state_bad, res, msg)
    old_state = state
    if res then
      state = state_ok
    else
      state = state_bad
    end
    out:write(string.format("%-10s -> %-10s, because of operation %s: %s\n", old_state, state, res and "success" or "failure", msg))
  end

  while true do

    if state == "FATAL" then
      os.exit(1)
    end

    if state == "INIT" then
      set_state("LOADED", "FATAL", bin.load())
    end

    if state == "RESTART" then
      if retry then
        sleep(3000) -- Prevent from looping too fast.
        set_state("LOADED", "FATAL", b.cleanup())
      else
        set_state("FATAL", "FATAL", b.cleanup())
      end
    end

    if state == "LOADED" then
      set_state("OPENED", "RESTART", b.open_uart())
    end

    if state == "OPENED" then
      set_state("RESET", "RESTART", b.reset_node())
    end

    if state == "RESET" then
      set_state("SYNCED", "RESTART", b.sync())
    end

    if state == "SYNCED" then
      set_state("IDENTIFIED", "RESTART", b.identify())
    end

    if state == "IDENTIFIED" then
      set_state("SANE", "FATAL", bin.sanity_check(b.flash_size))
    end

    if state == "SANE" then
      set_state("ERASED", "RESTART", b.erase())
    end

    if state == "ERASED" then
      set_state("WRITTEN", "RESTART", b.write_image(bin.segments))
    end

    if state == "WRITTEN" then
      set_state("LAUNCH", "RESTART", b.check_crc(bin.segments))
    end

    if state == "LAUNCH" then
      set_state("FINISHED", "RESTART", b.launch())
    end

    if state == "FINISHED" then
      b.cleanup()
      break
    end
  end
end

function bootloader()
  if #arg < 1 then
    out:write("Usage: ./bootloader image_file [platform] [port]\n")
    os.exit(1)
  end
  file_name = arg[1]
  platform = arg[2] or 'cherry'
  port = arg[3]
  StateMachine(platform, file_name, port, true)
end

function sanity_check()
  if #arg < 1 then
    out:write("Usage: ./sanity_check image_file [platform] [flash_size]\n")
    os.exit(1)
  end
  file_name = arg[1]
  platform = arg[2] or 'cherry'
  flash_size = arg[3] or 128*1024
  bin = Binary(file_name, platform)
  res, msg = bin.load()
  if not res then
    print(msg)
    os.exit(1)
  end

  res, msg = bin.sanity_check(flash_size)
  print(msg)
  if not res then
    os.exit(1)
  end
  os.exit(0)
end

if string.match(arg[0], '.*sanity[-_]check$') then
 sanity_check()
else
 bootloader()
end
