// This file is part of the 64k demo project.
// It implements the Renderer3D class.

#include "3d/renderer.h"
#include <algorithm>
#include <cassert>
#include <cstring>
#include <iostream>

#if !defined(STRIP_ALL)
bool Renderer3D::s_debug_enabled_ = false;
#endif

static const char* kShaderCode = R"(
struct GlobalUniforms {
    view_proj: mat4x4<f32>,
    camera_pos_time: vec4<f32>,
    params: vec4<f32>,
};

struct ObjectData {
    model: mat4x4<f32>,
    model_inv_tr: mat4x4<f32>,
    color: vec4<f32>,
    params: vec4<f32>,
};

struct ObjectsBuffer {
    objects: array<ObjectData>,
};

@group(0) @binding(0) var<uniform> globals: GlobalUniforms;
@group(0) @binding(1) var<storage, read> object_data: ObjectsBuffer;
@group(0) @binding(2) var noise_tex: texture_2d<f32>;
@group(0) @binding(3) var noise_sampler: sampler;

struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) local_pos: vec3<f32>,
    @location(1) color: vec4<f32>,
    @location(2) @interpolate(flat) instance_index: u32,
    @location(3) world_pos: vec3<f32>,
};

@vertex
fn vs_main(@builtin(vertex_index) vertex_index: u32, 
           @builtin(instance_index) instance_index: u32) -> VertexOutput {
    
    var pos = array<vec3<f32>, 36>(
        vec3(-1.0, -1.0,  1.0), vec3( 1.0, -1.0,  1.0), vec3( 1.0,  1.0,  1.0),
        vec3(-1.0, -1.0,  1.0), vec3( 1.0,  1.0,  1.0), vec3(-1.0,  1.0,  1.0),
        vec3(-1.0, -1.0, -1.0), vec3(-1.0,  1.0, -1.0), vec3( 1.0,  1.0, -1.0),
        vec3(-1.0, -1.0, -1.0), vec3( 1.0,  1.0, -1.0), vec3( 1.0, -1.0, -1.0),
        vec3(-1.0,  1.0, -1.0), vec3(-1.0,  1.0,  1.0), vec3( 1.0,  1.0,  1.0),
        vec3(-1.0,  1.0, -1.0), vec3( 1.0,  1.0,  1.0), vec3( 1.0,  1.0, -1.0),
        vec3(-1.0, -1.0, -1.0), vec3( 1.0, -1.0, -1.0), vec3( 1.0, -1.0,  1.0),
        vec3(-1.0, -1.0, -1.0), vec3( 1.0, -1.0,  1.0), vec3(-1.0, -1.0,  1.0),
        vec3( 1.0, -1.0, -1.0), vec3( 1.0,  1.0, -1.0), vec3( 1.0,  1.0,  1.0),
        vec3( 1.0, -1.0, -1.0), vec3( 1.0,  1.0,  1.0), vec3( 1.0, -1.0,  1.0),
        vec3(-1.0, -1.0, -1.0), vec3(-1.0, -1.0,  1.0), vec3(-1.0,  1.0,  1.0),
        vec3(-1.0, -1.0, -1.0), vec3(-1.0,  1.0,  1.0), vec3(-1.0,  1.0, -1.0)
    );

    let p = pos[vertex_index];
    let obj = object_data.objects[instance_index];
    let world_pos = obj.model * vec4<f32>(p, 1.0);
    let clip_pos = globals.view_proj * world_pos;

    var out: VertexOutput;
    out.position = clip_pos;
    out.local_pos = p; 
    out.color = obj.color;
    out.instance_index = instance_index;
    out.world_pos = world_pos.xyz;
    return out;
}

fn sdSphere(p: vec3<f32>, r: f32) -> f32 {
    return length(p) - r;
}

fn sdBox(p: vec3<f32>, b: vec3<f32>) -> f32 {
    let q = abs(p) - b;
    return length(max(q, vec3<f32>(0.0))) + min(max(q.x, max(q.y, q.z)), 0.0);
}

