#include <stdio.h>
#include <util/delay.h>
#include "jtag.h"
#include "avriotools.h"

/* It is supposed to work at 1Mbps */
#define CLK_DELAY() _delay_us(1)

void jtag_init(void) {
    pin_mode(&JTAG_PORT, JTAG_TCK, PM_OUTPUT);
    pin_mode(&JTAG_PORT, JTAG_TMS, PM_OUTPUT);
    pin_mode(&JTAG_PORT, JTAG_TDI, PM_OUTPUT);
    pin_mode(&JTAG_PORT, JTAG_TDO, PM_INPUT);
    pin_mode(&JTAG_PORT, JTAG_TRST, PM_OUTPUT);
    pin_lo(JTAG_PORT, JTAG_TCK);
    pin_lo(JTAG_PORT, JTAG_TMS);
    pin_lo(JTAG_PORT, JTAG_TDI);
    pin_hi(JTAG_PORT, JTAG_TRST);
}

#define TCK_LO() pin_lo(JTAG_PORT, JTAG_TCK)
#define TCK_HI() pin_hi(JTAG_PORT, JTAG_TCK)
#define SEND_BIT(pin, bit)			\
    do {					\
	if(bit)					\
	    pin_hi(JTAG_PORT, pin);		\
	else					\
	    pin_lo(JTAG_PORT, pin);		\
	CLK_DELAY();				\
	TCK_HI();				\
	CLK_DELAY();				\
	TCK_LO();				\
    } while(0)
#define GET_TDO() (JTAG_PIN & _BV(JTAG_TDO))
#define SEND_GET_BIT(pin, bit, out)		\
    do {					\
	if(bit)					\
	    pin_hi(JTAG_PORT, pin);		\
	else					\
	    pin_lo(JTAG_PORT, pin);		\
	CLK_DELAY();				\
	TCK_HI();				\
	CLK_DELAY();				\
	out = GET_TDO();			\
	TCK_LO();				\
    } while(0)

void jtag_trst(void) {
    pin_lo(JTAG_PORT, JTAG_TRST);
    SEND_BIT(JTAG_TMS, 1);
    SEND_BIT(JTAG_TMS, 1);
    SEND_BIT(JTAG_TMS, 1);
    SEND_BIT(JTAG_TMS, 1);
    SEND_BIT(JTAG_TMS, 1);
    pin_hi(JTAG_PORT, JTAG_TRST);
}

/*!
 * Before shifting registers, TAP controller must move to last state before
 * shift state.  The state machine transite to shift state when shifting
 * starts.
 */
void jtag_tms(unsigned char *buf, int nbits) {
    int i;
    int nbytes, byte;
    int byteoff, bitoff;
    int bit;
    
    pin_hi(JTAG_PORT, JTAG_TDI);

    nbytes = nbits / 8;
    for(i = 0; i < nbytes; i++) {
	byte = buf[i];
	
	bit = byte & 0x01;
	SEND_BIT(JTAG_TMS, bit);
	
	bit = byte & 0x02;
	SEND_BIT(JTAG_TMS, bit);
	
	bit = byte & 0x04;
	SEND_BIT(JTAG_TMS, bit);
	
	bit = byte & 0x08;
	SEND_BIT(JTAG_TMS, bit);
	
	bit = byte & 0x10;
	SEND_BIT(JTAG_TMS, bit);
	
	bit = byte & 0x20;
	SEND_BIT(JTAG_TMS, bit);
	
	bit = byte & 0x40;
	SEND_BIT(JTAG_TMS, bit);
	
	bit = byte & 0x80;
	SEND_BIT(JTAG_TMS, bit);
    }
    
    byte = buf[i];
    nbits %= 8;
    for(i = 0; i < nbits; i++) {
	bit = byte & (1 << i);
	SEND_BIT(JTAG_TMS, bit);
    }
}

