// This file is part of the 64k demo project.
// It implements the Sequence management logic.

#include "effect.h"
#include "gpu/demo_effects.h"
#include "gpu/gpu.h"
#include <algorithm>
#include <cstdio>
#include <vector>

// --- PostProcessEffect ---
void PostProcessEffect::render(WGPURenderPassEncoder pass, float, float, float,
                               float) {
  if (pipeline_ && bind_group_) {
    wgpuRenderPassEncoderSetPipeline(pass, pipeline_);
    wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr);
    wgpuRenderPassEncoderDraw(pass, 3, 1, 0, 0); // Fullscreen triangle
  }
}

// --- Sequence Implementation ---
void Sequence::resize(int width, int height) {
  for (SequenceItem& item : items_) {
    item.effect->resize(width, height);
  }
}

void Sequence::init(MainSequence* demo) {
  for (SequenceItem& item : items_) {
    if (!item.effect->is_initialized) {
      item.effect->init(demo);
      item.effect->is_initialized = true;
    }
  }
}

void Sequence::add_effect(std::shared_ptr<Effect> effect, float start_time,
                          float end_time, int priority) {
  items_.push_back({effect, start_time, end_time, priority, false});
  is_sorted_ = false;
}

void Sequence::sort_items() {
  if (is_sorted_)
    return;
  std::sort(items_.begin(), items_.end(),
            [](const SequenceItem& a, const SequenceItem& b) {
              return a.priority < b.priority;
            });
  is_sorted_ = true;
}

void Sequence::update_active_list(float seq_time) {
  for (SequenceItem& item : items_) {
    bool should_be_active =
        (seq_time >= item.start_time && seq_time < item.end_time);

    if (should_be_active && !item.active) {
      item.effect->start();
      item.active = true;
    } else if (!should_be_active && item.active) {
      item.effect->end();
      item.active = false;
    }
  }
}

void Sequence::collect_active_effects(
    std::vector<SequenceItem*>& scene_effects,
    std::vector<SequenceItem*>& post_effects) {
  sort_items();
  for (SequenceItem& item : items_) {
    if (item.active) {
      if (item.effect->is_post_process()) {
        post_effects.push_back(&item);
      } else {
        scene_effects.push_back(&item);
      }
    }
  }
}

void Sequence::reset() {
  for (SequenceItem& item : items_) {
    if (item.active) {
      item.effect->end();
      item.active = false;
    }
  }
}

// --- MainSequence Implementation ---

MainSequence::MainSequence() = default;
MainSequence::~MainSequence() = default;

void MainSequence::create_framebuffers(int width, int height) {
  // In test mode, this would be skipped or mocked.
  // For now, it will only be called by the real init.
  WGPUTextureDescriptor desc = {};
  desc.usage =
      WGPUTextureUsage_RenderAttachment | WGPUTextureUsage_TextureBinding;
  desc.dimension = WGPUTextureDimension_2D;
  desc.size = {(uint32_t)width, (uint32_t)height, 1};
  desc.format = format;
  desc.mipLevelCount = 1;
  desc.sampleCount = 1;

  framebuffer_a_ = wgpuDeviceCreateTexture(device, &desc);
  framebuffer_b_ = wgpuDeviceCreateTexture(device, &desc);

  WGPUTextureViewDescriptor view_desc = {};
  view_desc.dimension = WGPUTextureViewDimension_2D;
  view_desc.format = format;
  view_desc.mipLevelCount = 1;
  view_desc.arrayLayerCount = 1;

  framebuffer_view_a_ = wgpuTextureCreateView(framebuffer_a_, &view_desc);
  framebuffer_view_b_ = wgpuTextureCreateView(framebuffer_b_, &view_desc);

  // Depth Buffer
  WGPUTextureDescriptor depth_desc = {};
  depth_desc.usage = WGPUTextureUsage_RenderAttachment;
  depth_desc.dimension = WGPUTextureDimension_2D;
  depth_desc.size = {(uint32_t)width, (uint32_t)height, 1};
  depth_desc.format = WGPUTextureFormat_Depth24Plus;
  depth_desc.mipLevelCount = 1;
  depth_desc.sampleCount = 1;
  depth_texture_ = wgpuDeviceCreateTexture(device, &depth_desc);

  WGPUTextureViewDescriptor depth_view_desc = {};
  depth_view_desc.format = WGPUTextureFormat_Depth24Plus;
  depth_view_desc.dimension = WGPUTextureViewDimension_2D;
  depth_view_desc.aspect = WGPUTextureAspect_DepthOnly;
  depth_view_desc.mipLevelCount = 1;
  depth_view_desc.arrayLayerCount = 1;
  depth_view_ = wgpuTextureCreateView(depth_texture_, &depth_view_desc);
}

void MainSequence::init_test(WGPUDevice d, WGPUQueue q, WGPUTextureFormat f) {
  device = d;
  queue = q;
  format = f;
  // No framebuffers or passthrough effect created in test mode.
  // Test effects should not rely on these being real.
}