fn sdTorus(p: vec3<f32>, t: vec2<f32>) -> f32 {
    let q = vec2<f32>(length(p.xz) - t.x, p.y);
    return length(q) - t.y;
}

fn get_dist(p: vec3<f32>, obj_type: f32) -> f32 {
    if (obj_type == 1.0) { return sdSphere(p, 0.9); }
    if (obj_type == 2.0) { return sdBox(p, vec3<f32>(0.7)); }
    if (obj_type == 3.0) { return sdTorus(p, vec2<f32>(0.6, 0.25)); }
    return 100.0;
}

fn map_scene(p: vec3<f32>) -> f32 {
    var d = 1000.0;
    let count = u32(globals.params.x);
    
    // Brute force loop over all objects
    for (var i = 0u; i < count; i = i + 1u) {
        let obj = object_data.objects[i];
        let obj_type = obj.params.x;
        if (obj_type <= 0.0) { continue; } // Skip non-sdf objects

        // Transform world p to local q
        // Assuming uniform scale
        let center = vec3<f32>(obj.model[3].x, obj.model[3].y, obj.model[3].z);
        let scale = length(vec3<f32>(obj.model[0].x, obj.model[0].y, obj.model[0].z));
        let mat3 = mat3x3<f32>(obj.model[0].xyz/scale, obj.model[1].xyz/scale, obj.model[2].xyz/scale);
        let q = transpose(mat3) * (p - center) / scale;
        
        let dist = get_dist(q, obj_type) * scale;
        d = min(d, dist);
    }
    return d;
}

fn calc_shadow(ro: vec3<f32>, rd: vec3<f32>, tmin: f32, tmax: f32) -> f32 {
    var t = tmin;
    var res = 1.0;
    for (var i = 0; i < 30; i = i + 1) {
        let h = map_scene(ro + rd * t);
        if (h < 0.001) {
            return 0.0; // Hard shadow hit
        }
        res = min(res, 8.0 * h / t); // Soft shadow k=8
        t = t + h;
        if (t > tmax) { break; }
    }
    return res;
}

