import os
import random
import socket
import subprocess
import telnetlib
import time
import threading

__all__ = ['OpenOCD']
openocd_path = '/usr/local/bin/openocd'
openocd_libpath = '/usr/local/share/openocd/scripts'

possible_interfaces = os.listdir(os.path.join(openocd_libpath, 'interface'))
possible_interfaces = [(s.endswith('.cfg') and s[:-4] or s) for s in possible_interfaces]

class InterfaceNotFound(Exception):
    pass

class OpenOCD(object):
    config = {}
    def __init__(self, config=None):
        if config:
            self.config = config

    def scan_interface(self):
        ints = set()
        for interface in possible_interfaces:
            if self.try_interface(interface):
                self.config['interface'] = interface
                ints.add(interface)
        return ints

    def try_interface(self, interface):
        p = subprocess.Popen([openocd_path, '-c', 'source [find interface/%s.cfg]' % interface, 'init', 'exit'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        def WaitToKill():
            for t in range(25):
                time.sleep(.2)
                if p.poll():
                    return
            try:
                p.terminate()
            except AttributeError: # Python is too old - emulate it
                try:
                    os.kill(p.pid, 15)
                    for t in range(25):
                        time.sleep(.2)
                        if p.poll():
                            return
                    os.kill(p.pid, 9)
                except OSError:
                    pass
        t = threading.Thread(target=WaitToKill)
        t.start()
        data = str(p.communicate()[0]) # Output
        t.join()
        if ('Info' in data):
            return True
        return False

    def scan_chain(self, khz=1000):
        random.seed()
        if not 'interface' in self.config:
            raise InterfaceNotFound('No interface specified')
        while True:
            port = random.randint(1024, 65535)
            try:
                s = socket.socket()
                s.bind(('0.0.0.0', port))
                s.close()
                del s
                break
            except socket.error:
                continue
        p = subprocess.Popen([openocd_path, '-c', 'telnet_port %d' % port, '-c', 'noinit'], stderr=subprocess.STDOUT, stdout=file('/dev/null','w'))
        try:
            time.sleep(1)
            t = telnetlib.Telnet('127.0.0.1', port)
            startup = t.read_until('>')
            t.write('source [find interface/%s.cfg]\r\n' % self.config['interface'])
            interface_stat = t.read_until('>').split('\n')
            for l in interface_stat:
                if "Can't find" in l:
                    raise InterfaceNotFound(l.trim())
            t.write('jtag_rclk %d\r\n' % khz)
            t.read_until('>')
            t.write('init\r\n')
            chain = t.read_until('>').split('\n')
            chain = [s for s in chain if s.startswith('AUTO')]
            return chain
        finally:
            os.kill(p.pid, 9)

if __name__ == '__main__':
    test = OpenOCD()
    if False: # Disable slow scanning test
        test.config['interface'] = 'arm-usb-ocd'
    else:
        print "Finding interfaces..."
        ints = test.scan_interface()
        for i in ints:
            print "Found " + i
        if not ints:
            print "No interfaces found"
            raise SystemExit()

    chain = test.scan_chain()
    for s in chain:
        print s