#include <stdio.h>
#include <port_io.h>
#include <tty.h>
#include <stdlib.h>


#include <drivers/ps2_keyboard.h>


/*
  PS/2 Controller IO Ports

  The PS/2 Controller itself uses 2 IO ports (IO ports 0x60 and 0x64). Like many IO ports, reads and writes may access different internal registers.

  Historical note: The PC-XT PPI had used port 0x61 to reset the keyboard interrupt request signal (among other unrelated functions). Port 0x61 has no keyboard related functions on AT and PS/2 compatibles.
  IO Port  Access Type  Purpose
  0x60 	 Read/Write   Data Port
  0x64 	 Read 	     Status Register
  0x64 	 Write 	     Command Register
*/

/*
  Note:
   The interrupt number for keyboard input in 1.
*/

#define PS2_DATA_PORT   0x60
#define PS2_STATUS_PORT 0x64

#define KEYBOARD_IRQ    1

#define RETURN_STI() do { _sti_asm(); return; } while (0)

bool ps2keyboard_initialized = false;

/* State for shift key */
static bool shift_pressed = false;

/* State for caps-lock key */
static bool capslock_pressed = false;

/* Used for arrow keys among others */
static bool extended = false;

static bool is_new_char = false;
static bool is_new_key  = false;

static bool gets_called            = false;
static volatile bool gets_finished = false;

volatile unsigned char current_char;
volatile char* current_string   = NULL;
volatile char* gets_string      = NULL;
volatile int32_t current_length = 0;
volatile int32_t gets_length    = 0;
volatile int32_t capacity       = 0;
volatile int32_t gets_capacity  = 0;
volatile uint16_t current_key;

volatile ps2_hook_t* hooks = NULL;
volatile int hook_count = 0;

extern void _cli_asm(void);
extern void _sti_asm(void);

static const char scancode_map[128] = {
  0,  27, '1','2','3','4','5','6','7','8','9','0','-','=','\b', /* Backspace */
  '\t', /* Tab */
  'q','w','e','r','t','y','u','i','o','p','[',']','\n', /* Enter */
  0, /* Control */
  'a','s','d','f','g','h','j','k','l',';','\'','`',
  0, /* Left shift */
  '\\','z','x','c','v','b','n','m',',','.','/',
  0, /* Right shift */
  '*',
  0, /* Alt */
  ' ', /* Spacebar */
  0, /* Caps lock */
  /* The rest are function and control keys */
};


void keyboard_init(void)
{
#ifdef _DEBUG
  printf("[ PS/2 KBD ] Initializing the PS/2 keyboard...\n");
#endif

  outb(0x64, 0xAE); /* Enable keyboard interface (often optional, but here for safety) */
  
  hooks = (ps2_hook_t*) malloc(sizeof(ps2_hook_t) * hook_count);
  
  ps2keyboard_initialized = true;
  
#ifdef _DEBUG
  printf("[ PS/2 KBD ] PS/2 Keyboard initialized\n");
#endif
}

bool setup_hook(ps2_hook_t func)
{
  _cli_asm();

  ps2_hook_t* copy = malloc(sizeof(ps2_hook_t) * hook_count);
  
  if (hook_count > 0 && copy == NULL)
  {
    return false;
  }
  
  if (hook_count > 0)
  {
    memcpy(copy, hooks, sizeof(ps2_hook_t) * hook_count);
  }

  free((void*) hooks);

  hooks = malloc(sizeof(ps2_hook_t) * (hook_count + 1));
  
  if (hooks == NULL)
  {
    return false;
  }

  if (hook_count > 0)
  {
    memcpy((void*) hooks, copy, sizeof(ps2_hook_t) * hook_count);
  }

  free(copy);

  hooks[hook_count] = func;
  hook_count++;
  
  _sti_asm();
  
  return true;
}