fn get_normal(p: vec3<f32>, obj_type: f32) -> vec3<f32> {
    if (obj_type == 1.0) { return normalize(p); }
    let e = vec2<f32>(0.001, 0.0);
    return normalize(vec3<f32>(
        get_dist(p + e.xyy, obj_type) - get_dist(p - e.xyy, obj_type),
        get_dist(p + e.yxy, obj_type) - get_dist(p - e.yxy, obj_type),
        get_dist(p + e.yyx, obj_type) - get_dist(p - e.yyx, obj_type)
    ));
}

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    let obj = object_data.objects[in.instance_index];
    let obj_type = obj.params.x;

    var p: vec3<f32>;
    var normal: vec3<f32>;
    var base_color = in.color.rgb;
    let light_dir = normalize(vec3<f32>(0.2, 1.0, 0.2)); // More vertical light for easier shadow debugging

    if (obj_type == 0.0) { // Rasterized object
        p = in.world_pos;
        let local_normal = normalize(cross(dpdx(in.local_pos), dpdy(in.local_pos)));
        
        // Correct normal transformation using inverse transpose
        let mat3_it = mat3x3<f32>(obj.model_inv_tr[0].xyz, obj.model_inv_tr[1].xyz, obj.model_inv_tr[2].xyz);
        normal = normalize(mat3_it * local_normal);

        // Apply grid texture for color to floor
        let uv = p.xz * 0.1;
        let grid_val = textureSample(noise_tex, noise_sampler, uv).r;
        base_color = base_color * (0.3 + 0.7 * grid_val);
    } else { // SDF object
        let center = vec3<f32>(obj.model[3].x, obj.model[3].y, obj.model[3].z);
        let scale = length(vec3<f32>(obj.model[0].x, obj.model[0].y, obj.model[0].z));
        let ro = globals.camera_pos_time.xyz;
        let rd = normalize(in.world_pos - ro);
        var t = length(in.world_pos - ro); 
        p = ro + rd * t;
        
        let mat3 = mat3x3<f32>(obj.model[0].xyz/scale, obj.model[1].xyz/scale, obj.model[2].xyz/scale);
        
        var hit = false;
        for (var i = 0; i < 40; i = i + 1) {
            let q = transpose(mat3) * (p - center) / scale;
            let d_world = get_dist(q, obj_type) * scale;
            if (d_world < 0.001) { hit = true; break; }
            if (d_world > 3.0 * scale) { break; }
            p = p + rd * d_world;
        }
        
        if (!hit) { discard; }
        
        let q_hit = transpose(mat3) * (p - center) / scale;
        
        let e = vec2<f32>(0.005, 0.0);
        let disp_strength = 0.05;
        
        // Calculate normal with bump mapping
        let q_x1 = q_hit + e.xyy;
        let uv_x1 = vec2<f32>(atan2(q_x1.x, q_x1.z) / 6.28 + 0.5, acos(clamp(q_x1.y / length(q_x1), -1.0, 1.0)) / 3.14);
        let h_x1 = textureSample(noise_tex, noise_sampler, uv_x1).r;
        let d_x1 = get_dist(q_x1, obj_type) - disp_strength * h_x1;
        
        let q_x2 = q_hit - e.xyy;
        let uv_x2 = vec2<f32>(atan2(q_x2.x, q_x2.z) / 6.28 + 0.5, acos(clamp(q_x2.y / length(q_x2), -1.0, 1.0)) / 3.14);
        let h_x2 = textureSample(noise_tex, noise_sampler, uv_x2).r;
        let d_x2 = get_dist(q_x2, obj_type) - disp_strength * h_x2;
        
        let q_y1 = q_hit + e.yxy;
        let uv_y1 = vec2<f32>(atan2(q_y1.x, q_y1.z) / 6.28 + 0.5, acos(clamp(q_y1.y / length(q_y1), -1.0, 1.0)) / 3.14);
        let h_y1 = textureSample(noise_tex, noise_sampler, uv_y1).r;
        let d_y1 = get_dist(q_y1, obj_type) - disp_strength * h_y1;
        
        let q_y2 = q_hit - e.yxy;
        let uv_y2 = vec2<f32>(atan2(q_y2.x, q_y2.z) / 6.28 + 0.5, acos(clamp(q_y2.y / length(q_y2), -1.0, 1.0)) / 3.14);
        let h_y2 = textureSample(noise_tex, noise_sampler, uv_y2).r;
        let d_y2 = get_dist(q_y2, obj_type) - disp_strength * h_y2;
        
        let q_z1 = q_hit + e.yyx;
        let uv_z1 = vec2<f32>(atan2(q_z1.x, q_z1.z) / 6.28 + 0.5, acos(clamp(q_z1.y / length(q_z1), -1.0, 1.0)) / 3.14);
        let h_z1 = textureSample(noise_tex, noise_sampler, uv_z1).r;
        let d_z1 = get_dist(q_z1, obj_type) - disp_strength * h_z1;
        
        let q_z2 = q_hit - e.yyx;
        let uv_z2 = vec2<f32>(atan2(q_z2.x, q_z2.z) / 6.28 + 0.5, acos(clamp(q_z2.y / length(q_z2), -1.0, 1.0)) / 3.14);
        let h_z2 = textureSample(noise_tex, noise_sampler, uv_z2).r;
        let d_z2 = get_dist(q_z2, obj_type) - disp_strength * h_z2;
        
        let n_local = normalize(vec3<f32>(d_x1 - d_x2, d_y1 - d_y2, d_z1 - d_z2));
        
        normal = normalize(mat3 * n_local);
    }

    let shadow = calc_shadow(p + normal * 0.01, light_dir, 0.0, 20.0);
    let lighting = (max(dot(normal, light_dir), 0.0) * shadow) + 0.1;
    return vec4<f32>(base_color * lighting, 1.0);
}

)";

