import tms_ptns
import cmd_proto

class jtagdev(object):
    ST_REST = 0
    ST_IDLE = 1
    ST_SHIFT_DR = 2
    ST_SHIFT_IR = 3
    
    def __init__(self, fo):
        self.fo = fo
        self.state = self.IDLE
        self.seq = 0
        pass

    def _send_cmd(self, code, data):
        cmd = cmd_proto.cmd(self.seq, code, data)
        frame = cmd.to_frame()

        self.fo.write(frame)
        self.fo.flush()
        pass

    def _send_nbits_data_cmd(self, code, nbits, data):
        nbits_data = cmd_proto.prepend_nbits(nbits, data)
        self._send_cmd(code, nbits_data)
        pass

    def _send_tms(self, ptn, nbits):
        self._send_nbits_data_cmd(cmd_proto.CPCMD_SHIFT_TMS,
                                  nbits, ptn)
        pass

    def wait_reply(self, tmo=None):
        import select
        
        if tmo is None:
            select.select((self.fo,), (), ())
        else:
            select.select((self.fo,), (), (), tmo)
            pass
        frame = self.fo.read()
        cmd = cmd_proto.cmd()
        cmd.from_frame(frame)
        return cmd

    def idle(self):
        self.seq = (self.seq + 1) % 256
        
        ptn, nbits = tms_ptns.TMS_IDLE_SEQ
        self.send_tms(ptn, nbits)
        
        self.state = self.ST_IDLE
        pass

    def go_shift_IR(self):
        if self.state == self.ST_SHIFT_IR:
            return
        
        if self.state == self.ST_IDLE:
            ptn, ptn_nbits = tms_ptns.TMS_CAP_IR
            self.send_tms(ptn, nbits)
        elif self.state == self.ST_SHIFT_DR:
            ptn, ptn_nbits = tms_ptns.TMS_SHIFT_DR_2_CAP_IR
            self.send_tms(ptn, nbits)
        else:
            raise RuntimeError, 'Transite to shift IR state from invalid state'
        self.state = self.ST_SHIFT_IR
        pass

    def go_shift_DR(self):
        if self.state == self.ST_SHIFT_DR:
            return
        
        if self.state == self.ST_IDLE:
            ptn, ptn_nbits = tms_ptns.TMS_CAP_DR
            self.send_tms(ptn, nbits)
        elif self.state == self.ST_SHIFT_IR:
            ptn, ptn_nbits = tms_ptns.TMS_SHIFT_IR_2_CAP_DR
            self.send_tms(ptn, nbits)
        else:
            raise RuntimeError, 'Transite to shift DR state from invalid state'
            pass
        self.state = self.ST_SHIFT_DR
        pass

    def shift_IR(self, data, nbits):
        self.seq = (self.seq + 1) % 256

        if self.state != self.ST_SHIFT_IR:
            raise RuntimeError, 'Invalid state'

        self._send_nbits_data_cmd(cmd_proto.CPCMD_SHIFT_TDI,
                                  nbits, data)
        pass

    def shift_DR(self, data, nbits):
        self.seq = (self.seq + 1) % 256

        if self.state != self.ST_SHIFT_DR:
            raise RuntimeError, 'Invalid state'

        self._send_nbits_data_cmd(cmd_proto.CPCMD_SHIFT_TDI,
                                  nbits, data)
        pass

    def shift_DR_n_out(self, data, nbits):
        self.seq = (self.seq + 1) % 256

        if self.state != self.ST_SHIFT_DR:
            raise RuntimeError, 'Invalid state'

        self._send_nbits_data_cmd(cmd_proto.CPCMD_SHIFT_TDI_TDO,
                                  nbits, data)
        pass

    def reset(self):
        self._send_cmd(cmd_proto.CPCMD_TRST, '')
        self.state = self.ST_RESET
        pass
    pass

def _extract_nbits_data(nbits_data):
    nbits = ord(nbits_data[0]) | (ord(nbits_data[1] << 8))
    return nbits, nbits_data[2:]

def _get_bit(data, bit_idx):
    byte_off = bit_idx / 8
    bit_off = bit_idx % 8
    return (ord(data[byte_off]) >> bit_off) & 0x1

def identify_components(dev):
    dev.go_shift_IR()
    reply = dev.wait_reply()
    if reply.code != cmd_proto.CPCMD_ACK:
        raise RutimeError, 'invliad reply code 0x%02x' % (reply.code)

    dev.go_shift_DR()
    dev.shift_DR('\xff' * 253, 253 * 8)
    reply = dev.wait_reply()
    if reply.code != cmd_proto.CPCMD_DATA:
        raise RutimeError, 'invliad reply code 0x%02x' % (reply.code)

    nbits, data = _extract_nbits_data(reply.data)
    
    components = []
    i = 0
    while i < nbits:
        bit = _get_bit(data, i)
        if bit == 0:
            components.append(0)
            i = i + 1
        else:
            comp_id = 0
            for j in range(32):
                bit = _get_bit(data, i + j)
                comp_id = comp_id | (bit << j)
                pass
            if comp_id == 0xffffffff:
                break
            components.append(comp_id)
            i = i + 32
            pass
        pass
    return components