#include <acpi.h>
#include <stdbool.h>
#include <io.h>
#include <printf.h>
#include <libc.h>
#include <kernel.h>
#include <cpuid.h>
#include <panic.h>
#include <int.h>
#include <isv.h>


#define PIT_COMMAND   0x43
#define PIT_FREQ      1193182

#define PIT_CHANNEL_0 0x40
#define PIT_CHANNEL_2 0x42
#define PIT_INPUT_2   0x61
//you can move these if it makes more sense to put them in a lapic.h
#define LAPIC_TIMER_LVT             200
#define LAPIC_TIMER_INIT_COUNT      224
#define LAPIC_TIMER_CURRENT_COUNT   228
#define LAPIC_TIMER_DIVIDE          248

#define LAPIC_TIMER_MODE_ONESHOT       0b00
#define LAPIC_TIMER_MODE_PERIODIC      0b01
#define LAPIC_TIMER_MODE_TSC_DEADLINE  0b10

#define IA32_TSC_DEADLINE 0x6E0 


struct acpi_addr_fmt {
  uint8_t addr_space; //0 is system memory, 1 is system io
  uint8_t reg_bit_width;
  uint8_t reg_bit_offset; 
  uint8_t resrved;
  uint64_t addr;
} __attribute__((packed));

struct hpet_sdt {
  sdt_head header;
  uint8_t hw_rev_id;
  unsigned int comparator_cnt:5;
  unsigned int counter_size_64bit:1;
  unsigned int reserved:1;
  unsigned int legacy_capable:1;
  uint16_t vendor_id;
  struct acpi_addr_fmt acpi_addr;
  uint8_t hpet_number;
  uint16_t min_tick;
  unsigned int protection:4; //0 = no protection, 1 = 4kb, 2 = 64kb
  unsigned int oem_attribute:4;
} __attribute__((packed));

struct hpet_info_struct { 
  uint64_t *hpet_reg;
  unsigned int hpet_period;
  double hpet_freq_mhz; //improves speed with us quantum
  unsigned int timer_cnt;
  unsigned int irqs; 
  bool long_mode;
} hpet_info;


struct lapic_lvt {
  uint8_t vector;
  unsigned int reserved:4;
  unsigned int delivery_status:1;
  unsigned int reserved_1:3;
  unsigned int mask:1;
  unsigned int timer_mode:2;
  unsigned int reserved_2:13;
} __attribute__((packed));

extern lapic_t lapic;

#define TSC_TIME_TSC_EVENT        0
#define TSC_TIME_LAPIC_EVENT      1
#define LAPIC_TIME_HPET_EVENT     2

static int clock_mode;

//time quantum is 1us
static double tsc_freq;
static double apic_freq;

static unsigned int early_event_vector;

static struct lapic_lvt lapic_timer_lvt;

uint64_t timestamp() {
  if((clock_mode == TSC_TIME_TSC_EVENT) || (clock_mode == TSC_TIME_LAPIC_EVENT)) {
    return read_tsc() / tsc_freq;
  }
  else {
    return UINT32_MAX - (lapic[LAPIC_TIMER_CURRENT_COUNT] / apic_freq);
  }
  return 0;
}

//timer test code
//there is a race condition here.
//I'm leaving it unattended, as I'm only expecting this to be used for SMP booting.
void usleep(unsigned int us) {
  if(clock_mode == TSC_TIME_TSC_EVENT) {
    write_msr(IA32_TSC_DEADLINE, ((uint64_t)((us * (double)tsc_freq) + 0.5)) + read_tsc());
  }
  else if(clock_mode == TSC_TIME_LAPIC_EVENT) {
    lapic[LAPIC_TIMER_INIT_COUNT] = (uint32_t)((us * (double)apic_freq) + 0.5);
  }
  else {
    hpet_info.hpet_reg[33] = hpet_info.hpet_reg[30] + 
      (uint32_t)((us * hpet_info.hpet_freq_mhz) + 0.5);
  }
  asm("hlt\n");
}