void Renderer3D::init(WGPUDevice device, WGPUQueue queue,
                      WGPUTextureFormat format) {
  device_ = device;
  queue_ = queue;
  format_ = format;

  WGPUSamplerDescriptor sampler_desc = {};
  sampler_desc.addressModeU = WGPUAddressMode_Repeat;
  sampler_desc.addressModeV = WGPUAddressMode_Repeat;
  sampler_desc.magFilter = WGPUFilterMode_Linear;
  sampler_desc.minFilter = WGPUFilterMode_Linear;
  sampler_desc.maxAnisotropy = 1;
  default_sampler_ = wgpuDeviceCreateSampler(device_, &sampler_desc);

  create_default_resources();
  create_pipeline();

#if !defined(STRIP_ALL)
  visual_debug_.init(device_, format_);
#endif
}

void Renderer3D::shutdown() {
#if !defined(STRIP_ALL)
  visual_debug_.shutdown();
#endif

  if (default_sampler_)
    wgpuSamplerRelease(default_sampler_);
  if (pipeline_)
    wgpuRenderPipelineRelease(pipeline_);
  if (bind_group_)
    wgpuBindGroupRelease(bind_group_);
  if (global_uniform_buffer_)
    wgpuBufferRelease(global_uniform_buffer_);
  if (object_storage_buffer_)
    wgpuBufferRelease(object_storage_buffer_);
  if (depth_view_)
    wgpuTextureViewRelease(depth_view_);
  if (depth_texture_)
    wgpuTextureRelease(depth_texture_);
}

