#include <tigcclib.h>
#include "equates.h"
#include "usb.h"

unsigned int state;															//Jailbreak state
int port_cur;																		//Current port/device we're emulating
int port_addr[7];																//Mapping between hub port and USB function address
unsigned int port_status[6];										//Current status for each hub port
unsigned int port_change[6];										//Current unacknowledged changes for each hub port
unsigned int sent_interrupt_data;								//Flag to indicate whether we've sent hub port status change once before (for troublesome ports)
unsigned int hub_int_response;									//Value sent to host that indicates which hub ports have status changes
unsigned int hub_int_ready;											//Flag to indicate whether we're done sending jailbreak-state-changing control request response
unsigned int device_1_ready;										//Hack to get around that stupid device 1 descriptor crapping out early...
unsigned int last_port_conn_clear;							//Last hub port whose connection was cleared by host
unsigned int last_port_reset_clear;							//Last hub port reset by host
unsigned int last_port_status;									//Last hub port queried by host about status
unsigned int jigBytesReceived;									//Number of jig challenge bytes received so far from PS3
short lastDescriptorType;												//Last descriptor type (01 = device, 02 = config) sent
short useExternalPayload;												//Flag to indicate whether to use built-in payload or external one
unsigned char device4ConfigDescriptor3[48];		//Buffer used for device 4 config descriptor 3
unsigned char jigResponseData[64];							//Buffer used for jig response data

char stage1FileName[41];
char stage2FileName[41];
char compatFileName[41];
short useDevice6;
FILES fhandle;

void USB_SetFunctionAddress(int address);
void LoadPS3JBOptions();
void LoadPayloadData(int wLength);
void LoadCompatibilityData();

void switch_port(int port)
{
	if (port_cur == port) return;

	port_cur = port;
	if (port_addr[port] < 0)
		port_addr[port] = 0;

	USB_SetFunctionAddress(port_addr[port] & 0x7F);
}

void connect_port(int port)
{
	last_port_reset_clear = 0;
	hub_int_response = (1 << port);
	port_status[port - 1] = PORT_FULL;
	port_change[port - 1] = C_PORT_CONN;
}

void disconnect_port(int port)
{
	last_port_conn_clear = 0;
	hub_int_response = (1 << port);
	port_status[port - 1] = PORT_EMPTY;
	port_change[port - 1] = C_PORT_CONN;
}

void HandleConnected()
{
	//Reset all our jailbreak-related variables since the PS3 just initialized us
  state = init;
	sent_interrupt_data = 0;
	hub_int_response = 0;
	hub_int_ready = 0;
	jigBytesReceived = 0;
	device_1_ready = 0;

	//Set all the hub's ports to their default states
  port_cur = 0;
	last_port_conn_clear = 0;
	last_port_reset_clear = 0;
	last_port_status = 0;
  unsigned int i;
  for( i = 0; i < 7; i++ )
  	port_addr[i] = -1;
  for( i = 0; i < 6; i++ )
  {
  	port_status[i] = PORT_EMPTY;
  	port_change[i] = C_PORT_NONE;
  }
  
  //Load options
  LoadPS3JBOptions();
}

void HandleSetAddress(int address)
{
	//Save the address that was set for the current port
	port_addr[port_cur] = address & 0xFF;
}

void HandleHubSetFeature(int feature, int index)
{
	switch(feature)
	{
		case 0x0008: //PORT_POWER
		{
			if( index == 6 && state == init)
				state = hub_ready;
			break;
		}
		case 0x0004: //PORT_RESET
		{
			hub_int_response = (1 << index);
			port_change[index-1] |= C_PORT_RESET;
			break;
		}
	}
}

void HandleHubClearFeature(int feature, int index)
{
	switch(feature)
	{
		case 0x0010: //C_PORT_CONNECTION
		{
			port_change[index-1] &= ~C_PORT_CONN;
			last_port_conn_clear = index;
			break;
		}
		case 0x0014: //C_PORT_RESET
		{
			port_change[index-1] &= ~C_PORT_RESET;
			last_port_reset_clear = index;
			break;
		}
	}
}