void MainSequence::init(WGPUDevice d, WGPUQueue q, WGPUTextureFormat f,
                        int width, int height) {
  device = d;
  queue = q;
  format = f;
  width_ = width;
  height_ = height;

  create_framebuffers(width, height);
  passthrough_effect_ =
      std::make_unique<PassthroughEffect>(device, queue, format);
  passthrough_effect_->resize(width, height);

  for (ActiveSequence& entry : sequences_) {
    entry.seq->init(this);
    entry.seq->resize(width, height);
  }
}

void MainSequence::add_sequence(std::shared_ptr<Sequence> seq, float start_time,
                                int priority) {
  sequences_.push_back({seq, start_time, priority});
  // If MainSequence is already initialized, init the new sequence immediately
  if (device) {
    seq->init(this);
    seq->resize(width_, height_);
  }
  std::sort(sequences_.begin(), sequences_.end(),
            [](const ActiveSequence& a, const ActiveSequence& b) {
              return a.priority < b.priority;
            });
}

void MainSequence::resize(int width, int height) {
  width_ = width;
  height_ = height;
  // Release old resources
  if (framebuffer_view_a_)
    wgpuTextureViewRelease(framebuffer_view_a_);
  if (framebuffer_a_)
    wgpuTextureRelease(framebuffer_a_);
  if (framebuffer_view_b_)
    wgpuTextureViewRelease(framebuffer_view_b_);
  if (framebuffer_b_)
    wgpuTextureRelease(framebuffer_b_);
  if (depth_view_)
    wgpuTextureViewRelease(depth_view_);
  if (depth_texture_)
    wgpuTextureRelease(depth_texture_);

  // Recreate with new size
  create_framebuffers(width, height);

  if (passthrough_effect_) {
    passthrough_effect_->resize(width, height);
  }

  // Propagate to all sequences
  for (ActiveSequence& entry : sequences_) {
    entry.seq->resize(width, height);
  }
}

void MainSequence::render_frame(float global_time, float beat, float peak,
                                float aspect_ratio, WGPUSurface surface) {
  WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device, nullptr);

  std::vector<SequenceItem*> scene_effects;
  std::vector<SequenceItem*> post_effects;
  for (ActiveSequence& entry : sequences_) {
    if (global_time >= entry.start_time) {
      float seq_time = global_time - entry.start_time;
      entry.seq->update_active_list(seq_time);
      entry.seq->collect_active_effects(scene_effects, post_effects);
    }
  }
  std::sort(scene_effects.begin(), scene_effects.end(),
            [](const SequenceItem* a, const SequenceItem* b) {
              return a->priority < b->priority;
            });
  std::sort(post_effects.begin(), post_effects.end(),
            [](const SequenceItem* a, const SequenceItem* b) {
              return a->priority < b->priority;
            });

  // 1. Compute
  for (const SequenceItem* item : scene_effects) {
    item->effect->compute(encoder, global_time - item->start_time, beat, peak,
                          aspect_ratio);
  }

  // 2. Scene Pass (to A)
  WGPURenderPassColorAttachment scene_attachment = {};
  scene_attachment.view = framebuffer_view_a_;
  scene_attachment.resolveTarget = nullptr;
  scene_attachment.loadOp = WGPULoadOp_Clear;
  scene_attachment.storeOp = WGPUStoreOp_Store;
  scene_attachment.clearValue = {0, 0, 0, 1};
#if !defined(DEMO_CROSS_COMPILE_WIN32)
  scene_attachment.depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;