void Renderer3D::resize(int width, int height) {
  if (width == width_ && height == height_)
    return;
  width_ = width;
  height_ = height;

  if (depth_view_)
    wgpuTextureViewRelease(depth_view_);
  if (depth_texture_)
    wgpuTextureRelease(depth_texture_);

  WGPUTextureDescriptor desc = {};
  desc.usage = WGPUTextureUsage_RenderAttachment;
  desc.dimension = WGPUTextureDimension_2D;
  desc.size = {(uint32_t)width, (uint32_t)height, 1};
  desc.format = WGPUTextureFormat_Depth24Plus;
  desc.mipLevelCount = 1;
  desc.sampleCount = 1;
  depth_texture_ = wgpuDeviceCreateTexture(device_, &desc);

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

void Renderer3D::create_default_resources() {
  global_uniform_buffer_ =
      gpu_create_buffer(device_, sizeof(GlobalUniforms),
                        WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst,
                        nullptr)
          .buffer;
  object_storage_buffer_ =
      gpu_create_buffer(device_, sizeof(ObjectData) * kMaxObjects,
                        WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst,
                        nullptr)
          .buffer;
}

void Renderer3D::set_noise_texture(WGPUTextureView noise_view) {
  noise_texture_view_ = noise_view;
  // Note: Bind group needs recreation if texture changes, but we'll do it in
  // render for simplicity or just once at init if it's static. For this demo,
  // let's recreate in render if changed.
}

void Renderer3D::create_pipeline() {
  WGPUBindGroupLayoutEntry entries[4] = {};
  entries[0].binding = 0;
  entries[0].visibility = WGPUShaderStage_Vertex | WGPUShaderStage_Fragment;
  entries[0].buffer.type = WGPUBufferBindingType_Uniform;
  entries[0].buffer.minBindingSize = sizeof(GlobalUniforms);

  entries[1].binding = 1;
  entries[1].visibility = WGPUShaderStage_Vertex | WGPUShaderStage_Fragment;
  entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
  entries[1].buffer.minBindingSize = sizeof(ObjectData) * kMaxObjects;

  entries[2].binding = 2;
  entries[2].visibility = WGPUShaderStage_Fragment;
  entries[2].texture.sampleType = WGPUTextureSampleType_Float;
  entries[2].texture.viewDimension = WGPUTextureViewDimension_2D;

  entries[3].binding = 3;
  entries[3].visibility = WGPUShaderStage_Fragment;
  entries[3].sampler.type = WGPUSamplerBindingType_Filtering;

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = 4;
  bgl_desc.entries = entries;
  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device_, &bgl_desc);

  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bgl;
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(device_, &pl_desc);

#if defined(DEMO_CROSS_COMPILE_WIN32)
  WGPUShaderModuleWGSLDescriptor wgsl_desc = {};
  wgsl_desc.chain.sType = WGPUSType_ShaderModuleWGSLDescriptor;
  wgsl_desc.code = kShaderCode;
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = (const WGPUChainedStruct*)&wgsl_desc.chain;
#else
  WGPUShaderSourceWGSL wgsl_desc = {};
  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_desc.code = {kShaderCode, strlen(kShaderCode)};
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = (const WGPUChainedStruct*)&wgsl_desc.chain;
#endif
  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(device_, &shader_desc);

  WGPUDepthStencilState depth_stencil = {};
  depth_stencil.format = WGPUTextureFormat_Depth24Plus;
  depth_stencil.depthWriteEnabled = WGPUOptionalBool_True;
  depth_stencil.depthCompare = WGPUCompareFunction_Less;

  WGPURenderPipelineDescriptor desc = {};
  desc.layout = pipeline_layout;
  desc.vertex.module = shader_module;
#if defined(DEMO_CROSS_COMPILE_WIN32)
  desc.vertex.entryPoint = "vs_main";
#else
  desc.vertex.entryPoint = {"vs_main", 7};
#endif
  WGPUColorTargetState color_target = {};
  color_target.format = format_;
  color_target.writeMask = WGPUColorWriteMask_All;
  WGPUFragmentState fragment = {};
  fragment.module = shader_module;
#if defined(DEMO_CROSS_COMPILE_WIN32)
  fragment.entryPoint = "fs_main";
#else
  fragment.entryPoint = {"fs_main", 7};
#endif
  fragment.targetCount = 1;
  fragment.targets = &color_target;
  desc.fragment = &fragment;
  desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  desc.primitive.cullMode = WGPUCullMode_Back;
  desc.primitive.frontFace = WGPUFrontFace_CCW;
  desc.depthStencil = &depth_stencil;
  desc.multisample.count = 1;
  desc.multisample.mask = 0xFFFFFFFF;

  pipeline_ = wgpuDeviceCreateRenderPipeline(device_, &desc);
  wgpuBindGroupLayoutRelease(bgl);
  wgpuPipelineLayoutRelease(pipeline_layout);
  wgpuShaderModuleRelease(shader_module);
}

void Renderer3D::update_uniforms(const Scene& scene, const Camera& camera,
                                 float time) {
  GlobalUniforms globals;
  globals.view_proj = camera.get_projection_matrix() * camera.get_view_matrix();
  globals.camera_pos_time = vec4(camera.position.x, camera.position.y, camera.position.z, time);
  globals.params = vec4((float)std::min((size_t)kMaxObjects, scene.objects.size()), 0.0f, 0.0f, 0.0f);
  
  wgpuQueueWriteBuffer(queue_, global_uniform_buffer_, 0, &globals,
                       sizeof(GlobalUniforms));

  std::vector<ObjectData> obj_data;
  for (const auto& obj : scene.objects) {
    ObjectData data;
    data.model = obj.get_model_matrix();
    
    // Calculate Inverse Transpose for correct normal transformation
    mat4 inverse = data.model.inverse();
    data.model_inverse_transpose = mat4::transpose(inverse);

    data.color = obj.color;
    float type_id = 0.0f;
    if (obj.type == ObjectType::SPHERE)
      type_id = 1.0f;
    else if (obj.type == ObjectType::BOX)
      type_id = 2.0f;
    else if (obj.type == ObjectType::TORUS)
      type_id = 3.0f;
    data.params = vec4(type_id, 0, 0, 0);
    obj_data.push_back(data);
    if (obj_data.size() >= kMaxObjects)
      break;
  }
  if (!obj_data.empty()) {
    wgpuQueueWriteBuffer(queue_, object_storage_buffer_, 0, obj_data.data(),
                         obj_data.size() * sizeof(ObjectData));
  }
}

