#include "ps2mouse.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( ps2mouse_t *ps2mouse )
{
	while( ( port8_read( &ps2mouse->command_port ) & STATUS_REG_OUTPUT_BUF_FULL ) == 0 ) { }
	return port8_read( &ps2mouse->data_port );
}

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

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

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

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

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

static ps2mouse_vtable_t const ps2mouse_vtable = {
	{
		{
			ps2mouse_activate,
			ps2mouse_deactivate,
			driver_deinit,
			ps2mouse_print_name,
			ps2mouse_print_info
		},
		mouse_set_resolution,
		mouse_set_position
	}
};

void ps2mouse_init( ps2mouse_t *ps2mouse, mouse_event_handler_t handler, interrupt_t *interrupt, void *context )
{
	memset( ps2mouse, 0, sizeof( ps2mouse_t ) );

	mouse_init( (mouse_t *)ps2mouse, handler, interrupt, context );

	ps2mouse->nof_packets = DEFAULT_NOF_PACKETS;

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

	((driver_t *)ps2mouse)->vtable = (driver_vtable_t *)&ps2mouse_vtable;
}

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

	driver_t *driver = obj;
	ps2mouse_t *ps2mouse = obj;

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

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

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

	interrupt_handler_init( &ps2mouse->interrupt_handler, IRQ_BASE + 0x0C, driver->interrupt, ps2mouse_handle_interrupt, obj );
	interrupts_register_interrupt_handler( ps2mouse->interrupt_handler );

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

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

	ps2mouse_t *ps2mouse = obj;

	send_command( ps2mouse, COMMAND_DISABLE_PORT2 );
}

uint32_t ps2mouse_handle_interrupt( interrupt_handler_t *handler, uint32_t esp )
{
	ps2mouse_t *ps2mouse = (ps2mouse_t *)handler->driver;

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

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

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

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

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

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

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

				event.cursor_x = ps2mouse->base.cursor_x;
				event.cursor_y = ps2mouse->base.cursor_y;

				ps2mouse->base.handler( &event, ps2mouse->base.base.context );
			}

			break;

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

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

	return esp;
}

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

void ps2mouse_print_info( void *obj )
{
	// TODO: print IRQ, mouse type, wheel, buttons, etc.
}