
#include <multiboot.h>
#include <lib/stdint.h>
#include <lib/vga.h>
#include <lib/stdio.h>
#include <gdt.h>
#include <memorymanagement.h>
#include <hardwarecommunication/interrupts.h>
#include <syscalls.h>
#include <hardwarecommunication/pci.h>
#include <drivers/driver.h>
#include <drivers/keyscancodes.h>
#include <drivers/keyboard.h>
#include <drivers/mouse.h>
#include <drivers/ata.h>
#include <drivers/serial.h>
#include <drivers/speaker.h>
#include <drivers/rtc.h>
#include <drivers/pit.h>
#include <multitasking.h>
#include <filesystem/fat.h>
#include <c64/c64.h>
#include <c64/monitor.h>

uint32_t current_milli = 0;

multiboot_info_t* mboot_hdr;     
multiboot_info_t *verified_mboot_hdr;
uint32_t mboot_reserved_start;
uint32_t mboot_reserved_end;
uint32_t next_free_frame;

using namespace myos;
using namespace myos::drivers;
using namespace myos::hardwarecommunication;
using namespace myos::filesystem;

// for accessing and sending keyboard or 
// other external data to c64's IO system

C64* c64ptr;	
uint32_t* framebuffer_addr;

// ---------------------------------------------------------------------------
// Bochs / QEMU stdvga (dispi) VBE bring-up.
// QEMU's -kernel multiboot loader does NOT provide a VBE framebuffer tag
// ("multiboot knows VBE. we don't"), so we set the mode ourselves via the
// Bochs DISPI registers and locate the linear framebuffer through PCI BAR0.
// ---------------------------------------------------------------------------
#define VBE_DISPI_IOPORT_INDEX  0x01CE
#define VBE_DISPI_IOPORT_DATA   0x01CF
#define VBE_DISPI_INDEX_ID      0x0
#define VBE_DISPI_INDEX_XRES    0x1
#define VBE_DISPI_INDEX_YRES    0x2
#define VBE_DISPI_INDEX_BPP     0x3
#define VBE_DISPI_INDEX_ENABLE  0x4
#define VBE_DISPI_ENABLED       0x01
#define VBE_DISPI_LFB_ENABLED   0x40

// Globals holding the real framebuffer geometry we set up.
uint32_t g_fb_addr   = 0;
uint32_t g_fb_width  = 0;
uint32_t g_fb_height = 0;
uint32_t g_fb_pitch  = 0;
uint8_t  g_fb_bpp    = 0;

static inline void vbe_write(uint16_t index, uint16_t data)
{
    hardwarecommunication::Port16Bit idx(VBE_DISPI_IOPORT_INDEX);
    hardwarecommunication::Port16Bit dat(VBE_DISPI_IOPORT_DATA);
    idx.Write(index);
    dat.Write(data);
}

static inline uint16_t vbe_read(uint16_t index)
{
    hardwarecommunication::Port16Bit idx(VBE_DISPI_IOPORT_INDEX);
    hardwarecommunication::Port16Bit dat(VBE_DISPI_IOPORT_DATA);
    idx.Write(index);
    return dat.Read();
}

// Program the DISPI mode. Returns true if the DISPI interface is present.
static bool vbe_set_mode(uint16_t w, uint16_t h, uint16_t bpp)
{
    uint16_t id = vbe_read(VBE_DISPI_INDEX_ID);
    if(id < 0xB0C0 || id > 0xB0CF)
        return false; // no Bochs DISPI present
    vbe_write(VBE_DISPI_INDEX_ENABLE, 0);           // disable while programming
    vbe_write(VBE_DISPI_INDEX_XRES, w);
    vbe_write(VBE_DISPI_INDEX_YRES, h);
    vbe_write(VBE_DISPI_INDEX_BPP,  bpp);
    vbe_write(VBE_DISPI_INDEX_ENABLE, VBE_DISPI_ENABLED | VBE_DISPI_LFB_ENABLED);
    return true;
}

// Find the stdvga linear framebuffer via PCI. QEMU stdvga = 1234:1111,
// BAR0 is the LFB. Fall back to the classic Bochs LFB base if not found.
static uint32_t vbe_find_lfb()
{
    hardwarecommunication::PeripheralComponentInterconnectController pci;
    for(int bus=0; bus<8; bus++)
      for(int dev=0; dev<32; dev++)
        for(int fn=0; fn<8; fn++)
        {
            uint32_t vend = pci.Read(bus,dev,fn,0x00);
            uint16_t vid = vend & 0xFFFF;
            uint16_t did = (vend>>16) & 0xFFFF;
            if(vid==0xFFFF || vid==0x0000) continue;
            // QEMU stdvga (1234:1111) or bochs vga
            if((vid==0x1234 && did==0x1111) || (vid==0x1111))
            {
                uint32_t bar0 = pci.Read(bus,dev,fn,0x10);
                if(bar0 & 0x1) continue;            // I/O BAR, skip
                return bar0 & 0xFFFFFFF0;           // memory BAR base
            }
        }
    return 0xE0000000; // common QEMU stdvga LFB fallback
}