void Renderer3D::draw(WGPURenderPassEncoder pass, const Scene& scene,
                      const Camera& camera, float time) {
  update_uniforms(scene, camera, time);

  // Lazy Bind Group creation

  if (bind_group_)
    wgpuBindGroupRelease(bind_group_);

  WGPUBindGroupEntry bg_entries[4] = {};

  bg_entries[0].binding = 0;

  bg_entries[0].buffer = global_uniform_buffer_;

  bg_entries[0].size = sizeof(GlobalUniforms);

  bg_entries[1].binding = 1;

  bg_entries[1].buffer = object_storage_buffer_;

  bg_entries[1].size = sizeof(ObjectData) * kMaxObjects;

  bg_entries[2].binding = 2;

  bg_entries[2].textureView = noise_texture_view_;

  bg_entries[3].binding = 3;

  bg_entries[3].sampler = default_sampler_;

  WGPUBindGroupDescriptor bg_desc = {};

  bg_desc.layout = wgpuRenderPipelineGetBindGroupLayout(pipeline_, 0);

  bg_desc.entryCount = 4;

  bg_desc.entries = bg_entries;

  bind_group_ = wgpuDeviceCreateBindGroup(device_, &bg_desc);

  wgpuBindGroupLayoutRelease(bg_desc.layout);

  wgpuRenderPassEncoderSetPipeline(pass, pipeline_);

  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr);

  uint32_t instance_count =
      (uint32_t)std::min((size_t)kMaxObjects, scene.objects.size());

  if (instance_count > 0) {
    wgpuRenderPassEncoderDraw(pass, 36, instance_count, 0, 0);
  }

#if !defined(STRIP_ALL)
  if (s_debug_enabled_) {
    for (const auto& obj : scene.objects) {
      // Simple AABB approximation from scale
      visual_debug_.add_box(obj.position, obj.scale, vec3(1.0f, 1.0f, 0.0f)); // Yellow boxes
    }
    
    // Calculate ViewProj matrix for the debug renderer
    mat4 view_proj = camera.get_projection_matrix() * camera.get_view_matrix();
    visual_debug_.render(pass, view_proj);
  }
#endif
}

void Renderer3D::render(const Scene& scene, const Camera& camera, float time,

                        WGPUTextureView target_view,
                        WGPUTextureView depth_view_opt) {
  WGPUTextureView depth_view = depth_view_opt ? depth_view_opt : depth_view_;

  if (!depth_view)
    return;

  WGPURenderPassColorAttachment color_attachment = {};

  gpu_init_color_attachment(color_attachment, target_view);

  color_attachment.clearValue = {0.05, 0.05, 0.1, 1.0};

  WGPURenderPassDepthStencilAttachment depth_attachment = {};

  depth_attachment.view = depth_view;

  depth_attachment.depthLoadOp = WGPULoadOp_Clear;

  depth_attachment.depthStoreOp = WGPUStoreOp_Store;

  depth_attachment.depthClearValue = 1.0f;

  WGPURenderPassDescriptor pass_desc = {};

  pass_desc.colorAttachmentCount = 1;

  pass_desc.colorAttachments = &color_attachment;

  pass_desc.depthStencilAttachment = &depth_attachment;

  WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device_, nullptr);

  WGPURenderPassEncoder pass =
      wgpuCommandEncoderBeginRenderPass(encoder, &pass_desc);

  draw(pass, scene, camera, time);

  wgpuRenderPassEncoderEnd(pass);

  WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, nullptr);

  wgpuQueueSubmit(queue_, 1, &commands);

  wgpuRenderPassEncoderRelease(pass);

  wgpuCommandBufferRelease(commands);

  wgpuCommandEncoderRelease(encoder);
}