#include "renderer.h"
#include "gpu.h"
#include "types.h"
#include <SDL3/SDL.h>
#include <SDL3/SDL_gpu.h>
#include <SDL3/SDL_gamepad.h>
#include <SDL3/SDL_oldnames.h>
#include <SDL3/SDL_pixels.h>
#include <SDL3/SDL_timer.h>
#include <SDL3_image/SDL_image.h>
#include <iomanip>
#include <iostream>
#include "bsp/BSP.h"

#include "imgui.h"
#include "imgui_impl_sdl3.h"
#include "imgui_impl_sdlgpu3.h"


extern int post_mode;
extern vec2 atlas_adjust;

bool Renderer::init(SDL_Window* w) {
  win = w;
  dev = SDL_CreateGPUDevice(SDL_GPU_SHADERFORMAT_SPIRV, true, NULL);
  if (!dev) { std::cerr << "GPU device failed: " << SDL_GetError() << std::endl; return false; }

  SDL_ClaimWindowForGPUDevice(dev, win);

  printf("LOAD BIN?\n");
  load_bin();
  printf("BIN LOADED\n");

  render_sampler = create_linear_sampler(dev);
  SDL_GPUSamplerCreateInfo depth_sampler_info {
    .min_filter    = SDL_GPU_FILTER_NEAREST,
    .mag_filter    = SDL_GPU_FILTER_NEAREST,
    .mipmap_mode   = SDL_GPU_SAMPLERMIPMAPMODE_NEAREST,
    .address_mode_u = SDL_GPU_SAMPLERADDRESSMODE_MIRRORED_REPEAT,
    .address_mode_v = SDL_GPU_SAMPLERADDRESSMODE_MIRRORED_REPEAT,
    .compare_op    = SDL_GPU_COMPAREOP_INVALID,
    .enable_compare = false,
  };
  
  depth_sampler = SDL_CreateGPUSampler(dev, &depth_sampler_info);
  atlas_sampler = SDL_CreateGPUSampler(dev, &depth_sampler_info);
  
  resize_render_texture(960, 540);

  setup_imgui();

  fprintf(stderr, "Pipeline creation....");
  scene.create(dev, win);
  ray.create(dev, win);
  post.create(dev, win);
  fprintf(stderr, "finished\n");

  return scene.pipeline && ray.pipeline;
}

void Renderer::destroy() {
  if (render_tex)     SDL_ReleaseGPUTexture(dev, render_tex);
  if (post_tex)       SDL_ReleaseGPUTexture(dev, post_tex);
  if (render_sampler) SDL_ReleaseGPUSampler(dev, render_sampler);
  if (vert_buff)      SDL_ReleaseGPUBuffer(dev, vert_buff);
  if (depth_tex)      SDL_ReleaseGPUTexture(dev, depth_tex);
  if (gamepad)        SDL_CloseGamepad(gamepad);
  scene.destroy(dev);
  ray.destroy(dev);
  post.destroy(dev);
  SDL_DestroyGPUDevice(dev);
}

bool Renderer::load_bin() {
  const char* path = "assets/bin/hl1.bin";
  const char* tex_path = "assets/bin/hl1.bin.png";
  /*if(!std::filesystem::exists(path) || !std::filesystem::exists(tex_path)) {
    SDL_Log("Failed to find %s (or maybe %s)", path, tex_path);
    return false;
  }*/

  size_t bin_size;
  const uint8_t* bin_data = (uint8_t*)SDL_LoadFile(path, &bin_size);
  if(!bin_data) {
    SDL_Log("Failed to read %s %s", path, tex_path);
    return false;
  }
  vertices.resize(bin_size/sizeof(Vertex));
  memcpy(vertices.data(), bin_data, bin_size);
  SDL_Log("File size: %zu, %zu vertices", bin_size, bin_size/3);


  SDL_GPUCommandBuffer* cmd = SDL_AcquireGPUCommandBuffer(dev);
  SDL_GPUCopyPass *copy_pass = SDL_BeginGPUCopyPass(cmd);
  int w,h;
  atlas_tex = IMG_LoadGPUTexture(dev, copy_pass, tex_path, &w, &h);

  SDL_EndGPUCopyPass(copy_pass);
  SDL_SubmitGPUCommandBuffer(cmd);
  SDL_GPUBufferCreateInfo info {
    .usage = SDL_GPU_BUFFERUSAGE_VERTEX,
    .size  = static_cast<uint32_t>(vertices.size() * sizeof(Vertex)),
  };
  vert_buff = SDL_CreateGPUBuffer(dev, &info);
  upload_buffer(dev, vert_buff, vertices.data(), info.size);
  return true;
}

