#include "mouse.h"

#include "stdio.h"

// status register on command port (read)
#define STATUS_REG_OUTPUT_BUF_FULL 0x01
#define STATUS_REG_INPUT_BUF_FULL 0x02
#define STATUS_REG_OUTPUT_BUF2_FULL 0x20

// status control register flags
#define STATUS_CONTROL_CONFIG_ENABLE_IRQ_PORT2 0x02
#define STATUS_CONTROL_CONFIG_CLOCK_PORT2 0x20

// generic PS/2 commands
#define COMMAND_GET_STATE 0x20
#define COMMAND_SET_STATE 0x60
#define COMMAND_DISABLE_PORT2 0xA7
#define COMMAND_ENABLE_PORT2 0xA8
#define COMMAND_SEND_TO_PORT2 0xD4

// mouse commands
#define MOUSE_COMMAND_GET_DEVICE_ID 0xF2
#define MOUSE_COMMAND_SET_SAMPLE_RATE 0xF3
#define MOUSE_COMMAND_ENABLE_DATAREPORTING 0xF4
#define MOUSE_COMMAND_SET_DEFAULTS 0xF6
#define DEFAULT_NOF_PACKETS 3

// mouse data packet one
#define MOUSE_BYTE1_LEFT_BUTTON 0x01
#define MOUSE_BYTE1_RIGHT_BUTTON 0x02
#define MOUSE_BYTE1_MIDDLE_BUTTON 0x04
#define MOUSE_BYTE1_ID_BIT 0x08
#define MOUSE_BYTE1_X_SIGN 0x10
#define MOUSE_BYTE1_Y_SIGN 0x20
#define MOUSE_BYTE1_X_OVERFLOW 0x40
#define MOUSE_BYTE1_Y_OVERFLOW 0x80

#undef DEBUG

static uint8_t read_data( mouse_t *mouse )
{
	while( ( port8_read( &mouse->command_port ) & STATUS_REG_OUTPUT_BUF_FULL ) == 0 ) { }
	return port8_read( &mouse->data_port );
}

static void send_command( mouse_t *mouse, uint8_t command )
{
	while( port8_read( &mouse->command_port ) & STATUS_REG_INPUT_BUF_FULL ) { }
	port8_write( &mouse->command_port, command );
}

static void write_data( mouse_t *mouse, uint8_t data )
{
	while( port8_read( &mouse->command_port ) & STATUS_REG_INPUT_BUF_FULL ) { }
	port8_write( &mouse->data_port, data );
}

static void read_ack( mouse_t *mouse )
{
	uint8_t data = read_data( mouse );
	if( data == 0xFA ) {
		return;
	} else {
		printf( "ERR PS/2 mouse: acknoledgment failed 0x%X\n", data );
	}
}

static void set_sample_rate( mouse_t *mouse, uint8_t samples )
{
	send_command( mouse, COMMAND_SEND_TO_PORT2 );
	write_data( mouse, MOUSE_COMMAND_SET_SAMPLE_RATE );
	read_ack( mouse );

	send_command( mouse, COMMAND_SEND_TO_PORT2 );
	write_data( mouse, samples );
	read_ack( mouse );
}

static mouse_vtable_t const mouse_vtable = {
	{
	driver_register_interrupt_handler,
	mouse_activate,
	mouse_deactivate,
	mouse_deinit,
	mouse_print_info
	}
};

void mouse_init( mouse_t *mouse, mouse_event_handler_t handler, const uint32_t res_x, const uint32_t res_y, void *context )
{
	memset( mouse, 0, sizeof( mouse_t ) );

	// TODO: we should probe for wheel mouse and more than 3 packets
	mouse->nof_packets = DEFAULT_NOF_PACKETS;

	mouse_set_resolution( mouse, res_x, res_y );
	mouse->handler = handler;
	mouse->context = context;

	port8_init( &mouse->command_port, 0x64 );
	port8_init( &mouse->data_port, 0x60 );

	((driver_t *)mouse)->vtable = (driver_vtable_t *)&mouse_vtable;
}

