#!/usr/bin/env python3
# This file is part of the 64k demo project.
# Converts ShaderToy shader to demo effect boilerplate.
#
# Usage:
#   ./tools/shadertoy/convert_shadertoy.py <shader.txt> <EffectName>
#
# Example:
#   ./tools/shadertoy/convert_shadertoy.py tunnel.txt Tunnel
#   ./tools/shadertoy/convert_shadertoy.py tools/shadertoy/example.txt Rainbow
#
# Generates:
#   - src/gpu/effects/<effect>_effect.h
#   - src/gpu/effects/<effect>_effect.cc
#   - workspaces/main/shaders/<effect>.wgsl
#
# The script performs basic ShaderToy→WGSL conversion:
#   - Converts types (float→f32, vec2→vec2<f32>, etc.)
#   - Converts uniforms (iTime→uniforms.time, etc.)
#   - Extracts mainImage() body into fs_main()
#   - Generates boilerplate C++ effect class
#
# Manual fixes usually needed:
#   - fragColor assignments → return statements
#   - Variable name conflicts (e.g., shadowing 'p')
#   - Complex type inference
#   - Texture channel mappings
#   - Helper function signatures

import sys
import os
import re
from pathlib import Path

def to_snake_case(name):
    """Convert CamelCase to snake_case."""
    s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()

def to_upper_snake_case(name):
    """Convert CamelCase to UPPER_SNAKE_CASE."""
    return to_snake_case(name).upper()

def to_camel_case(name):
    """Convert snake_case to CamelCase."""
    return ''.join(word.capitalize() for word in name.split('_'))

def convert_shadertoy_to_wgsl(shader_code):
    """Basic ShaderToy to WGSL conversion."""
    # Replace common ShaderToy uniforms
    conversions = [
        (r'\biResolution\b', 'uniforms.resolution'),
        (r'\biTime\b', 'uniforms.time'),
        (r'\bfragCoord\b', 'p.xy'),

        # Type conversions
        (r'\bfloat\b', 'f32'),
        (r'\bvec2\b', 'vec2<f32>'),
        (r'\bvec3\b', 'vec3<f32>'),
        (r'\bvec4\b', 'vec4<f32>'),
        (r'\bmat2\b', 'mat2x2<f32>'),
        (r'\bmat3\b', 'mat3x3<f32>'),
        (r'\bmat4\b', 'mat4x4<f32>'),

        # Function declarations
        (r'\bfloat\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvec2\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvec3\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvec4\s+(\w+)\s*\(', r'fn \1('),

        # Function parameters (simple cases)
        (r'\(\s*vec2\s+(\w+)\s*\)', r'(\1: vec2<f32>)'),
        (r'\(\s*vec3\s+(\w+)\s*\)', r'(\1: vec3<f32>)'),
        (r'\(\s*float\s+(\w+)\s*\)', r'(\1: f32)'),

        # Return types
        (r'\)\s*{', r') -> f32 {'),  # Assume f32 return, may need manual fix

        # Texture sampling
        (r'\btexture\s*\(\s*iChannel0\s*,', 'textureSample(txt, smplr,'),
        (r'\btexture\s*\(\s*iChannel1\s*,', 'textureSample(txt, smplr,'),
    ]

    converted = shader_code
    for pattern, replacement in conversions:
        converted = re.sub(pattern, replacement, converted)

    # Convert variable declarations
    converted = re.sub(r'\bfloat\s+(\w+)\s*=', r'var \1: f32 =', converted)
    converted = re.sub(r'\bvec2\s+(\w+)\s*=', r'var \1: vec2<f32> =', converted)
    converted = re.sub(r'\bvec3\s+(\w+)\s*=', r'var \1: vec3<f32> =', converted)
    converted = re.sub(r'\bvec4\s+(\w+)\s*=', r'var \1: vec4<f32> =', converted)

    return converted