/**
we either use:
1: tsc
2: tsc

1: tsc
2: lapic (no tsc deadline)

1: lapic
2: hpet

PIT not implimented, I've yet to find a computer that doesn't have HPET
**/

void calibrate_lapic() {
  unsigned int timer_idt = alloc_idt(&KERNEL_IDT_GATE(lapic_timer_racefixer));
  uint32_t lapic_ticks;
  uint64_t hpet_ticks = (100000000000000 * (1 / (double)hpet_info.hpet_period)) + 0.5;
  uint64_t timer_loops = 0;
  struct lapic_lvt timer_lvt = {
    .vector = timer_idt,
    .delivery_status = 0,
    .mask = 0,
    .timer_mode = LAPIC_TIMER_MODE_PERIODIC
  };
  printf("Starting LAPIC timer calibration...\n");
  
  lapic[LAPIC_TIMER_LVT] = *(uint32_t *)&timer_lvt;
  lapic[LAPIC_TIMER_DIVIDE] = 0b1011;
  asm(".global calibrate_lapic_waiting\n"
      "xor rbx, rbx\n"
      "mov rax, %2\n"
      "add rax, %3\n"
      "mov [%4], rax\n"       //hpet timer
      "mov [%5], 0xffffffff\n"     //lapic timer
      "calibrate_lapic_waiting:\n"
      "hlt\n"
      "mov %0, %6\n"          //save lapic time
      "mov [%5], 0\n"
      "mov %1, rbx\n"
      :"=r"(lapic_ticks), "=r"(timer_loops)
      :"r"(hpet_ticks), "m"(hpet_info.hpet_reg[30]), "m"(hpet_info.hpet_reg[33]), 
      "m"(lapic[LAPIC_TIMER_INIT_COUNT]), "m"(lapic[LAPIC_TIMER_CURRENT_COUNT])
      :"rax","rbx" );
  apic_freq = ((UINT32_MAX * timer_loops) + (UINT32_MAX - lapic_ticks)) / (double)100000;

  printf("LAPIC timer frequency: %f MHz\n", apic_freq);
  free_idt(timer_idt);
}