void mouse_activate( void *obj )
{
	puts( "Activating driver for PS/2 mouse.." );

	mouse_t *mouse = obj;

	// enable port 2
	send_command( mouse, COMMAND_ENABLE_PORT2 );
	
	// enable interrupts for port 2
	send_command( mouse, COMMAND_GET_STATE );
	uint8_t status = read_data( mouse );
	status |= STATUS_CONTROL_CONFIG_ENABLE_IRQ_PORT2;
	status &= ~STATUS_CONTROL_CONFIG_CLOCK_PORT2;
	send_command( mouse, COMMAND_SET_STATE );
	write_data( mouse, status );

	// detecting mouse wheel (and 4th data packet)
	set_sample_rate( mouse, 200 );
	set_sample_rate( mouse, 100 );
	set_sample_rate( mouse, 80 );
	send_command( mouse, COMMAND_SEND_TO_PORT2 );
	write_data( mouse, MOUSE_COMMAND_GET_DEVICE_ID );
	read_ack( mouse );
	uint8_t id = read_data( mouse );
	printf( "PS/2 mouse model id is 0x%X\n", id );
	if( id == 0x03 ) {
		mouse->nof_packets = 4;
	}

	// set mouse defaults
//	send_command( mouse, COMMAND_SEND_TO_PORT2 );
//	write_data( mouse, MOUSE_COMMAND_SET_DEFAULTS );

	// enable mouse on second port
	send_command( mouse, COMMAND_SEND_TO_PORT2 );
	write_data( mouse, MOUSE_COMMAND_ENABLE_DATAREPORTING );
	read_ack( mouse );
}

void mouse_deactivate( void *obj )
{
	puts( "Dectivating driver for PS/2 mouse.." );

	mouse_t *mouse = obj;

	send_command( mouse, COMMAND_DISABLE_PORT2 );
}

void mouse_deinit( void *obj )
{
	// nothing to do
}

uint32_t mouse_handle_interrupt( interrupt_handler_t *handler, uint32_t esp )
{
	mouse_t *mouse = (mouse_t *)handler->driver;

	uint8_t status = port8_read( &mouse->command_port );
	if( !( status & STATUS_REG_OUTPUT_BUF2_FULL ) ) {
		return esp;
	}

	mouse->buf[mouse->offset] = port8_read( &mouse->data_port );

#ifdef DEBUG
	printf( "MOUSE: %d 0x%X\n", mouse->offset, mouse->buf[mouse->offset] );
#endif
	
	switch( mouse->offset ) {
		case 0:
			if( mouse->buf[0] & MOUSE_BYTE1_ID_BIT ) {
				for( int i = 0; i < 3; i++ ) {
					if( ( mouse->buf[0] & ( 0x01 << i ) ) != ( mouse->buttons & ( 0x01 << i ) ) ) {
						mouse_event_t event;
						event.button = i + 1;
						if( mouse->buttons & ( 0x01 << i ) ) {
							event.type = MOUSE_EVENT_TYPE_BUTTON_UP;
						} else {
							event.type = MOUSE_EVENT_TYPE_BUTTON_DOWN;
						}
						event.cursor_x = mouse->cursor_x;
						event.cursor_y = mouse->cursor_y;
						mouse->handler( &event, mouse->context );
					}
				}
				mouse->buttons = mouse->buf[0];
			}
			break;

		case 1:
			// delay evaluation of move events till second byte
			break;

		case 2:
			if( mouse->buf[1] !=0 || mouse->buf[2] != 0 ) {

				mouse_event_t event;
				event.type = MOUSE_EVENT_TYPE_MOVE;
				event.old_cursor_x = mouse->cursor_x;
				event.old_cursor_y = mouse->cursor_y;

				mouse->cursor_x += (int8_t)mouse->buf[1];
				mouse->cursor_y -= (int8_t)mouse->buf[2];
				if( mouse->cursor_x < 0 ) {
					mouse->cursor_x = 0;
				}
				if( mouse->cursor_x > mouse->res_x - 1 ) {
					mouse->cursor_x = mouse->res_x - 1;
				}
				if( mouse->cursor_y < 0 ) {
					mouse->cursor_y = 0;
				}
				if( mouse->cursor_y > mouse->res_y - 1 ) {
					mouse->cursor_y = mouse->res_y - 1;
				}

				event.cursor_x = mouse->cursor_x;
				event.cursor_y = mouse->cursor_y;

				mouse->handler( &event, mouse->context );
			}

			break;

		case 3:
			// TODO: read fourth data packet, but don't handle
			// it yet
			break;
	}

	// advance offset
	mouse->offset = ( mouse->offset + 1 ) % mouse->nof_packets;

	return esp;
}

void mouse_print_info( void *obj )
{
	puts( "PS/2 mouse driver" );
}

void mouse_set_resolution( mouse_t *mouse, const uint32_t res_x, const uint32_t res_y )
{
	mouse->res_x = res_x;
	mouse->res_y = res_y;
	mouse->cursor_x = mouse->res_x / 2;
	mouse->cursor_y = mouse->res_y / 2;
}

void mouse_set_position( mouse_t *mouse, const uint32_t x, const uint32_t y )
{
	mouse->cursor_x = x;
	mouse->cursor_y = y;
}