def extract_main_image(shader_code):
    """Extract mainImage function body from ShaderToy code."""
    # Try to find mainImage function
    match = re.search(r'void\s+mainImage\s*\([^)]+\)\s*\{(.*)\}', shader_code, re.DOTALL)
    if match:
        return match.group(1).strip()

    # If no mainImage found, return whole shader
    return shader_code

def generate_header(effect_name):
    """Generate .h file content."""
    class_name = f"{effect_name}Effect"
    snake_name = to_snake_case(effect_name)
    upper_name = to_upper_snake_case(effect_name)

    return f"""// This file is part of the 64k demo project.
// {effect_name} effect - ShaderToy conversion
// Generated by convert_shadertoy.py

#ifndef {upper_name}_EFFECT_H_
#define {upper_name}_EFFECT_H_

#include "gpu/effect.h"
#include "gpu/effects/post_process_helper.h"
#include "gpu/uniform_helper.h"

class {class_name} : public Effect {{
 public:
  {class_name}(const GpuContext& ctx);
  ~{class_name}() override;

  void init(MainSequence* demo) override;
  void render(WGPURenderPassEncoder pass, float time, float beat,
              float intensity, float aspect_ratio) override;

 private:
  // Effect-specific parameters - adjust as needed
  struct {effect_name}Params {{
    float param1;
    float param2;
    float _pad[2];
  }};
  static_assert(sizeof({effect_name}Params) == 16,
                "{effect_name}Params must be 16 bytes for WGSL alignment");

  MainSequence* demo_ = nullptr;
  WGPURenderPipeline pipeline_ = nullptr;
  WGPUBindGroup bind_group_ = nullptr;
  WGPUSampler sampler_ = nullptr;
  UniformBuffer<{effect_name}Params> params_;
}};

#endif /* {upper_name}_EFFECT_H_ */
"""