void init_timer() {
  struct hpet_sdt *hpet_desc = find_sdt(SDT_HPET);
  early_event_vector = alloc_idt(&KERNEL_IDT_GATE(kernel_block));
  if(!hpet_desc) PANIC(KERNEL_PANIC_HPET_REQUIRED);

  unsigned int first_irq;

  struct apic_vt hpet_redirect;
  bzero(&hpet_redirect, sizeof(hpet_redirect));
  hpet_redirect.vector = early_event_vector;
  hpet_redirect.polarity = 1;
  

  hpet_info.hpet_reg = PHYS_TO_VIRT(hpet_desc->acpi_addr.addr);
  hpet_info.hpet_period = (hpet_info.hpet_reg[0] >> 32) & UINT32_MAX; 
  hpet_info.long_mode = (hpet_info.hpet_reg[0] >> 13) & 1;
  hpet_info.irqs = (hpet_info.hpet_reg[32] >> 32) & UINT32_MAX;
  hpet_info.hpet_freq_mhz = (1 / (double)hpet_info.hpet_period) * 1000000000;
  //get first irq available
  first_irq = __builtin_ctzl(hpet_info.irqs) + 1;

  printf("hpet frequency: %f MHz\n", hpet_info.hpet_freq_mhz);

  //do we need to worry about bypassing irq remappings?
  create_fixed_interrupt(first_irq, &hpet_redirect); 

  //set irq, enable triggering of interrupts
  hpet_info.hpet_reg[32] = (first_irq << 9) | (1 << 2) | 1;

  //enable main counter
  hpet_info.hpet_reg[2] = 1; 

  asm("sti\n");

  //detect consistent TSC
  uint32_t unused, cpuid_reg;
  __get_cpuid(0x80000007, &unused, &unused, &unused, &cpuid_reg); 

  //goto debug_tsc;
  if((cpuid_reg >> 8) & 1) {
    printf("Detected invariant TSC\n");
    //.1 second to calibrate, TODO do we need to check if the register is big enough?
    uint64_t hpet_ticks = (100000000000000 * (1 / (double)hpet_info.hpet_period)) + 0.5;
    printf("Starting TSC calibration...\n");
    uint64_t volatile start_tsc, end_tsc;
    //this may not be the fastest, but I'll improve if needed further into development
     asm volatile(
        "xor rax, rax\n"
        "rdtsc\n"
        "shl rdx, 32\n"
        "or rax, rdx\n"
        "mov %0, rax\n" 
        "mov rbx, %4\n" 
        "add rbx, [%3]\n"
        "mov [%2], rbx\n"
        //if you're worried about a single instruction being a race condition, it's time for an upgrade
        "xor rax, rax\n" 
        "hlt\n"
        "rdtsc\n"
        "shl rdx, 32\n"
        "or rax, rdx\n"
        "mov %1, rax\n"
         :"=&r"(start_tsc), "=&r"(end_tsc)
         :"m"(hpet_info.hpet_reg[33]), "m"(hpet_info.hpet_reg[30]),"r"(hpet_ticks)
         :"rax","rbx","rdx");
    tsc_freq = (end_tsc - start_tsc) / (double)100000;
    printf("TSC: Detected frequency is %f MHz\n", tsc_freq);
    __get_cpuid(0x1, &unused, &unused, &cpuid_reg, &unused);
    if((cpuid_reg >> 24) & 1) {
      clock_mode = TSC_TIME_TSC_EVENT;

      lapic_timer_lvt.vector = early_event_vector; 
      lapic_timer_lvt.timer_mode = LAPIC_TIMER_MODE_TSC_DEADLINE;
      lapic_timer_lvt.delivery_status = 0;
      lapic_timer_lvt.mask = 0;
    }
    else {
      clock_mode = TSC_TIME_LAPIC_EVENT;
      calibrate_lapic();

      lapic_timer_lvt.vector = early_event_vector; 
      lapic_timer_lvt.timer_mode = LAPIC_TIMER_MODE_ONESHOT;
      lapic_timer_lvt.delivery_status = 0;
      lapic_timer_lvt.mask = 0;
    }
    lapic[LAPIC_TIMER_LVT] = *(uint32_t *)&lapic_timer_lvt;
  }
  else {
    uint32_t apic_div;
    clock_mode = LAPIC_TIME_HPET_EVENT;
    calibrate_lapic();

    //because lapic is used as scheduler, we want to cause as little interrupts as possible
    //while retaining 1 mhz accuracy

    apic_div = __builtin_clz((uint32_t)apic_freq);
    if(!apic_div) {
      apic_div = 0b1011;
    }
    else if(apic_div >= 7) {
      apic_div = 0b1010;
      apic_freq /= 128;
    }
    else {
      apic_freq /= (1 << --apic_div);
      apic_div = (((apic_div & 0b100) << 1) | (apic_div * 0b1011));
    }

    lapic_timer_lvt.vector = SPURRIOUS_VECTOR; //TODO CHANGE ME
    lapic_timer_lvt.timer_mode = LAPIC_TIMER_MODE_PERIODIC;
    lapic_timer_lvt.delivery_status = 0;
    lapic_timer_lvt.mask = 1;
    lapic[LAPIC_TIMER_DIVIDE] = apic_div;
    lapic[LAPIC_TIMER_LVT] = *(uint32_t *)&lapic_timer_lvt;
    lapic[LAPIC_TIMER_INIT_COUNT] = UINT32_MAX;
  }
  switch(clock_mode) {
    case TSC_TIME_TSC_EVENT:
      printf("Clock source: tsc\nEvent source: tsc-deadline\n");
      break;
    case TSC_TIME_LAPIC_EVENT:
      printf("Clock source: tsc\nEvent source: lapic\n");
      break;
    case LAPIC_TIME_HPET_EVENT:
      printf("Clock source: lapic\nEvent source: hpet\n");
      break;
  }
}