void Renderer::setup_imgui() {
  IMGUI_CHECKVERSION();
  ImGui::CreateContext();
  ImGuiIO& io = ImGui::GetIO();
  io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
  io.ConfigFlags |= ImGuiConfigFlags_NavEnableGamepad;

  ImGui::StyleColorsDark();

  ImGui_ImplSDL3_InitForSDLGPU(win);
  ImGui_ImplSDLGPU3_InitInfo init_info = {
    .Device = dev,
    .ColorTargetFormat = SDL_GetGPUSwapchainTextureFormat(dev, win),
    .MSAASamples = SDL_GPU_SAMPLECOUNT_1,
    .SwapchainComposition = SDL_GPU_SWAPCHAINCOMPOSITION_SDR,
    .PresentMode = SDL_GPU_PRESENTMODE_VSYNC,
  };

  ImGui_ImplSDLGPU3_Init(&init_info);
}

bool Renderer::load_bsp() {
  const std::string& path = "assets/hl1.bsp";
  HLBSP::BSP bsp(path);
  bsp.load_vertices();
  vertices = bsp.textured_vertices;

  SDL_GPUBufferCreateInfo info {
    .usage = SDL_GPU_BUFFERUSAGE_VERTEX,
    .size  = static_cast<uint32_t>(vertices.size() * sizeof(Vertex)),
  };
  vert_buff = SDL_CreateGPUBuffer(dev, &info);
  upload_buffer(dev, vert_buff, vertices.data(), info.size);
  return true;
}

void Renderer::resize_render_texture(uint32_t w, uint32_t h) {
  SDL_WaitForGPUIdle(dev);
  if (render_tex) SDL_ReleaseGPUTexture(dev, render_tex);
  if (post_tex) SDL_ReleaseGPUTexture(dev, post_tex);
  if (depth_tex) SDL_ReleaseGPUTexture(dev, depth_tex);

  SDL_GPUTextureCreateInfo depth_info {
    .type   = SDL_GPU_TEXTURETYPE_2D,
    .format = SDL_GPU_TEXTUREFORMAT_D32_FLOAT,
    .usage  = SDL_GPU_TEXTUREUSAGE_DEPTH_STENCIL_TARGET | SDL_GPU_TEXTUREUSAGE_SAMPLER,
    .width  = w,
    .height = h,
    .layer_count_or_depth = 1,
    .num_levels = 1,
  };

  depth_tex = SDL_CreateGPUTexture(dev, &depth_info);
  render_tex = create_render_texture(dev, win, w, h);
  post_tex = create_render_texture(dev, win, w, h);
  render_w = w;
  render_h = h;
}


void Renderer::pass_scene(SDL_GPUCommandBuffer* cmd, const CameraUBO& ubo) {
  SDL_PushGPUVertexUniformData(cmd, 0, &ubo, sizeof(ubo));

  SDL_GPUColorTargetInfo tgt {
    .texture     = render_tex,
    .clear_color = { 0.0, 0.0, 0.0, 1.0},
    .load_op     = SDL_GPU_LOADOP_CLEAR,
    .store_op    = SDL_GPU_STOREOP_STORE,
  };
  
  SDL_GPUDepthStencilTargetInfo depth_tgt {
    .texture          = depth_tex,
    .clear_depth      = 1.0f,
    .load_op          = SDL_GPU_LOADOP_CLEAR,
    .store_op         = SDL_GPU_STOREOP_STORE,
    .stencil_load_op  = SDL_GPU_LOADOP_DONT_CARE,
    .stencil_store_op = SDL_GPU_STOREOP_STORE,
  };

  SDL_GPURenderPass* pass = SDL_BeginGPURenderPass(cmd, &tgt, 1, &depth_tgt);
  SDL_BindGPUGraphicsPipeline(pass, scene.pipeline);
  SDL_GPUBufferBinding binding { .buffer = vert_buff, .offset = 0 };
  SDL_BindGPUVertexBuffers(pass, 0, &binding, 1);
  
  SDL_GPUTextureSamplerBinding tex[1] = {
    { .texture = atlas_tex, .sampler = atlas_sampler },
  };
  SDL_BindGPUFragmentSamplers(pass, 0, tex, 1);
  SDL_DrawGPUPrimitives(pass, vertices.size(), 1, 0, 0);
  SDL_EndGPURenderPass(pass);
}

void Renderer::pass_ray(SDL_GPUCommandBuffer* cmd, const CameraUBO& ubo) {
  SDL_PushGPUFragmentUniformData(cmd, 0, &ubo, sizeof(ubo));

  SDL_GPUColorTargetInfo tgt {
    .texture  = post_tex,
    .load_op  = SDL_GPU_LOADOP_CLEAR,
    .store_op = SDL_GPU_STOREOP_STORE,
  };

  SDL_GPURenderPass* pass = SDL_BeginGPURenderPass(cmd, &tgt, 1, NULL);
  SDL_BindGPUGraphicsPipeline(pass, ray.pipeline);
  SDL_GPUTextureSamplerBinding texs[2] = {
    { .texture = render_tex, .sampler = render_sampler },
    { .texture = depth_tex, .sampler = depth_sampler },
  };
  SDL_BindGPUFragmentSamplers(pass, 0, texs, 2);
  SDL_DrawGPUPrimitives(pass, 3, 1, 0, 0);
  SDL_EndGPURenderPass(pass);
}

