#!/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."""
    # Extract mainImage first
    main_match = re.search(r'void\s+mainImage\s*\([^)]+\)\s*\{(.*)\}', shader_code, re.DOTALL)
    if main_match:
        main_body = main_match.group(1).strip()
        helpers = shader_code[:main_match.start()]
    else:
        main_body = ""
        helpers = shader_code

    # Replace common ShaderToy defines
    conversions = [
        (r'#define\s+TIME\s+iTime', ''),
        (r'#define\s+RESOLUTION\s+iResolution', ''),
        (r'#define\s+PI\s+[\d.]+', 'const PI: f32 = 3.141592654;'),
        (r'#define\s+TAU\s+\([^)]+\)', 'const TAU: f32 = 6.283185307;'),
        (r'#define\s+ROT\(a\)\s+mat2\([^)]+\)', ''),  # Will be converted to function

        # Common ShaderToy uniforms
        (r'\bTIME\b', 'uniforms.time'),
        (r'\biTime\b', 'uniforms.time'),
        (r'\bRESOLUTION\b', 'uniforms.resolution'),
        (r'\biResolution\b', 'uniforms.resolution'),
        (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 (preserve return type context)
        (r'\bf32\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvec2<f32>\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvec3<f32>\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvec4<f32>\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvoid\s+(\w+)\s*\(', r'fn \1('),

        # Const declarations
        (r'\bconst\s+f32\s+(\w+)\s*=', r'const \1: f32 ='),
        (r'\bconst\s+vec2<f32>\s+(\w+)\s*=', r'const \1 ='),
        (r'\bconst\s+vec3<f32>\s+(\w+)\s*=', r'const \1 ='),
        (r'\bconst\s+vec4<f32>\s+(\w+)\s*=', r'const \1 ='),

        # Function calls that need fixing
        (r'\bfract\s*\(', 'fract('),
        (r'\bmod\s*\(([^,]+),\s*([^)]+)\)', r'(\1 % \2)'),
    ]

    converted_helpers = helpers
    for pattern, replacement in conversions:
        converted_helpers = re.sub(pattern, replacement, converted_helpers)

    # Convert mainImage body
    converted_main = main_body
    for pattern, replacement in conversions:
        converted_main = re.sub(pattern, replacement, converted_main)

    # Fix fragColor assignments -> returns
    converted_main = re.sub(r'\bfragColor\s*=\s*([^;]+);', r'return \1;', converted_main)

    # Indent main body
    indented_main = '\n'.join('  ' + line if line.strip() else '' for line in converted_main.split('\n'))

    # Build fragment function with Y-flip for ShaderToy convention
    fragment = f"""@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {{
  // Flip Y to match ShaderToy convention (origin at bottom-left)
  let flipped = vec2<f32>(p.x, uniforms.resolution.y - p.y);
  let q = flipped / uniforms.resolution;
  var coord = -1.0 + 2.0 * q;
  coord.x *= uniforms.resolution.x / uniforms.resolution.y;

{indented_main}
}}"""

    return converted_helpers + '\n\n' + fragment

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, is_post_process=False):
    """Generate .h file content."""
    class_name = f"{effect_name}Effect"
    upper_name = to_upper_snake_case(effect_name)

    if is_post_process:
        return f"""// This file is part of the 64k demo project.
// {effect_name} effect - ShaderToy conversion (post-process)
// 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"

class {class_name} : public PostProcessEffect {{
 public:
  {class_name}(const GpuContext& ctx);
  void render(WGPURenderPassEncoder pass, float time, float beat,
              float intensity, float aspect_ratio) override;
  void update_bind_group(WGPUTextureView input_view) override;
}};

#endif /* {upper_name}_EFFECT_H_ */
"""
    else:
        # Scene effect (simpler, like HeptagonEffect)
        return f"""// This file is part of the 64k demo project.
// {effect_name} effect - ShaderToy conversion (scene)
// Generated by convert_shadertoy.py

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

#include "gpu/effect.h"

class {class_name} : public Effect {{
 public:
  {class_name}(const GpuContext& ctx);
  void render(WGPURenderPassEncoder pass, float time, float beat,
              float intensity, float aspect_ratio) override;

 private:
  RenderPass pass_;
}};

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

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

    if is_post_process:
        return f"""// This file is part of the 64k demo project.
// {effect_name} effect - ShaderToy conversion (post-process)
// Generated by convert_shadertoy.py

#include "gpu/demo_effects.h"
#include "gpu/effects/post_process_helper.h"
#include "gpu/gpu.h"

{class_name}::{class_name}(const GpuContext& ctx) : PostProcessEffect(ctx) {{
  pipeline_ = create_post_process_pipeline(ctx_.device, ctx_.format, {snake_name}_shader_wgsl);
}}