def generate_implementation(effect_name):
    """Generate .cc file content."""
    class_name = f"{effect_name}Effect"
    snake_name = to_snake_case(effect_name)
    upper_name = to_upper_snake_case(effect_name)

    return f"""// This file is part of the 64k demo project.
// {effect_name} effect implementation - ShaderToy conversion
// Generated by convert_shadertoy.py

#include "gpu/effects/{snake_name}_effect.h"
#include "gpu/effects/shader_composer.h"
#include "generated/assets.h"

{class_name}::{class_name}(const GpuContext& ctx) : Effect(ctx) {{
}}

{class_name}::~{class_name}() {{
  if (sampler_)
    wgpuSamplerRelease(sampler_);
  if (bind_group_)
    wgpuBindGroupRelease(bind_group_);
  if (pipeline_)
    wgpuRenderPipelineRelease(pipeline_);
}}

void {class_name}::init(MainSequence* demo) {{
  demo_ = demo;
  params_.init(ctx_.device);

  WGPUSamplerDescriptor sampler_desc = {{}};
  sampler_desc.addressModeU = WGPUAddressMode_ClampToEdge;
  sampler_desc.addressModeV = WGPUAddressMode_ClampToEdge;
  sampler_desc.magFilter = WGPUFilterMode_Linear;
  sampler_desc.minFilter = WGPUFilterMode_Linear;
  sampler_desc.mipmapFilter = WGPUMipmapFilterMode_Linear;
  sampler_desc.maxAnisotropy = 1;
  sampler_ = wgpuDeviceCreateSampler(ctx_.device, &sampler_desc);

  size_t shader_size;
  const char* shader_code = (const char*)GetAsset(
      AssetId::ASSET_{upper_name}_SHADER, &shader_size);

  std::string composed = ShaderComposer::Get().Compose({{}}, shader_code);

  WGPUShaderSourceWGSL wgsl = {{}};
  wgsl.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl.code = str_view(composed.c_str());

  WGPUShaderModuleDescriptor desc = {{}};
  desc.nextInChain = &wgsl.chain;
  WGPUShaderModule module = wgpuDeviceCreateShaderModule(ctx_.device, &desc);

  const WGPUColorTargetState target = {{
      .format = ctx_.format,
      .writeMask = WGPUColorWriteMask_All,
  }};
  WGPUFragmentState frag = {{}};
  frag.module = module;
  frag.entryPoint = str_view("fs_main");
  frag.targetCount = 1;
  frag.targets = &target;

  const WGPUDepthStencilState depth_stencil = {{
      .format = WGPUTextureFormat_Depth24Plus,
      .depthWriteEnabled = WGPUOptionalBool_False,
      .depthCompare = WGPUCompareFunction_Always,
  }};

  WGPURenderPipelineDescriptor pipeline_desc = {{}};
  pipeline_desc.label = label_view("{class_name}");
  pipeline_desc.vertex.module = module;
  pipeline_desc.vertex.entryPoint = str_view("vs_main");
  pipeline_desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  pipeline_desc.primitive.cullMode = WGPUCullMode_None;
  pipeline_desc.depthStencil = &depth_stencil;
  pipeline_desc.multisample.count = 1;
  pipeline_desc.multisample.mask = 0xFFFFFFFF;
  pipeline_desc.fragment = &frag;

  pipeline_ = wgpuDeviceCreateRenderPipeline(ctx_.device, &pipeline_desc);
  wgpuShaderModuleRelease(module);

  WGPUTextureView prev_view = demo_->get_prev_texture_view();
  const WGPUBindGroupEntry entries[] = {{
      {{.binding = 0, .sampler = sampler_}},
      {{.binding = 1, .textureView = prev_view}},
      {{.binding = 2,
       .buffer = uniforms_.get().buffer,
       .size = sizeof(CommonPostProcessUniforms)}},
      {{.binding = 3,
       .buffer = params_.get().buffer,
       .size = sizeof({effect_name}Params)}},
  }};
  const WGPUBindGroupDescriptor bg_desc = {{
      .layout = wgpuRenderPipelineGetBindGroupLayout(pipeline_, 0),
      .entryCount = 4,
      .entries = entries,
  }};
  bind_group_ = wgpuDeviceCreateBindGroup(ctx_.device, &bg_desc);
}}

void {class_name}::render(WGPURenderPassEncoder pass, float time,
                          float beat, float intensity, float aspect_ratio) {{
  const CommonPostProcessUniforms uniforms = {{
      .resolution = {{static_cast<float>(width_), static_cast<float>(height_)}},
      .aspect_ratio = aspect_ratio,
      .time = time,
      .beat = beat,
      .audio_intensity = intensity,
  }};
  uniforms_.update(ctx_.queue, uniforms);

  const {effect_name}Params params = {{
      .param1 = 1.0f,
      .param2 = beat,
  }};
  params_.update(ctx_.queue, params);

  wgpuRenderPassEncoderSetPipeline(pass, pipeline_);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, 3, 1, 0, 0);
}}
"""

def generate_shader(effect_name, shadertoy_code):
    """Generate .wgsl file content."""
    # Extract mainImage body
    main_body = extract_main_image(shadertoy_code)

    # Convert to WGSL
    converted = convert_shadertoy_to_wgsl(main_body)

    # Indent the converted code
    indented = '\n'.join('    ' + line if line.strip() else ''
                         for line in converted.split('\n'))

    return f"""// {effect_name} effect shader - ShaderToy conversion
// Generated by convert_shadertoy.py
// TODO: Review and fix conversion issues

@group(0) @binding(0) var smplr: sampler;
@group(0) @binding(1) var txt: texture_2d<f32>;

#include "common_uniforms"

@group(0) @binding(2) var<uniform> uniforms: CommonUniforms;

struct {effect_name}Params {{
    param1: f32,
    param2: f32,
    _pad0: f32,
    _pad1: f32,
}}

@group(0) @binding(3) var<uniform> params: {effect_name}Params;

@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> {{
    var pos = array<vec2<f32>, 3>(
        vec2<f32>(-1.0, -1.0),
        vec2<f32>(3.0, -1.0),
        vec2<f32>(-1.0, 3.0)
    );
    return vec4<f32>(pos[i], 0.0, 1.0);
}}

// ============================================================================
// CONVERTED SHADERTOY CODE (may need manual fixes)
// ============================================================================
// Original helper functions (if any) appear above mainImage conversion

@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {{
    // Converted from mainImage()
{indented}

    // TODO: Fix the return statement if needed
    // ShaderToy used: fragColor = ...
    // WGSL needs: return vec4<f32>(...);
}}
"""

