#!/bin/python

import sys
import os
import re

class BaseFlasher:
    def flash(self, micro, filename):
        pass
    def getSuffix(self):
        pass
    def getFamilyName(self):
        pass

class AVRFlasher:
    def __init__(self):
       self.microDict = False

    def flash(self, micro, filename):
        """
        Flashes AVR microcontroller
        @param micro Microcontroller name understood by GCC
        @param filename Path to hex file to flash
        @returns True on success
        """
        #avrdude -c avrftdi -pusb162 -U flash:w:out/avrboard-cnc.hex
        if not self.microDict:
            d = {}
            f = os.popen("avrdude -cavrftdi 2>&1")
            for line in f:
                if line.find('=') == -1:
                    continue
                parts = line.split('=')
                if len(parts) != 2:
                    continue
                m = parts[1].strip().lower()
                code = parts[0].strip()
                d[m] = code
            f.close()
            self.microDict = d
        if micro not in self.microDict:
            print("Unknown microcontroller: " + micro + " for avrdude")
            return False
        command = "avrdude -cavrftdi -p" + self.microDict[micro] + " -U flash:w:" + filename
        print "Executing '" + command + "'"
        os.execvp("avrdude", ['avrdude', '-cavrftdi', '-p', self.microDict[micro], '-U', 'flash:w:' + filename])
        return True
    def getSuffix(self):
        return 'avr.hex'
    def getFamilyName(self):
        return 'avr'

class STM32Flasher:
    def flash(self, micro, filename):
        """
        Flashes STM32 microcontroller. Only STM32F4Discovery
        board is supported at the moment
        @param micro Microcontroller name (ignored currently)
        @param filename Path to hex file to flash
        @returns True on success
        """
        args = [
            'openocd',
            '-c', 'source [find interface/stlink-v2.cfg]',
            '-c', 'source [find target/stm32f4x_stlink.cfg]',
            '-c', 'reset_config srst_only srst_nogate',
            '-c', 'init',
            '-c', 'reset halt',
            '-c', 'flash write_image erase ' + filename,
            '-c', 'reset',
            '-c', 'shutdown',
        ]
        print "Executing '" + ' '.join(args) + "'"
        os.execvp(args[0], args)
        return True
    def getSuffix(self):
        return 'stm32.elf'
    def getFamilyName(self):
        return 'stm32'

flashersList = [AVRFlasher(), STM32Flasher()]
flashers = {}
for f in flashersList:
    flashers[f.getFamilyName()] = f

def searchFirmware(dirs):
    """
    Searches directories for files that can be used
    to flash microcontrollers
    @param dirs List of directories to search in
    @returns list of filenames
    """
    if not hasattr(searchFirmware, 'p'):
        suffixes = [f.getSuffix() for f in flashersList]
        suffixes = '|'.join(suffixes)
        searchFirmware.p = re.compile('.*\.[^.]*\.(' + suffixes + ')$')
    lst = []
    for d in dirs:
        if not os.path.isdir(d):
            continue
        l = os.listdir(d)
        for i in l:
            lst.append(d + '/' + i)

    return [i for i in lst if searchFirmware.p.match(i)]

def filterListByPattern(lst, pattern):
    """
    Filters list by pattern
    @param lst Input list
    @param pattern Match this string
    @returns List of items of lst that have
             pattern as substring
    """
    return [i for i in lst if i.find(pattern) != -1]

def flashFile(filename):
    """
    Flash hex file to microcontroller
    @param filename
    @returns True on success
    """
    parts = filename.split('/')
    parts = parts[-1].split('.')
    # name.micro.architecture.hex
    if len(parts) < 4:
        print("Wrong filename. Should be name.architecture.micro.hex")
        return False
    family = parts[-2]
    micro = parts[-3]

    if not family in flashers:
        print("Unsupported architecture: " + architecture)
        return False
    return flashers[family].flash(micro, filename)

if len(sys.argv) > 2:
    print("Usage: indemic_flash [part_of_filename]")
    print("       File will be searched in current directory and in")
    print("       _build directory. Only *.*.*.hex files are considered")
    sys.exit(1)

dirs = ['.', '_build']
arg = None
trytoflash = None
pattern = None

if len(sys.argv) == 2:
    arg = sys.argv[1]
    if os.path.isfile(arg):
        trytoflash = arg
    elif os.path.isdir(arg):
        dirs.append(arg)
    else:
        pattern = arg

if trytoflash is None:
    listOfFirmware = searchFirmware(dirs)
    if pattern:
        listOfFirmware = filterListByPattern(listOfFirmware, pattern)
    if len(listOfFirmware) == 0:
        print("Found nothing matching.")
        sys.exit(1)
    if len(listOfFirmware) > 1:
        print("Too many matches: ")
        for l in listOfFirmware:
            print("\t" + l)
        sys.exit(1)
    trytoflash = listOfFirmware[0]

flashFile(trytoflash)