void {class_name}::render(WGPURenderPassEncoder pass, float time, float beat,
                          float intensity, float aspect_ratio) {{
  const CommonPostProcessUniforms u = {{
      .resolution = {{(float)width_, (float)height_}},
      ._pad = {{0.0f, 0.0f}},
      .aspect_ratio = aspect_ratio,
      .time = time,
      .beat = beat,
      .audio_intensity = intensity,
  }};
  uniforms_.update(ctx_.queue, u);

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

void {class_name}::update_bind_group(WGPUTextureView input_view) {{
  pp_update_bind_group(ctx_.device, pipeline_, &bind_group_, input_view, uniforms_.get());
}}
"""
    else:
        # Scene effect (simpler pattern like HeptagonEffect)
        return f"""// This file is part of the 64k demo project.
// {effect_name} effect - ShaderToy conversion (scene)
// Generated by convert_shadertoy.py

#include "gpu/demo_effects.h"
#include "gpu/gpu.h"

{class_name}::{class_name}(const GpuContext& ctx) : Effect(ctx) {{
  ResourceBinding bindings[] = {{{{uniforms_.get(), WGPUBufferBindingType_Uniform}}}};
  pass_ = gpu_create_render_pass(ctx_.device, ctx_.format, {snake_name}_shader_wgsl,
                                 bindings, 1);
  pass_.vertex_count = 3;
}}

void {class_name}::render(WGPURenderPassEncoder pass, float t, float b,
                          float i, float a) {{
  CommonPostProcessUniforms u = {{
      .resolution = {{(float)width_, (float)height_}},
      ._pad = {{0.0f, 0.0f}},
      .aspect_ratio = a,
      .time = t,
      .beat = b,
      .audio_intensity = i,
  }};
  uniforms_.update(ctx_.queue, u);
  wgpuRenderPassEncoderSetPipeline(pass, pass_.pipeline);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, pass_.bind_group, 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, pass_.vertex_count, 1, 0, 0);
}}
"""

def generate_shader(effect_name, shadertoy_code, is_post_process=False):
    """Generate .wgsl file content."""
    # Convert to WGSL (full shader, not just mainImage)
    converted = convert_shadertoy_to_wgsl(shadertoy_code)

    if is_post_process:
        bindings = """@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;"""
    else:
        # Scene effect - only uniforms, no texture input
        bindings = """#include "common_uniforms"

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

    return f"""// {effect_name} effect shader - ShaderToy conversion
// Generated by convert_shadertoy.py
// NOTE: Manual review recommended - conversion is basic

{bindings}

@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}
"""

def main():
    if len(sys.argv) < 3:
        print("Usage: convert_shadertoy.py <shader.txt> <EffectName> [--post-process] [--shader-only]")
        print()
        print("Examples:")
        print("  ./tools/shadertoy/convert_shadertoy.py tunnel.txt Tunnel")
        print("  ./tools/shadertoy/convert_shadertoy.py blur.txt Blur --post-process")
        print("  ./tools/shadertoy/convert_shadertoy.py tunnel.txt Tunnel --shader-only")
        print()
        print("Options:")
        print("  --post-process  Generate post-process effect (operates on previous frame)")
        print("                  Default: scene effect (renders geometry)")
        print("  --shader-only   Only regenerate .wgsl shader (skip .h/.cc files)")
        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]
    is_post_process = '--post-process' in sys.argv
    shader_only = '--shader-only' in sys.argv

    # 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)
    upper_name = to_upper_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
    if shader_only:
        print(f"Regenerating shader only: {effect_name}")
        print(f"  Shader:     {shader_path}")
        print()
        shader_path.write_text(generate_shader(effect_name, shadertoy_code, is_post_process))
        print(f"✓ Shader regenerated")
        return

    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, is_post_process))
    impl_path.write_text(generate_implementation(effect_name, is_post_process))
    shader_path.write_text(generate_shader(effect_name, shadertoy_code, is_post_process))

    effect_type = "post-process" if is_post_process else "scene"
    print(f"✓ Files generated ({effect_type} effect)")
    print()
    print("Next steps (see doc/EFFECT_WORKFLOW.md for details):")
    print()
    print("1. Add shader to workspaces/main/assets.txt:")
    print(f"   SHADER_{upper_name}, NONE, shaders/{snake_name}.wgsl, \"{effect_name} effect\"")
    print()
    print()
    print("2. Add shader declaration to src/gpu/effects/shaders.h:")
    print(f"   extern const char* {snake_name}_shader_wgsl;")
    print()
    print("3. Add shader definition to src/gpu/effects/shaders.cc:")
    print(f"   const char* {snake_name}_shader_wgsl = SafeGetAsset(AssetId::ASSET_SHADER_{upper_name});")
    print()
    print("4. Include header in src/gpu/demo_effects.h:")
    print(f'   #include "gpu/effects/{snake_name}_effect.h"')
    print()
    print("5. Add to timeline in workspaces/main/timeline.seq:")
    print(f"   EFFECT + {effect_name}Effect 0.0 10.0")
    print()
    print("6. Add to CMakeLists.txt GPU_SOURCES (both headless and normal mode):")
    print(f"   src/gpu/effects/{snake_name}_effect.cc")
    print()
    print("7. Update src/tests/gpu/test_demo_effects.cc:")
    test_list = "post_process_effects" if is_post_process else "scene_effects"
    print(f'   - Add "{{{effect_name}Effect", std::make_shared<{effect_name}Effect>(fixture.ctx())}}" to {test_list} list')
    print()
    print("8. Build and test:")
    print("   cmake --build build -j4")
    print("   ./build/demo64k")
    print()
    print("Note: Review generated shader for const expression issues (normalize, etc)")

if __name__ == '__main__':
    main()