void HandleHubGetStatus(unsigned char* data, int index)
{
	hub_int_ready = 0;
	last_port_status = index;
	
	if (index == 0)
	{
		data[0] = 0x00;
		data[1] = 0x00;
		data[2] = 0x00;
		data[3] = 0x00;
	}
	else
	{
		data[0] = (port_status[index-1] & 0xFF);
		data[1] = (port_status[index-1] >> 8) & 0xFF;
		data[2] = (port_change[index-1] & 0xFF);
		data[3] = (port_change[index-1] >> 8) & 0xFF;
	}
}

int HandleDescriptorRequest(int wValue, int wLength, void** descAddress)
{
	int descType = (wValue >> 8);
	int descLength = 0;
	void* address = NULL;

	hub_int_ready = 0;
	lastDescriptorType = descType;
	
	LoadPayloadData(wLength);
	
	if (descType == 0x01)
	{
		switch(port_cur)
		{
			case 0:
			{
				address = deviceDescriptor;
				descLength = sizeof(deviceDescriptor);
				break;
			}
			case 1:
			{
				address = port1DeviceDescriptor;
				descLength = sizeof(port1DeviceDescriptor);
				break;
			}
			case 2:
			{
				address = port2_device_descriptor;
				descLength = sizeof(port2_device_descriptor);
				break;
			}
			case 3:
			{
				address = port3_device_descriptor;
				descLength = sizeof(port3_device_descriptor);
				break;
			}
			case 4:
			{
				address = port4_device_descriptor;
				descLength = sizeof(port4_device_descriptor);
				break;
			}
			case 5:
			{
				address = port5_device_descriptor;
				descLength = sizeof(port5_device_descriptor);
				break;
			}
			case 6:
			{
				address = port6_device_descriptor;
				descLength = sizeof(port6_device_descriptor);
				break;
			}
		}
	}
	else if(descType == 0x02)
	{
		int num = (wValue & 0xFF);
	
		switch(port_cur)
		{
			case 0:
			{
				address = configDescriptor;
				descLength = sizeof(configDescriptor);
				break;
			}
			case 1:
			{
				if (num < 4)
				{
					if (wLength == 8)
					{
						address = port1_short_config_descriptor;
						descLength = sizeof(port1_short_config_descriptor);
					}
					else
					{
						//HACK: We're going to want to make this configurable...
						address = port1_config_descriptor;
						descLength = sizeof(port1_config_descriptor);
					}
	
					if (num == 3 && wLength > 8)
					{
						device_1_ready = 1;
					}
				}
				break;
			}
			case 2:
			{
				address = port2_config_descriptor;
				descLength = sizeof(port2_config_descriptor);
				
				if (wLength > 8)
				{
					state = p2_ready;
				}
				break;
			}
			case 3:
			{
				address = port3_config_descriptor;
				descLength = sizeof(port3_config_descriptor);
				if (num == 1 && wLength > 8)
				{
					state = p3_ready;
				}
				break;
			}
			case 4:
			{
				if (num == 0)
				{
					address = port4_config_descriptor_1;
					descLength = sizeof(port4_config_descriptor_1);
				}
				else if (num == 1)
				{
					if (wLength == 8)
					{
						address = port4_short_config_descriptor_2;
						descLength = sizeof(port4_short_config_descriptor_2);
					}
					else
					{
						address = port4_config_descriptor_2;
						descLength = sizeof(port4_config_descriptor_2);
					}
				}
				else if (num == 2)
				{
					LoadCompatibilityData();
					address = device4ConfigDescriptor3;
					descLength = sizeof(device4ConfigDescriptor3);
					if (wLength > 8)
					{
						state = p4_ready;
					}
				}
				break;
			}
			case 5:
			{
				address = port5_config_descriptor;
				descLength = sizeof(port5_config_descriptor);
				break;
			}
			case 6:
			{
				address = port6_config_descriptor;
				descLength = sizeof(port6_config_descriptor);
				break;
			}
		}
	}

	*descAddress = address;
	return (descLength > wLength)? wLength : descLength;
}

