#include "pipelines.h"
#include "gpu.h"
#include "types.h"
#include <iostream>

void ScenePipeline::create(SDL_GPUDevice* dev, SDL_Window* win) {
  vert = load_shader(dev, "shaders/scene.vert.spv", SDL_GPU_SHADERSTAGE_VERTEX,   0, 1);
  frag = load_shader(dev, "shaders/scene.frag.spv", SDL_GPU_SHADERSTAGE_FRAGMENT, 0, 0);

  SDL_GPUVertexBufferDescription vert_buf {
    .slot               = 0,
    .pitch              = sizeof(Vertex),
    .input_rate         = SDL_GPU_VERTEXINPUTRATE_VERTEX,
    .instance_step_rate = 0,
  };

  SDL_GPUVertexAttribute attrs[3] = {
    { .location = 0, .buffer_slot = 0, .format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3, .offset = 0 },
    { .location = 1, .buffer_slot = 0, .format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4, .offset = (uint32_t)offsetof(Vertex, color) },
    { .location = 2, .buffer_slot = 0, .format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3, .offset = (uint32_t)offsetof(Vertex, normal) },
  };

  SDL_GPUColorTargetDescription color_tgt {
    .format = SDL_GetGPUSwapchainTextureFormat(dev, win),
    .blend_state = {
      .src_color_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
      .dst_color_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
      .color_blend_op        = SDL_GPU_BLENDOP_ADD,
      .src_alpha_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
      .dst_alpha_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
      .alpha_blend_op        = SDL_GPU_BLENDOP_ADD,
      .enable_blend          = false,
    }
  };

  SDL_GPUGraphicsPipelineCreateInfo info {
    .vertex_shader   = vert,
    .fragment_shader = frag,
    .vertex_input_state = {
      .vertex_buffer_descriptions = &vert_buf,
      .num_vertex_buffers         = 1,
      .vertex_attributes          = attrs,
      .num_vertex_attributes      = sizeof(attrs)/sizeof(attrs[0]),
    },
    .primitive_type = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,
    .rasterizer_state = {
      .fill_mode  = SDL_GPU_FILLMODE_FILL,
      .cull_mode  = SDL_GPU_CULLMODE_NONE,
      .front_face = SDL_GPU_FRONTFACE_COUNTER_CLOCKWISE,
    },
    .depth_stencil_state = {
      .compare_op = SDL_GPU_COMPAREOP_LESS,
      .enable_depth_test = true,
      .enable_depth_write = true,
    },
    .target_info = {
      .color_target_descriptions = &color_tgt,
      .num_color_targets         = 1,
      .depth_stencil_format      = SDL_GPU_TEXTUREFORMAT_D32_FLOAT,
      .has_depth_stencil_target  = true,
    },
  };

  pipeline = SDL_CreateGPUGraphicsPipeline(dev, &info);
  if (!pipeline)
    std::cerr << "Scene pipeline failed: " << SDL_GetError() << std::endl;
}

void ScenePipeline::destroy(SDL_GPUDevice* dev) {
  if (pipeline) SDL_ReleaseGPUGraphicsPipeline(dev, pipeline);
  if (vert)     SDL_ReleaseGPUShader(dev, vert);
  if (frag)     SDL_ReleaseGPUShader(dev, frag);
}

void BlitPipeline::create(SDL_GPUDevice* dev, SDL_Window* win) {
  vert = load_shader(dev, "shaders/blit.vert.spv", SDL_GPU_SHADERSTAGE_VERTEX,   0, 0);
  frag = load_shader(dev, "shaders/blit.frag.spv", SDL_GPU_SHADERSTAGE_FRAGMENT, 2, 1);

  SDL_GPUColorTargetDescription color_tgt {
    .format = SDL_GetGPUSwapchainTextureFormat(dev, win),
  };

  SDL_GPUGraphicsPipelineCreateInfo info {
    .vertex_shader   = vert,
    .fragment_shader = frag,
    .vertex_input_state = {
      .vertex_buffer_descriptions = nullptr,
      .num_vertex_buffers         = 0,
      .vertex_attributes          = nullptr,
      .num_vertex_attributes      = 0,
    },
    .primitive_type = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,
    .rasterizer_state = {
      .fill_mode  = SDL_GPU_FILLMODE_FILL,
      .cull_mode  = SDL_GPU_CULLMODE_NONE,
      .front_face = SDL_GPU_FRONTFACE_COUNTER_CLOCKWISE,
    },
    .target_info = {
      .color_target_descriptions = &color_tgt,
      .num_color_targets         = 1,
    },
  };

  pipeline = SDL_CreateGPUGraphicsPipeline(dev, &info);
  if (!pipeline)
    std::cerr << "Blit pipeline failed: " << SDL_GetError() << std::endl;
}

void BlitPipeline::destroy(SDL_GPUDevice* dev) {
  if (pipeline) SDL_ReleaseGPUGraphicsPipeline(dev, pipeline);
  if (vert)     SDL_ReleaseGPUShader(dev, vert);
  if (frag)     SDL_ReleaseGPUShader(dev, frag);
}

void PostPipeline::create(SDL_GPUDevice* dev, SDL_Window* win) {
  vert = load_shader(dev, "shaders/blit.vert.spv",  SDL_GPU_SHADERSTAGE_VERTEX,   0, 0);
  frag = load_shader(dev, "shaders/post.frag.spv",  SDL_GPU_SHADERSTAGE_FRAGMENT, 1, 1);

  SDL_GPUColorTargetDescription color_tgt {
    .format = SDL_GetGPUSwapchainTextureFormat(dev, win),
  };

  SDL_GPUGraphicsPipelineCreateInfo info {
    .vertex_shader   = vert,
    .fragment_shader = frag,
    .vertex_input_state = {
      .vertex_buffer_descriptions = nullptr,
      .num_vertex_buffers         = 0,
      .vertex_attributes          = nullptr,
      .num_vertex_attributes      = 0,
    },
    .primitive_type = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,
    .rasterizer_state = {
      .fill_mode  = SDL_GPU_FILLMODE_FILL,
      .cull_mode  = SDL_GPU_CULLMODE_NONE,
      .front_face = SDL_GPU_FRONTFACE_COUNTER_CLOCKWISE,
    },
    .target_info = {
      .color_target_descriptions = &color_tgt,
      .num_color_targets         = 1,
    },
  };

  pipeline = SDL_CreateGPUGraphicsPipeline(dev, &info);
  if (!pipeline)
    std::cerr << "Post pipeline failed: " << SDL_GetError() << std::endl;
}

void PostPipeline::destroy(SDL_GPUDevice* dev) {
  if (pipeline) SDL_ReleaseGPUGraphicsPipeline(dev, pipeline);
  if (vert)     SDL_ReleaseGPUShader(dev, vert);
  if (frag)     SDL_ReleaseGPUShader(dev, frag);
}