class IOKeyboardEventHandler : public KeyboardEventHandler
{
private:
  uint8_t mode;  // 0 = emulation, 1 = terminal
public:
    IOKeyboardEventHandler()
    {
      mode = 0;
    }

    void OnKeyDown(uint8_t c)
    {
      { extern volatile unsigned int g_hk_count, g_hk_last; g_hk_count++; g_hk_last=c; }
      // ESC key will toggle between text mode and emulation
      if(c == 0x01) 
      {
	if (mode == 0) 
	{ 
	  mode = 1; 
	  c64ptr->stop();
	  return;
	}
	else 
	{ 
	  mode = 0; 
	  c64ptr->mon_->Stop();
	  return;
	}
      }
      
      if(c == 0x03) // Reset machine (F10)
      {
	c64ptr->reset = true;
	return;
      }
      
      if(c == 0x04) // toggle skipframes (F12)
      {
	if(c64ptr->io_->SkipFrames == 10)
	  c64ptr->io_->SkipFrames = 0;
	else
	  c64ptr->io_->SkipFrames = 10;
      }
      
      switch(mode)
      {
	case 0: c64ptr->io_->OnKeyDown(c); break;
	case 1: c64ptr->mon_->OnKeyDown(c); break;
      }
    }
    
    void OnKeyUp(uint8_t c)
    {
      switch(mode)
      {
	case 0: c64ptr->io_->OnKeyUp(c); break;
	case 1: c64ptr->mon_->OnKeyUp(c); break;
      }
    }
};


class IOPITEventHandler : public PITEventHandler
{
private:
public:
  IOPITEventHandler() 
  {
  }
  
  void OnTick() 
  {
    current_milli++;
    // ── 1000Hz KEYBOARD POLL ──────────────────────────────────────────────
    // The PIT fires at 1000Hz. Drain the 8042 controller output buffer into
    // the scancode ring here, so keyboard input is sampled every millisecond
    // REGARDLESS of how slow the C64 emulation loop runs. This decouples input
    // latency from CPU work entirely — a true 1000Hz tick rate.
    {
      extern volatile unsigned char g_scan_ring[256];
      extern volatile unsigned int  g_scan_head;
      for(int guard=0; guard<16; guard++){
        uint8_t st = myos::hardwarecommunication::Port8Bit(0x64).Read();
        if(!(st & 0x01)) break;
        uint8_t sc = myos::hardwarecommunication::Port8Bit(0x60).Read();
        g_scan_ring[g_scan_head & 0xFF] = sc; g_scan_head++;
      }
    }
  }
};

// Set up C++ object constructors.  This has to be set up and called manually
typedef void (*constructor)();
extern "C" constructor start_ctors;
extern "C" constructor end_ctors;
extern "C" void callConstructors()
{
    for(constructor* i = &start_ctors; i != &end_ctors; i++)
        (*i)();
}