bool remove_hook(ps2_hook_t func)
{
  _cli_asm();

  int index = -1;

  for (int i = 0; i < hook_count; i++)
  {
    if (hooks[i] == func)
    {
      index = i;
      break;
    }
  }

  if (index == -1)
  {
    return false;
  }

  for (int i = index; i < hook_count - 1; i++)
  {
    hooks[i] = hooks[i + 1];
  }

  hook_count--;

  if (hook_count == 0)
  {
    free((void*) hooks);
    hooks = NULL;
    return false;
  }

  ps2_hook_t* smaller = malloc(sizeof(ps2_hook_t) * hook_count);
  
  if (smaller == NULL)
  {
    return false;
  }

  memcpy(smaller, (void*) hooks, sizeof(ps2_hook_t) * hook_count);
  free((void*) hooks);
  hooks = smaller;
  
  _sti_asm();
  
  return true;
}

void call_hooks(char c)
{
  _cli_asm();

  if (hook_count == 0)
  {
    return;
  }

  for (int i = 0; i < hook_count; i++)
  {
    hooks[i](c);
  }
  
  _sti_asm();
}

char get_char(void)
{
  /*
    ASCII code 5 (Enquiry) is/was used in legacy systems meant for requesting a response from a remote terminal or device.
    For example, one system might send ENQ (ASCII 5), and the receiver could reply with ACK (ASCII 6) to indicate it's ready or still online.
    
    In modern computing, ENQ is largely obsolete:

    - Terminal emulators, shells, operating systems, and network protocols generally do not use ENQ.

    - It's not used in text files, programming, or standard communications.

    - It may still appear in legacy systems, embedded devices, or proprietary serial protocols, but that's niche.
  */
  if (!is_new_char)
  {
    return 5;
  }

  is_new_char = false;
  return current_char;
}

uint16_t get_key(void)
{
  uint16_t temp = 0xFFFA;
  
  if (is_new_key)
  {
    temp = current_key;
    is_new_key = false;
  }
  else if (is_new_char)
  {
    temp = (uint16_t)current_char;
    is_new_char = false;
  }
  
  return temp;
}

char* get_string(void)
{
  return current_string;
}

char* kbd_gets(void)
{
  gets_called = true;
  gets_finished = false;
  
  while (gets_finished != true)
  {
    sleep(1);
  }
  
  char* result = strdup(gets_string); /* Could be NULL, that is checked below */
  
  free(gets_string);
  
  gets_string = NULL;
  gets_capacity = 0;
  gets_length = 0;
  gets_called = false;
  
  return result == NULL ? "" : result;
}

void gets_append_char(unsigned char c)
{
  _cli_asm();
  
  if (!gets_string && gets_capacity > 0)
  {
    printf("[keyboard] ERROR: gets_string is NULL but capacity > 0!\n");
    RETURN_STI();
  }

  
  if (c == KEY_ARROW_UP || c == KEY_ARROW_DOWN || c == KEY_ARROW_RIGHT || c == KEY_ARROW_LEFT)
  {
    gets_string[1] = '\0';
    
    char tc = '\0';
    
    switch (c)
    {
      case KEY_ARROW_UP:
        tc = KEY_UP;
        break;
      case KEY_ARROW_DOWN:
        tc = KEY_DOWN;
        break;
      case KEY_ARROW_RIGHT:
        tc = KEY_RIGHT;
        break;
      case KEY_ARROW_LEFT:
        tc = KEY_LEFT;
        break;
      default:
        tc = '\0';
        break;
    }
    
    gets_string[0] = tc;
    
    gets_finished = true;
    
    RETURN_STI();
  }
  else if (c == '\n')
  {
    /* The returned string must/will not have newlines in it, so we return here after setting gets_finished to 'true'. */
  
    gets_finished = true;
    
    printf("\n");
    
    RETURN_STI();
  }
  else if (c == 27)  /* ASCII escape, here it's used to cancel input. */
  {
    gets_finished = true;
    gets_string[0] = '\0';
    
    RETURN_STI();
  }
  else if (c == '\b')
  {
    if (gets_length < 1)
    {
      RETURN_STI();
    }

    gets_length--;
    gets_string[gets_length] = '\0';
    
    printf("\b \b");
    
    RETURN_STI();
  }

  if ((gets_length) >= gets_capacity)
  {
    int new_capacity = (gets_capacity == 0) ? 64 : gets_capacity * 2;

    char* new_str = (char*) malloc(new_capacity);
    
    if (!new_str)
    {
      RETURN_STI();
    }
    
    if (gets_string)
    {
      for (int i = 0; i < gets_length; i++)
      {
        new_str[i] = gets_string[i];
      }
      
      new_str[gets_length] = '\0';
      
      free(gets_string);
    }

    gets_string = new_str;
    gets_capacity = new_capacity;
  }
  
  if (!gets_string)
  {
    RETURN_STI();
  }

  gets_string[gets_length] = (char) c;
  gets_length++;
  gets_string[gets_length] = '\0';
  
  printf("%c", c);
  
  RETURN_STI();
}