int HandleHubDescriptorRequest(int wLength, void** descAddress)
{
	void* address = NULL;
	int descLength = 0;
	
	hub_int_ready = 0;

	address = hubDescriptor;
	descLength = sizeof(hubDescriptor);

	*descAddress = address;
	return (descLength > wLength)? wLength : descLength;
}

void HandleControlOutputDone()
{
	//Only indicate device 1 (the really big one) is enumerated once we're done sending the last descriptor
	if (port_cur == 1 && device_1_ready > 0)
	{
		state = p1_ready;
	}

	hub_int_ready = 1;
}

void HandleIncomingData()
{
	//Receive 8 bytes of interrupt data and discard it
	unsigned char data[8];
	USB_ReceiveInterruptData(0x02, data, 8);
	jigBytesReceived += 8;
}

void HUB_Task()
{
	//unsigned int i;
	
	if (hub_int_response > 0)
	{
		//Send a port status change
		char data[1] = {hub_int_response};
		USB_SendInterruptData(0x01, data, 1);
		hub_int_response = 0;
	}

	if (state == hub_ready)
	{
		connect_port(1);
		state = p1_wait_reset;
	}

	if (state == p1_wait_reset && last_port_reset_clear == 1)
	{
		switch_port(1);
		state = p1_wait_enumerate;
	}

	if (state == p1_ready)
	{
		switch_port(0);
		connect_port(2);
		state = p2_wait_reset;
	}

	if (state == p2_wait_reset && last_port_reset_clear == 2)
	{
		switch_port(2);
		state = p2_wait_enumerate;
	}

	if (state == p2_ready)
	{
		switch_port(0);
		connect_port(3);
		state = p3_wait_reset;
	}

	if (state == p3_wait_reset && last_port_reset_clear == 3)
	{
		switch_port(3);
		state = p3_wait_enumerate;
	}

	if (state == p3_ready && hub_int_ready > 0)
	{
		switch_port(0);
		disconnect_port(2);
		state = p2_wait_disconnect;
	}

	if (state == p2_wait_disconnect && last_port_conn_clear == 2)
	{
		state = p4_wait_connect;
	}

	if (state == p4_wait_connect)
	{
		sent_interrupt_data = 0;
		last_port_status = 0;
		connect_port(4);
		state = p4_wait_reset;
	}
	
	if (state == p4_wait_reset)
	{
		if (sent_interrupt_data > 0)
		{
			if (last_port_reset_clear == 4)
			{
				switch_port(4);
				state = p4_wait_enumerate;
			}
		}
		else
		{
			if (last_port_status == 0x02)
			{
				sent_interrupt_data = 0x01;
				hub_int_response = 0x10;
			}
		}
	}

	if (state == p4_ready)
	{
		switch_port(0);
		connect_port(5);
		state = p5_wait_reset;
	}
	
	if (state == p5_wait_reset && last_port_reset_clear == 5)
	{
		switch_port(5);
		state = p5_wait_enumerate;
	}

	if (state == p5_wait_enumerate && jigBytesReceived >= 64)
	{
		//Send jig response
		LoadCompatibilityData();
		USB_SendInterruptData(0x01, jigResponseData, 8);
		USB_SendInterruptData(0x01, jigResponseData+8, 8);
		USB_SendInterruptData(0x01, jigResponseData+16, 8);
		USB_SendInterruptData(0x01, jigResponseData+24, 8);
		USB_SendInterruptData(0x01, jigResponseData+32, 8);
		USB_SendInterruptData(0x01, jigResponseData+40, 8);
		USB_SendInterruptData(0x01, jigResponseData+48, 8);
		USB_SendInterruptData(0x01, jigResponseData+56, 8);

		switch_port(0);
		disconnect_port(3);
		state = p3_wait_disconnect;
	}

	if (state == p3_wait_disconnect && last_port_conn_clear == 3)
	{
		state = p3_disconnected;
	}

	if (state == p3_disconnected)
	{
		sent_interrupt_data = 0;
		last_port_status = 0;
		switch_port(0);
		disconnect_port(5);
		state = p5_wait_disconnect;
	}

	if (state == p5_wait_disconnect)
	{
		if (last_port_conn_clear == 5 || sent_interrupt_data > 50)
		{
			if (last_port_conn_clear == 5)
				state = p5_disconnected;
		}
		else
		{
			if (last_port_status == 3 || last_port_status == 5)
			{
				//Wait for 100ms
				OSFreeTimer(0x09);
				OSRegisterTimer(0x09, 2);
				while (!OSTimerExpired(0x09));

				sent_interrupt_data += 1;
				hub_int_response = 0x20;
				*(unsigned char*)0x71008E = 0x01;
				unsigned char val = *(unsigned char*)0x710091;
				val &= 0xEF;
				val |= 0x40;
				*(unsigned char*)0x710091 = val;
			}
		}
	}

	if (state == p5_disconnected)
	{
		*(unsigned char*)0x71008E = 0x01;
		unsigned char val = *(unsigned char*)0x710091;
		val &= 0xEF;
		val |= 0x40;
		*(unsigned char*)0x710091 = val;
		sent_interrupt_data = 0;
		last_port_status = 0;
		switch_port(0);
		disconnect_port(4);
		state = p4_wait_disconnect;
	}

	if (state == p4_wait_disconnect)
	{
		if (last_port_conn_clear == 4 || sent_interrupt_data > 50)
		{
			if (last_port_conn_clear == 4)
				state = p4_disconnected;
		}
		else
		{
			if (last_port_status == 5 || last_port_status == 4)
			{
				//Wait for 100ms
				OSFreeTimer(0x09);
				OSRegisterTimer(0x09, 2);
				while (!OSTimerExpired(0x09));

				sent_interrupt_data += 1;
				hub_int_response = 0x10;
				*(unsigned char*)0x71008E = 0x01;
				unsigned char val = *(unsigned char*)0x710091;
				val &= 0xEF;
				val |= 0x40;
				*(unsigned char*)0x710091 = val;
			}
		}
	}

	if (state == p4_disconnected)
	{
		*(unsigned char*)0x71008E = 0x01;
		unsigned char val = *(unsigned char*)0x710091;
		val &= 0xEF;
		val |= 0x40;
		*(unsigned char*)0x710091 = val;
		sent_interrupt_data = 0;
		last_port_status = 0;
		switch_port(0);
		disconnect_port(1);
		state = p1_wait_disconnect;
	}
	
	if (state == p1_wait_disconnect)
	{
		if (last_port_conn_clear == 1 || sent_interrupt_data > 50)
		{
			if (last_port_conn_clear == 1)
			{
				if (useDevice6 == 2)
				{
					sent_interrupt_data = 0;
					last_port_status = 0;
					switch_port(0);
					connect_port(6);
					state = p6_wait_reset;
				}
				else
					state = done;
			}
		}
		else
		{
			if (last_port_status == 4 || last_port_status == 1)
			{
				//Wait for 100ms
				OSFreeTimer(0x09);
				OSRegisterTimer(0x09, 2);
				while (!OSTimerExpired(0x09));

				sent_interrupt_data += 1;
				hub_int_response = 0x02;
				*(unsigned char*)0x71008E = 0x01;
				unsigned char val = *(unsigned char*)0x710091;
				val &= 0xEF;
				val |= 0x40;
				*(unsigned char*)0x710091 = val;
			}
		}
	}
	
	if (state == p6_wait_reset)
	{
		if (last_port_reset_clear == 6 || sent_interrupt_data > 50)
		{
			if (last_port_reset_clear == 6)
			{
				switch_port(6);
				state = p6_wait_enumerate;
			}
		}
		else
		{
			if (last_port_status == 1 || last_port_status == 6)
			{
				//Wait for 100ms
				OSFreeTimer(0x09);
				OSRegisterTimer(0x09, 2);
				while (!OSTimerExpired(0x09));
	
				sent_interrupt_data += 1;
				hub_int_response = 0x40;
			}
		}
	}
}