extern "C" void kernelMain(const void* multiboot_structure, uint32_t /*multiboot_magic*/)
{
    mboot_hdr = (multiboot_info_t *)multiboot_structure;
    
    if ((mboot_hdr->flags & (1<<6)) == 0) {
        // No memory map (e.g. QEMU -kernel). Continue with a safe default
        // instead of halting; memupper below still yields a usable heap.
    }
    
    verified_mboot_hdr = mboot_hdr;
    mboot_reserved_start = (uint32_t)mboot_hdr;
    mboot_reserved_end = (uint32_t)(mboot_hdr + sizeof(multiboot_info_t));
    next_free_frame = 1;
    /* Framebuffer selection, metal-safe:
       1) Prefer the framebuffer GRUB actually set up and handed us in the
          multiboot info (this is the REAL address on real hardware).
       2) Only if GRUB gave us nothing usable, try Bochs/QEMU DISPI VBE
          (works under QEMU/Bochs; most real GPUs don't implement it). */
    {
      bool have_mb_fb =
        (mboot_hdr->flags & MULTIBOOT_INFO_FRAMEBUFFER_INFO) &&
        mboot_hdr->framebuffer_addr != 0 &&
        mboot_hdr->framebuffer_width >= 320 &&
        mboot_hdr->framebuffer_bpp  >= 15;   // graphics LFB, not text mode

      if(have_mb_fb)
      {
        g_fb_addr   = (uint32_t)mboot_hdr->framebuffer_addr;
        g_fb_width  = (uint32_t)mboot_hdr->framebuffer_width;
        g_fb_height = (uint32_t)mboot_hdr->framebuffer_height;
        g_fb_pitch  = (uint32_t)mboot_hdr->framebuffer_pitch;
        g_fb_bpp    = (uint8_t)mboot_hdr->framebuffer_bpp;
      }
      else
      {
        uint16_t W=800, H=600, BPP=16;
        bool ok = vbe_set_mode(W, H, BPP);
        uint32_t lfb = vbe_find_lfb();
        if(ok && lfb)
        {
          g_fb_addr   = lfb;
          g_fb_width  = W;
          g_fb_height = H;
          g_fb_pitch  = W * (BPP/8);
          g_fb_bpp    = BPP;
        }
        else
        {
          g_fb_addr   = (uint32_t)mboot_hdr->framebuffer_addr;
          g_fb_width  = (uint32_t)mboot_hdr->framebuffer_width;
          g_fb_height = (uint32_t)mboot_hdr->framebuffer_height;
          g_fb_pitch  = (uint32_t)mboot_hdr->framebuffer_pitch;
          g_fb_bpp    = (uint8_t)mboot_hdr->framebuffer_bpp;
        }
      }
    }

    framebuffer_addr = (uint32_t*)g_fb_addr;

    vga_init((uint32_t*)g_fb_addr, g_fb_width, g_fb_height, g_fb_pitch, g_fb_bpp);

    vga_clear();
    
   
    printf("OS/64 Operating System Starting...\n");

    GlobalDescriptorTable gdt;
    
    // Get start of memory that we can safely work from
    uint32_t* memupper = (uint32_t*)(((size_t)multiboot_structure) + 8);

    // ── PROOF OF LIFE (cannot fault): write to legacy VGA text memory at
    //    0xB8000, which is ALWAYS present on x86 regardless of graphics mode.
    //    If you see colored letters top-left, the kernel is alive and running;
    //    the problem is then specifically the graphics framebuffer. We spell
    //    "HDGL" in bright colors on the text plane. This runs before any
    //    graphics framebuffer access, so it can't triple-fault us. ──
    {
      volatile uint16_t* txt = (volatile uint16_t*)0x000B8000;
      const char* s = "HDGL BOOT OK";
      for(int i=0; s[i]; i++) txt[i] = (uint16_t)s[i] | 0x1F00; // white on blue
    }

    // ── BOOT PROGRESS MARKERS (metal diagnostic) ──
    // Paint a small square at a fixed screen slot for each boot stage. On real
    // hardware, the LAST square you see (frozen or before reboot) tells us how
    // far boot got. Each slot is 40px, left-to-right along the top.
    //   slot 0 GREEN  = heap about to init
    //   slot 1 CYAN   = heap done, about to construct C64
    //   slot 2 YELLOW = C64 constructed, display initialised
    //   slot 3 WHITE  = entering C64 run loop (first tick)
    #define BOOTMARK(slot,col) do{ \
        if(g_fb_addr && g_fb_pitch && g_fb_bpp>=15){ \
          volatile uint8_t* _fb=(volatile uint8_t*)g_fb_addr; \
          for(uint32_t _y=0;_y<32;_y++){ volatile uint8_t* _r=_fb+_y*g_fb_pitch; \
            for(uint32_t _x=(slot)*40;_x<(slot)*40+32;_x++){ \
              if(g_fb_bpp>=24){_r[_x*4]=((col)&0x1F)<<3;_r[_x*4+1]=(((col)>>5)&0x3F)<<2;_r[_x*4+2]=(((col)>>11)&0x1F)<<3;} \
              else {_r[_x*2]=(col)&0xFF;_r[_x*2+1]=((col)>>8)&0xFF;} } } } }while(0)

    BOOTMARK(0, 0x07E0);   // GREEN: heap about to init

    // Define 10MB of heap space for dynamic memory allocation
    size_t heap = 10*1024*1024;
    MemoryManager memoryManager(heap, (*memupper)*1024 - heap - 10*1024);

    BOOTMARK(1, 0x07FF);   // CYAN: heap constructed
    
    printf("\n%d Byte Heap @ 0x%08X", heap);
    printf("\n\n");
    
    // How to use memorymanager to allocate dynamic memory from the heap
    //uint8_t *ram1 = (uint8_t *) memoryManager.malloc(65535);
    //uint8_t* ram1 = new uint8_t[65535];
    //printf("\n64KB RAM BANK 1: %06X", ram1);
       
    TaskManager taskManager;
    
    InterruptManager interrupts(0x20, &gdt, &taskManager);
    SyscallHandler syscalls(&interrupts, 0x80);

    DriverManager drvManager;
    
    IOKeyboardEventHandler kbhandler;
    KeyboardDriver keyboard(&interrupts, &kbhandler);
    drvManager.AddDriver(&keyboard);
    
    SerialEventHandler serialhandler;
    SerialDriver serial(&interrupts, &serialhandler);
    drvManager.AddDriver(&serial);
    
    IOPITEventHandler timerhandler;
    PITDriver pit(&interrupts, &timerhandler);
    drvManager.AddDriver(&pit);
    
    PeripheralComponentInterconnectController PCIController;
    PCIController.SelectDrivers(&drvManager, &interrupts);
    drvManager.ActivateAll();
       
    printf("Initializing interrupts..........[OK]\n");
    interrupts.Activate();
       
    printf("\nATA pri master: ");
    AdvancedTechnologyAttachment ata0m(true, _ATA_FIRST);  
    ata0m.Identify();
    
    printf("\nATA pri slave : ");
    AdvancedTechnologyAttachment ata0s(false, _ATA_FIRST);  
    ata0s.Identify();
    
    printf("\nATA sec master: ");
    AdvancedTechnologyAttachment ata1m(true, _ATA_SECOND);  
    ata1m.Identify();
    
    printf("\nATA sec slave : ");
    AdvancedTechnologyAttachment ata1s(false, _ATA_SECOND);  
    ata1s.Identify();
    
    //uint32_t s1 = 0x0FFFFFFE;
    //s1 = 0;
    //ata0m.WriteSector(s1, (uint8_t*)"LBA48-Z", 7);  
    //ata0m.WriteSector(s1+1, (uint8_t*)"LBA48-V", 7);
    //ata0m.WriteSector(s1+2, (uint8_t*)"LBA48-O", 7);
    
    /*
    uint8_t buffer[512];
    ata0m.ReadSector(s1, buffer); displayMemory(buffer, 16);
    ata0m.ReadSector(s1+1, buffer); displayMemory(buffer, 16);
    ata0m.ReadSector(s1+2, buffer); displayMemory(buffer, 16);
    // ata0m.Read28(0, buffer); displayMemory(buffer, 16);
    
    while(1) {};*/
    
    //ata0m.Write28(268435456, (uint8_t*)"LBA28-0", 7);
    //ata0m.Write28(268435457, (uint8_t*)"LBA48-1", 7);
    //ata0m.Write28(268990000, (uint8_t*)"LBA28-2", 7);
    
    //while(1) {};
    
    //displayMemory(sector, 512);
    
    printf("\n\nInitializing filesystem driver...");
    Fat32 fat32(&ata0m,0);
    printf("[OK]");

    //fat32.ReadPartitions();
    //fat32.ReadDir();
    //fat32.WriteDir((uint8_t*)"12345678",(uint8_t*)"EXT", 32);

    //fat32.ReadFile("COMMAND .COM");
    //fat32.ReadFile("TEST0001.TXT");
    
    //while(1) {};
    
    printf("\nObtaining video mode.............[OK]");
    printf(" Screen %d X %d X %d - Pitch %d", (uint32_t)mboot_hdr->framebuffer_width,
	   (uint32_t)mboot_hdr->framebuffer_height, (uint8_t)mboot_hdr->framebuffer_bpp,
	   (uint32_t)mboot_hdr->framebuffer_pitch);
    printf("\n                                      VGA Memory @ $%08X",(uint32_t)mboot_hdr->framebuffer_addr);

    SpeakerDriver speaker;

    RTCDriver rtc;
    struct datetime* curDateTime;
    rtc.GetRTC(curDateTime);
    
    printf("\n\nGetting date: %d/%d/%d", curDateTime->month, curDateTime->day, curDateTime->year);
    printf("\nGetting time: %d:%d:%d", curDateTime->hour, curDateTime->minute, curDateTime->second);

    while(true)
    {
      C64 c64;
      c64ptr = &c64;
      c64ptr->sid_->speaker(&speaker);
      c64ptr->io_->init_display((uint32_t*)g_fb_addr, g_fb_width,
			  g_fb_height, g_fb_pitch,
			  g_fb_bpp);
      BOOTMARK(2, 0xFFE0);   // YELLOW: C64 constructed + display initialised
      
      c64ptr->io_->fat32(&fat32);
      c64ptr->io_->serial(&serial);
      c64ptr->io_->rtc(&rtc);
      c64ptr->mon_->fat32(&fat32);
      BOOTMARK(3, 0xFFFF);   // WHITE: about to enter the C64 run loop
      
      c64.start();
      //TODO: Memory leak possible here.  Calling destructor at the moment causes freeze
    }
    
}


