#include <stdio.h>
#include <stdlib.h>
#include "xil_printf.h"
#include "xil_types.h"
#include "xgpio.h"
#include "xparameters.h"
#include "xstatus.h"
#include "xspi.h"

#define TPM_ACCESS_REG       0
#define TPM_STS_REG       0x18
#define TPM_DATA_FIFO_REG 0x24
#define TPM_DID_VID_REG  0xf00
#define TPM_RID_REG      0xf04

/*
 * Code adapted from: Chromium/Android TPM FTDI MPSSE library
 *   Especially:
 *   1. https://chromium.googlesource.com/aosp/platform/system/trunks/+/master/trunks_ftdi_spi.h
 *   2. https://chromium.googlesource.com/aosp/platform/system/trunks/+/master/trunks_ftdi_spi.cc
 *
 *   NOTE: Maybe due to incorrect initialization of high-level SPI driver,
 *         the XSpi_Transfer function doesn't actually work. So, wrote a
 *         simple another function `tpm_read_did_vid` to read TPM_DID_VID_0 (locality 0) register
 */
XGpio tpm_reset;
#define TPM_RESET_BIT 0x01

XSpi tpm;

u8 reset_pin(u8);
u8 delay(u32);

u8 spi_init();
u8 tpm_init();


int main()
{
	print("\r\n\r\n\e[1;30mAller Artix Ultrascale+ TPM Test\e[0m\r\n");
	print("\e[1;30m______________\e[0m\r\n\r\n");
    XGpio_Initialize(&tpm_reset, XPAR_GPIO_0_DEVICE_ID);
    XGpio_SetDataDirection(&tpm_reset, 1, 0x00);
    XGpio_DiscreteWrite(&tpm_reset, 1, 0);

    spi_init();
	tpm_init();

    return 0;
}

u8 spi_init() {
    XSpi_Config *cfgPtr;
    //Lookup SPI peripheral configuration details
	cfgPtr = XSpi_LookupConfig(XPAR_SPI_0_DEVICE_ID);
	if (cfgPtr == NULL)
	{
		return 2;
	}

	 if(XSpi_CfgInitialize(&tpm, cfgPtr, cfgPtr->BaseAddress) != XST_SUCCESS)
	{
		return 1;
	}

	//Set up SPI controller. Master, manual slave select. The SPI peripheral
	//is configured with no FIFO
	XSpi_SetControlReg(&tpm, 0x86);

	//Disable interrupts
	XSpi_IntrGlobalDisable(&tpm);

	return 0;
}

u8 spiReadData()
{
    while(!(XSpi_ReadReg(tpm.BaseAddr, XSP_SR_OFFSET) & 0x02));
    return XSpi_ReadReg(tpm.BaseAddr, XSP_DRR_OFFSET);
}

//This function writes one byte to the SPI peripheral
void spiWriteData(u8 data)
{
    while(XSpi_GetStatusReg(&tpm) & 0x08);
    XSpi_WriteReg(tpm.BaseAddr, XSP_DTR_OFFSET, data);
}

u32 tpm_read_did_vid() {
	u8 sendbuf[4] = { 0, 0, 0, 0};

	// The first byte of the frame header encodes the transaction type (read or
	// write) and size (set to length - 1).
	sendbuf[0] = 0x80 | 0x40 | (4-1);
	u32 addr = TPM_DID_VID_REG;
	for (int i = 0; i < 3; i++) {
		sendbuf[i + 1] = (addr >> (8 * (2 - i))) & 0xff;
	}

	u8 rcvbuf[4] = {0, 0, 0, 0};

	for (int i=0; i<4; i++){
		spiWriteData(sendbuf[i]);
		rcvbuf[i] = spiReadData();
	}

	// The TCG TPM over SPI specification introduces the notion of SPI flow
	// control (Section "6.4.5 Flow Control" of the TCG issued "TPM Profile
	// (PTP) Specification Revision 00.43).

	// The slave (TPM device) expects each transaction to start with a 4 byte
	// header transmitted by master. If the slave needs to stall the transaction,
	// it sets the MOSI bit to 0 during the last clock of the 4 byte header. In
	// this case the master is supposed to start polling the line, byte at time,
	// until the last bit in the received byte (transferred during the last
	// clock of the byte) is set to 1.
	while (!(rcvbuf[3] & 0x01)) {
		spiWriteData(0x00); // dummy write
		rcvbuf[3] = spiReadData();
	}

	u8 did_vid[4] = {0, 0, 0, 0};
	for (int i=0; i<4; i++) {
		spiWriteData(0x00); // Write dummy data
		did_vid[i] = spiReadData(); // Read data shifted into the register
	}

	// TPM Vendor IDs are Big-endian. SPI-TPM has order of LSBytes first, MSbits first.
	u32 temp_didvid = 0;
	for (int i=0; i<4; i++) {
		temp_didvid = temp_didvid + (did_vid[i] << (8 * i));
	}

	return temp_didvid;
}

u8 tpm_init() {
	uint32_t did_vid;

	// Reset the TPM using TPM_RST# pin connected to Bit-0 of AXI GPIO
	xil_printf(">> Resetting TPM...");
	reset_pin(0);
	delay(1000);
	reset_pin(1);
	xil_printf("Done\r\n");

	//Cycle CS to reset the flash to known state
	xil_printf(">> Resetting TPM SPI state-machine...");
	XSpi_WriteReg(tpm.BaseAddr, XSP_SSR_OFFSET, 0x00);
	XSpi_WriteReg(tpm.BaseAddr, XSP_SSR_OFFSET, 0x01);
	XSpi_WriteReg(tpm.BaseAddr, XSP_SSR_OFFSET, 0x00);
	xil_printf("Done\r\n");

	xil_printf(">> Reading TPM DID_VID register in locality 0...");
	did_vid = tpm_read_did_vid();
	xil_printf("Successful!\r\n\n");
	xil_printf(">> Connected to device with VID:DID = \e[1;34m%04X:%04X\e[0m\r\n\n", did_vid & 0xffff, did_vid >> 16);

	u16 vid = did_vid & 0xffff;
	u16 did = did_vid >> 16;

	xil_printf("\e[1;30mTPM Information\e[0m\r\n");
	xil_printf("\e[1;30m***************\e[0m\r\n\n");
	xil_printf("Vendor ID: \e[1;31m0x%04X ", vid); xil_printf(vid==0x1114 ? "\e[1;34m(Atmel)\e[0m\r\n" : "\e[1;31m(UNKNOWN)\e[0m\r\n");
	xil_printf("Device ID: \e[1;31m0x%04X\e[0m\r\n\n", did);
	if (did == 0x3205) {
		xil_printf("Chip Found: \e[1;30mAtmel AT97SC\e[1;35m3205\e[0m\r\n", did);
	}

	return 0;
}

u8 reset_pin(u8 state) {
	XGpio_DiscreteWrite(&tpm_reset, 1, state);
	return 0;
}

u8 delay(u32 d) {
	u32 i = 0, j= 0, a = 0;
	for (i=0; i<d; i++) {
		for (j=0; j<d; j++) {
			a++;
		}
	}

	return 0;
}