void append_char(char c)
{
  _cli_asm();

  if (current_length + 1 >= capacity)
  {
    int new_capacity = (capacity == 0) ? 16 : capacity * 2;

    char* new_str = (char*) malloc(new_capacity);
    
    if (!new_str)
    {
      printf("[keyboard] ERROR: malloc failed for new input buffer\n");
      return;
    }


    if (current_string)
    {
      memcpy(new_str, current_string, current_length);
      free(current_string);
    }

    current_string = new_str;
    capacity = new_capacity;
  }

  if (!current_string)
  {
    return;
  }

  current_string[current_length] = c;
  current_length++;
  current_string[current_length] = '\0';
  
  _sti_asm();
}


void free_current_string(void)
{
  if (current_string)
  {
    free(current_string);
    current_string = NULL;
    current_length = 0;
    capacity = 0;
  }
}

void keyboard_handler(void)
{
  uint8_t scancode = inb(PS2_DATA_PORT);

  /* Handle shift press/release */
  if (scancode == 0x2A || scancode == 0x36)
  {
    shift_pressed = true;
    return;
  }
  else if (scancode == 0xAA || scancode == 0xB6)
  {
    shift_pressed = false;
    return;
  }
  else if (scancode == 0x3A)
  {
    capslock_pressed = !capslock_pressed;
    return;
  }
  else if (scancode == 0xE0)
  {
    extended = true;
    return;
  }


  if (scancode & 0x80)
  {
    extended = false;
  }
  else
  {
    uint16_t c = (uint16_t) scancode_map[scancode];
    if ((shift_pressed ^ capslock_pressed) && ((char) c >= 'a') && ((char) c <= 'z'))
    {
      c -= 32; /* Convert to uppercase */
    }
    else if (shift_pressed)
    {
      c = (uint16_t) terminal_get_shifted((unsigned char) c);
    }
    
    if (hook_count > 0)
    {
      call_hooks(c);
    }

    if (extended)
    {
      switch (scancode)
      {
        case 0x48:
          c = KEY_ARROW_UP;
          break;
        case 0x50:
          c = KEY_ARROW_DOWN;
          break;
        case 0x4B:
          c = KEY_ARROW_LEFT;
          break;
        case 0x4D:
          c = KEY_ARROW_RIGHT;
          break;
        default:
          c = KEY_NONE;
          break;
      }
      
      if (gets_called)
      {
        gets_append_char(c);
        
        return;
      }
      
      current_key = c;
      
      is_new_key = true;
      extended = false;
      
      return;
    }
    
    if (gets_called)
    {
      gets_append_char(c);
      
      return;
    }

    if (c)
    {
      if (c != '\n')
      {
        append_char(c);
      }

      current_char = c;
      is_new_char = true;

      if (c == '\n')
      {
        /*append_char(c); 
        current_char = c;
        is_new_char = true;*/

        free_current_string();
      }
    }

  }
}