void jtag_shift(unsigned char *buf, int nbits) {
    int i;
    int nbits_1;
    int nbytes, byte;
    int byteoff, bitoff;
    int bit, remain;

    if(nbits == 0)
	return;
    
    /* Transite to shift state.
     * \sa jtag_tms()
     */
    pin_lo(JTAG_PORT, JTAG_TMS);

    nbits_1 = nbits - 1;
    nbytes = nbits_1 / 8;
    for(i = 0; i < nbytes; i++) {
	byte = buf[i];
	
	bit = byte & 0x01;
	SEND_BIT(JTAG_TDI, bit);
	
	bit = byte & 0x02;
	SEND_BIT(JTAG_TDI, bit);
	
	bit = byte & 0x04;
	SEND_BIT(JTAG_TDI, bit);
	
	bit = byte & 0x08;
	SEND_BIT(JTAG_TDI, bit);
	
	bit = byte & 0x10;
	SEND_BIT(JTAG_TDI, bit);
	
	bit = byte & 0x20;
	SEND_BIT(JTAG_TDI, bit);
	
	bit = byte & 0x40;
	SEND_BIT(JTAG_TDI, bit);
	
	bit = byte & 0x80;
	SEND_BIT(JTAG_TDI, bit);
    }
    
    remain = nbits_1 % 8;
    byte = buf[i];
    for(i = 0; i < remain; i++) {
	bit = byte & (1 << i);
	SEND_BIT(JTAG_TDI, bit);
    }

    byte = buf[nbits / 8];
    bit = byte & (1 << (nbits_1 % 8));
    
    pin_hi(JTAG_PORT, JTAG_TMS);
    SEND_BIT(JTAG_TDI, bit);
}

void jtag_shift_inout(unsigned char *ibuf, unsigned char *obuf, int nbits) {
    int i, j;
    int nbits_1;
    int nbytes, byte, obyte;
    int tdo;
    int byteoff, bitoff;
    int bit, remain;
    
    if(nbits == 0)
	return;
    
    /* Transite to shift state.
     * \sa jtag_tms()
     */
    pin_lo(JTAG_PORT, JTAG_TMS);

    nbits_1 = nbits - 1;
    nbytes = nbits_1 / 8;
    for(i = 0; i < nbytes; i++) {
	byte = ibuf[i];
	obyte = 0;
	
	bit = byte & 0x01;
	SEND_GET_BIT(JTAG_TDI, bit, tdo);
	if(tdo)
	    obyte |= 0x01;
	else
	    obyte &= ~0x01;
	
	bit = byte & 0x02;
	SEND_GET_BIT(JTAG_TDI, bit, tdo);
	tdo = GET_TDO();
	if(tdo)
	    obyte |= 0x02;
	else
	    obyte &= ~0x02;
	
	bit = byte & 0x04;
	SEND_GET_BIT(JTAG_TDI, bit, tdo);
	if(tdo)
	    obyte |= 0x04;
	else
	    obyte &= ~0x04;
	
	bit = byte & 0x08;
	SEND_GET_BIT(JTAG_TDI, bit, tdo);
	if(tdo)
	    obyte |= 0x08;
	else
	    obyte &= ~0x08;
	
	bit = byte & 0x10;
	SEND_GET_BIT(JTAG_TDI, bit, tdo);
	if(tdo)
	    obyte |= 0x10;
	else
	    obyte &= ~0x10;
	
	bit = byte & 0x20;
	SEND_GET_BIT(JTAG_TDI, bit, tdo);
	if(tdo)
	    obyte |= 0x20;
	else
	    obyte &= ~0x20;
	
	bit = byte & 0x40;
	SEND_GET_BIT(JTAG_TDI, bit, tdo);
	if(tdo)
	    obyte |= 0x40;
	else
	    obyte &= ~0x40;
	
	bit = byte & 0x80;
	SEND_GET_BIT(JTAG_TDI, bit, tdo);
	if(tdo)
	    obyte |= 0x80;
	else
	    obyte &= ~0x80;

	obuf[i] = obyte;
    }
    
    remain = nbits_1 % 8;
    byte = ibuf[i];
    obyte = 0;
    for(j = 0; j < remain; j++) {
	bit = byte & (1 << j);
	SEND_GET_BIT(JTAG_TDI, bit, tdo);
	if(tdo)
	    obyte |= 1 << j;
	else
	    obyte &= ~(1 << j);
    }
    obuf[i] = obyte;

    byte = ibuf[nbits / 8];
    bit = byte & (1 << (nbits_1 % 8));

    pin_hi(JTAG_PORT, JTAG_TMS);
    SEND_GET_BIT(JTAG_TDI, bit, tdo);
    if(tdo)
	obuf[nbits / 8] |= 1 << j;
    else
	obuf[nbits / 8] &= ~(1 << j);
}