void HandleUnknownControlRequest(int bmRequestType, int bRequest)
{
	if (bmRequestType == 0x21 && bRequest == 0xAA)
	{
		USB_FinishControlRequest();
		state = done;
	}
}

void FetchOptions(char* requestBuffer, short* pulldownBuffer)
{
	int i;
	for (i = 0; i < 41; i++)
		requestBuffer[41*0+i] = stage1FileName[i];
	for (i = 0; i < 41; i++)
		requestBuffer[41*1+i] = stage2FileName[i];
	for (i = 0; i < 41; i++)
		requestBuffer[41*2+i] = compatFileName[i];
	pulldownBuffer[0] = useDevice6;
}

void LoadPS3JBOptions()
{
	//Set the default options first
	int i;
	for (i = 0; i < 41; i++)
	{
		stage1FileName[i] = 0;
		stage2FileName[i] = 0;
		compatFileName[i] = 0;
	}
	useDevice6 = 1;

	//Now read the data in
	FILES f;
	if (FOpen(configFileName, &f, FM_READ, "CONF") == FS_OK)
	{
		short id;
		FRead(&id, 2, &f);
		FRead(stage1FileName, 41, &f);
		FRead(stage2FileName, 41, &f);
		FRead(compatFileName, 41, &f);
		FRead(&useDevice6, 2, &f);
		
		FClose(&f);
	}
}