#endif /* !defined(DEMO_CROSS_COMPILE_WIN32) */

  WGPURenderPassDepthStencilAttachment depth_attachment = {};
  depth_attachment.view = depth_view_;
  depth_attachment.depthLoadOp = WGPULoadOp_Clear;
  depth_attachment.depthStoreOp = WGPUStoreOp_Store;
  depth_attachment.depthClearValue = 1.0f;

  WGPURenderPassDescriptor scene_desc = {.colorAttachmentCount = 1,
                                         .colorAttachments = &scene_attachment,
                                         .depthStencilAttachment =
                                             &depth_attachment};
  WGPURenderPassEncoder scene_pass =
      wgpuCommandEncoderBeginRenderPass(encoder, &scene_desc);
  wgpuRenderPassEncoderSetViewport(scene_pass, 0.0f, 0.0f, (float)width_,
                                   (float)height_, 0.0f, 1.0f);
  for (const SequenceItem* item : scene_effects) {
    item->effect->render(scene_pass, global_time - item->start_time, beat, peak,
                         aspect_ratio);
  }
  wgpuRenderPassEncoderEnd(scene_pass);

  // 3. Post Chain
  WGPUSurfaceTexture st = {};
  WGPUTextureView final_view = nullptr;

  if (post_effects.empty()) {
    wgpuSurfaceGetCurrentTexture(surface, &st);
    final_view = wgpuTextureCreateView(st.texture, nullptr);

    // Safely cast to PostProcessEffect to call update_bind_group
    PostProcessEffect* pp_effect =
        (PostProcessEffect*)passthrough_effect_.get();
    pp_effect->update_bind_group(framebuffer_view_a_);

    WGPURenderPassColorAttachment final_attachment = {};
    final_attachment.view = final_view;
    final_attachment.resolveTarget = nullptr;
    final_attachment.loadOp = WGPULoadOp_Load;
    final_attachment.storeOp = WGPUStoreOp_Store;
#if !defined(DEMO_CROSS_COMPILE_WIN32)
    final_attachment.depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;
#endif /* !defined(DEMO_CROSS_COMPILE_WIN32) */
    WGPURenderPassDescriptor final_desc = {
        .colorAttachmentCount = 1, .colorAttachments = &final_attachment};
    WGPURenderPassEncoder final_pass =
        wgpuCommandEncoderBeginRenderPass(encoder, &final_desc);
    wgpuRenderPassEncoderSetViewport(final_pass, 0.0f, 0.0f, (float)width_,
                                     (float)height_, 0.0f, 1.0f);
    passthrough_effect_->render(final_pass, 0, 0, 0, aspect_ratio);
    wgpuRenderPassEncoderEnd(final_pass);
  } else {
    WGPUTextureView current_input = framebuffer_view_a_;
    for (size_t i = 0; i < post_effects.size(); ++i) {
      bool is_last = (i == post_effects.size() - 1);
      WGPUTextureView current_output = nullptr;

      if (is_last) {
        wgpuSurfaceGetCurrentTexture(surface, &st);
        final_view = wgpuTextureCreateView(st.texture, nullptr);
        current_output = final_view;
      } else {
        current_output =
            (current_input == framebuffer_view_a_ ? framebuffer_view_b_
                                                  : framebuffer_view_a_);
      }

      PostProcessEffect* pp =
          (PostProcessEffect*)(post_effects[i]->effect.get());
      pp->update_bind_group(current_input);

      WGPURenderPassColorAttachment pp_attachment = {};
      pp_attachment.view = current_output;
      pp_attachment.resolveTarget = nullptr;
      pp_attachment.loadOp = WGPULoadOp_Load;
      pp_attachment.storeOp = WGPUStoreOp_Store;
#if !defined(DEMO_CROSS_COMPILE_WIN32)
      pp_attachment.depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;
#endif /* !defined(DEMO_CROSS_COMPILE_WIN32) */
      WGPURenderPassDescriptor pp_desc = {.colorAttachmentCount = 1,
                                          .colorAttachments = &pp_attachment};
      WGPURenderPassEncoder pp_pass =
          wgpuCommandEncoderBeginRenderPass(encoder, &pp_desc);
      wgpuRenderPassEncoderSetViewport(pp_pass, 0.0f, 0.0f, (float)width_,
                                       (float)height_, 0.0f, 1.0f);
      pp->render(pp_pass, global_time - post_effects[i]->start_time, beat, peak,
                 aspect_ratio);
      wgpuRenderPassEncoderEnd(pp_pass);
      current_input = current_output;
    }
  }

  WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, nullptr);
  wgpuQueueSubmit(queue, 1, &commands);

  if (st.texture) {
    wgpuTextureViewRelease(final_view);
    wgpuSurfacePresent(surface);
    wgpuTextureRelease(st.texture);
  }
}

void MainSequence::shutdown() {
  if (framebuffer_view_a_)
    wgpuTextureViewRelease(framebuffer_view_a_);
  if (framebuffer_a_)
    wgpuTextureRelease(framebuffer_a_);
  if (framebuffer_view_b_)
    wgpuTextureViewRelease(framebuffer_view_b_);
  if (framebuffer_b_)
    wgpuTextureRelease(framebuffer_b_);
  if (depth_view_)
    wgpuTextureViewRelease(depth_view_);
  if (depth_texture_)
    wgpuTextureRelease(depth_texture_);
  for (ActiveSequence& entry : sequences_) {
    entry.seq->reset();
  }
}

#if !defined(STRIP_ALL)
void MainSequence::simulate_until(float target_time, float step_rate) {
  const float bpm = 128.0f;
  const float aspect_ratio = 16.0f / 9.0f;
  for (float t = 0.0f; t < target_time; t += step_rate) {
    WGPUCommandEncoder encoder =
        wgpuDeviceCreateCommandEncoder(device, nullptr);
    float beat = fmodf(t * bpm / 60.0f, 1.0f);
    std::vector<SequenceItem*> scene_effects, post_effects;
    for (ActiveSequence& entry : sequences_) {
      if (t >= entry.start_time) {
        entry.seq->update_active_list(t - entry.start_time);
        entry.seq->collect_active_effects(scene_effects, post_effects);
      }
    }
    for (const SequenceItem* item : scene_effects) {
      item->effect->compute(encoder, t - item->start_time, beat, 0.0f,
                            aspect_ratio);
    }
    WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, nullptr);
    wgpuQueueSubmit(queue, 1, &commands);
  }
}
#endif /* !defined(STRIP_ALL) */