def main():
    if len(sys.argv) < 3:
        print("Usage: convert_shadertoy.py <shader.txt> <EffectName>")
        print()
        print("Examples:")
        print("  ./tools/shadertoy/convert_shadertoy.py tunnel.txt Tunnel")
        print("  ./tools/shadertoy/convert_shadertoy.py plasma.txt Plasma")
        print("  ./tools/shadertoy/convert_shadertoy.py tools/shadertoy/example.txt Rainbow")
        print()
        print("This will generate:")
        print("  src/gpu/effects/<effect>_effect.h")
        print("  src/gpu/effects/<effect>_effect.cc")
        print("  workspaces/main/shaders/<effect>.wgsl")
        sys.exit(1)

    shader_file = sys.argv[1]
    effect_name = sys.argv[2]

    # Ensure effect name is CamelCase
    if '_' in effect_name:
        effect_name = to_camel_case(effect_name)

    # Read shader code
    if not os.path.exists(shader_file):
        print(f"Error: {shader_file} not found")
        sys.exit(1)

    with open(shader_file, 'r') as f:
        shadertoy_code = f.read()

    # Generate file names
    snake_name = to_snake_case(effect_name)

    # Script is in tools/shadertoy/, so go up two levels to repo root
    repo_root = Path(__file__).parent.parent.parent
    header_path = repo_root / "src" / "gpu" / "effects" / f"{snake_name}_effect.h"
    impl_path = repo_root / "src" / "gpu" / "effects" / f"{snake_name}_effect.cc"
    shader_path = repo_root / "workspaces" / "main" / "shaders" / f"{snake_name}.wgsl"

    # Generate files
    print(f"Generating effect: {effect_name}")
    print(f"  Header:     {header_path}")
    print(f"  Impl:       {impl_path}")
    print(f"  Shader:     {shader_path}")
    print()

    # Write files
    header_path.write_text(generate_header(effect_name))
    impl_path.write_text(generate_implementation(effect_name))
    shader_path.write_text(generate_shader(effect_name, shadertoy_code))

    print("✓ Files generated")
    print()
    print("Next steps (see doc/EFFECT_WORKFLOW.md for details):")
    print()
    print("1. Add shader to workspaces/main/assets.txt:")
    print(f"   shaders/{snake_name}.wgsl")
    print()
    print("2. Include header in src/gpu/demo_effects.h:")
    print(f'   #include "gpu/effects/{snake_name}_effect.h"')
    print()
    print("3. Add to timeline in workspaces/main/timeline.seq:")
    print(f"   EFFECT + {effect_name}Effect 0.0 10.0")
    print()
    print("4. Add to CMakeLists.txt GPU_SOURCES (both headless and normal mode):")
    print(f"   src/gpu/effects/{snake_name}_effect.cc")
    print()
    print("5. Update src/tests/gpu/test_demo_effects.cc:")
    print(f'   - Add "{{{effect_name}Effect", std::make_shared<{effect_name}Effect>(fixture.ctx())}}" to appropriate list')
    print("   - Use post_process_effects (lines 80-93) for post-process effects")
    print("   - Use scene_effects (lines 125-137) for scene effects")
    print()
    print("6. Build and test:")
    print("   cmake --build build -j4")
    print("   ./build/demo64k")
    print()
    print("Note: WGSL conversion is basic - review and fix shader manually!")

if __name__ == '__main__':
    main()