void LoadCompatibilityData()
{
	//Load default device 4 config descriptor 3 data
	int count = sizeof(port4_config_descriptor_3);
	int i;
	
	for (i = 0; i < count; i++)
		device4ConfigDescriptor3[i] = port4_config_descriptor_3[i];

	//Load default jig response data
	count = sizeof(jigResponse);
	for (i = 0; i < count; i++)
		jigResponseData[i] = jigResponse[i];
		
	//Now try to read the data from the compatibility AppVar
	if (FOpen(compatFileName, &fhandle, FM_READ, "FWCB") == FS_OK)
	{
		//Read the new jig response data
		unsigned char size[2];
		FRead(size, 2, &fhandle);
		count = ((size[1] << 8) & 0xFF) | size[0];
		FRead(jigResponseData, count, &fhandle);
		
		//Read the new device 4 config descriptor 3 data
		FRead(size, 2, &fhandle);
		count = ((size[1] << 8) & 0xFF) | size[0];
		FRead(device4ConfigDescriptor3, count, &fhandle);
		
		FClose(&fhandle);
	}
}

void LoadPayloadData(int wLength)
{
	//Assume using built-in payload right now
	useExternalPayload = 0;

	//Do we even want to try loading an external file right now?
	if (port_cur == 1 && lastDescriptorType == 0x02 && wLength > 8)
	{
		//Yes, now read the data in
		if (FOpen(stage1FileName, &fhandle, FM_READ, "PAYL") == FS_OK)
			useExternalPayload = 1;
	}
}

void ReadPayloadData(unsigned char* buffer, unsigned int count)
{
	if (useExternalPayload)
	{
		FRead(buffer, count, &fhandle);
	}
	else
	{
		//Just go with the default...
	}
}

void ClosePayloadData()
{
	if (useExternalPayload)
		FClose(&fhandle);
}