#include <metal_stdlib>
using namespace metal;

struct Uniforms {
  float4x4 mvp;
  float4x4 model;
  float4x4 view;
  float4x4 projection;
  float4 ambient;       // rgb, unused
  float4 light_dir;     // xyz normalized, unused
  float4 light_color;   // rgb, intensity
  float4 fog_params;    // near, far, unused, enabled
  float4 fog_color;     // rgb, unused
  float4 tint;          // rgba (base_color_factor from glTF)
  float4 style_params;  // style_id, vertex_snap, affine, dither
  float4 resolution;    // w, h, unused, unused
  float4 material_params; // alpha_mode (0=opaque,1=mask,2=blend), alpha_cutoff, unlit, unused
};

// Joint palette: up to 64 joints (64 * 64 bytes = 4096 bytes)
struct JointPalette {
  float4x4 joints[64];
};

struct VertexIn {
  float3 position [[attribute(0)]];
  float3 normal   [[attribute(1)]];
  float2 uv       [[attribute(2)]];
  float4 color    [[attribute(3)]];
  float4 joints   [[attribute(4)]];  // Joint indices (as floats)
  float4 weights  [[attribute(5)]];  // Joint weights
};

struct VertexOut {
  float4 position [[position]];
  float3 world_normal;
  float2 uv;
  float4 color;
  float fog_factor;
  float3 noperspective_uv; // For affine texturing (PS1 style)
};

vertex VertexOut vertex_main(
  VertexIn in [[stage_in]],
  constant Uniforms &uniforms [[buffer(1)]],
  constant JointPalette &palette [[buffer(2)]]
) {
  VertexOut out;

  // Skinning: blend position and normal using joint weights
  int4 joint_indices = int4(in.joints);
  float4 weights = in.weights;

  // Compute skinned position
  float4 skinned_pos = float4(0.0);
  float3 skinned_normal = float3(0.0);

  for (int i = 0; i < 4; i++) {
    int idx = joint_indices[i];
    float w = weights[i];
    if (w > 0.0 && idx >= 0 && idx < 64) {
      float4x4 joint_mat = palette.joints[idx];
      skinned_pos += w * (joint_mat * float4(in.position, 1.0));

      // Transform normal (using upper 3x3 of joint matrix)
      float3x3 normal_mat = float3x3(
        joint_mat[0].xyz,
        joint_mat[1].xyz,
        joint_mat[2].xyz
      );
      skinned_normal += w * (normal_mat * in.normal);
    }
  }

  // If no weights, use original position
  float total_weight = weights.x + weights.y + weights.z + weights.w;
  if (total_weight < 0.001) {
    skinned_pos = float4(in.position, 1.0);
    skinned_normal = in.normal;
  }

  float4 world_pos = uniforms.model * skinned_pos;
  float4 clip_pos = uniforms.projection * uniforms.view * world_pos;

  // PS1-style vertex snapping
  float style_id = uniforms.style_params.x;
  float vertex_snap = uniforms.style_params.y;

  if (vertex_snap > 0.5 && style_id < 0.5) {
    // PS1 style: snap vertices to grid
    float2 res = uniforms.resolution.xy;
    float2 snap_res = res * 0.5;

    // Snap in clip space after perspective divide
    float4 snapped = clip_pos;
    if (snapped.w > 0.0) {
      float2 ndc = snapped.xy / snapped.w;
      ndc = floor(ndc * snap_res + 0.5) / snap_res;
      snapped.xy = ndc * snapped.w;
    }
    clip_pos = snapped;
  }

  out.position = clip_pos;

  // Transform normal to world space
  float3x3 model_normal_matrix = float3x3(
    uniforms.model[0].xyz,
    uniforms.model[1].xyz,
    uniforms.model[2].xyz
  );
  out.world_normal = normalize(model_normal_matrix * skinned_normal);

  // UV coordinates
  out.uv = in.uv;

  // For affine texturing (PS1), we pass UV * w to fragment shader
  // and divide by interpolated w there
  float affine = uniforms.style_params.z;
  if (affine > 0.5) {
    out.noperspective_uv = float3(in.uv * clip_pos.w, clip_pos.w);
  } else {
    out.noperspective_uv = float3(in.uv, 1.0);
  }

  // Vertex color
  out.color = in.color;

  // Fog calculation (linear fog)
  float fog_enabled = uniforms.fog_params.w;
  if (fog_enabled > 0.5) {
    float4 view_pos = uniforms.view * world_pos;
    float dist = length(view_pos.xyz);
    float fog_near = uniforms.fog_params.x;
    float fog_far = uniforms.fog_params.y;
    out.fog_factor = clamp((fog_far - dist) / (fog_far - fog_near), 0.0, 1.0);
  } else {
    out.fog_factor = 1.0;
  }

  return out;
}