void Renderer::pass_post(SDL_GPUCommandBuffer* cmd, SDL_GPUTexture* swapchain, const CameraUBO& ubo) {
  SDL_PushGPUFragmentUniformData(cmd, 0, &ubo, sizeof(ubo));

  SDL_GPUColorTargetInfo tgt {
    .texture  = swapchain,
    .load_op  = SDL_GPU_LOADOP_CLEAR,
    .store_op = SDL_GPU_STOREOP_STORE,
  };

  SDL_GPURenderPass* pass = SDL_BeginGPURenderPass(cmd, &tgt, 1, NULL);
  SDL_BindGPUGraphicsPipeline(pass, post.pipeline);
  SDL_GPUTextureSamplerBinding tex { .texture = post_tex, .sampler = render_sampler };
  SDL_BindGPUFragmentSamplers(pass, 0, &tex, 1);
  SDL_DrawGPUPrimitives(pass, 3, 1, 0, 0);
  SDL_EndGPURenderPass(pass);
}

void Renderer::draw_imgui(SDL_GPUCommandBuffer* cmd, SDL_GPUTexture* tex) {
  ImGui_ImplSDLGPU3_NewFrame();
  ImGui_ImplSDL3_NewFrame();
  ImGui::NewFrame();

  ImGui::Begin("Kale Renderer");
  ImGui::Text("Fps: %f", fps);

  if(ImGui::RadioButton("Standard", (post_mode == 0)))
    post_mode = 0;
  if(ImGui::RadioButton("Black/White", (post_mode == 1)))
    post_mode = 1;
  if(ImGui::RadioButton("Dithered", (post_mode == 2)))
    post_mode = 2;


  ImGui::SliderFloat("Atlas du", &atlas_adjust.x, 0.1f, 100.0f);
  ImGui::SliderFloat("Atlas dv", &atlas_adjust.y, 0.1f, 100.0f);
  
  ImGui::End();

  ImGui::Render();
  ImDrawData* draw_data = ImGui::GetDrawData();
  const bool is_minimized = (draw_data->DisplaySize.x <= 0.0f || draw_data->DisplaySize.y <= 0.0f);

  if(tex && !is_minimized) {
    ImGui_ImplSDLGPU3_PrepareDrawData(draw_data, cmd);

    SDL_GPUColorTargetInfo tgt_info = {
      .texture = tex,
      .mip_level = 0,
      .layer_or_depth_plane = 0,
      .load_op = SDL_GPU_LOADOP_LOAD,
      .store_op = SDL_GPU_STOREOP_STORE,
      .cycle = false
    };

    SDL_GPURenderPass* pass = SDL_BeginGPURenderPass(cmd, &tgt_info, 1, NULL);

    ImGui_ImplSDLGPU3_RenderDrawData(draw_data, cmd, pass);

    SDL_EndGPURenderPass(pass);
  }
}

void Renderer::draw(CameraUBO& ubo) {
  ubo.post_mode = post_mode;
  ubo.atlas_adjust = atlas_adjust;
  static std::array<float, 25> frame_times;

  static float last_time = (float)SDL_GetPerformanceCounter()/(float)SDL_GetPerformanceFrequency();

  float now = (float)SDL_GetPerformanceCounter()/(float)SDL_GetPerformanceFrequency();
  frame_times[frame % 25] = 1.0f/(25.0f*(now-last_time));
  last_time = now;

  if(frame > 25) {
    fps = 0.0f;
    for(const auto& t : frame_times) {
      fps += t;
    }
  }

  SDL_GPUCommandBuffer* cmd = SDL_AcquireGPUCommandBuffer(dev);

  SDL_GPUTexture* swapchain;
  uint32_t sw_w, sw_h;
  SDL_WaitAndAcquireGPUSwapchainTexture(cmd, win, &swapchain, &sw_w, &sw_h);

  if (!swapchain) {
    SDL_SubmitGPUCommandBuffer(cmd);
    return;
  }

  if (sw_w != render_w || sw_h != render_h)
    resize_render_texture(sw_w, sw_h);

  pass_scene(cmd, ubo);
  pass_ray(cmd, ubo);
  pass_post(cmd, swapchain, ubo);
  draw_imgui(cmd, swapchain);

  SDL_SubmitGPUCommandBuffer(cmd);

  frame++;
}
