#!/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/effects/<effect>_effect.h
#   - src/effects/<effect>_effect.cc
#   - workspaces/main/shaders/<effect>.wgsl
#
# The script performs basic ShaderToy→WGSL conversion:
#   - Converts types (float→f32, vec2→vec2f, 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', 'in.position.xy'),

        # Type conversions (use short form vec2f etc.)
        (r'\bfloat\b', 'f32'),
        (r'\bvec2\b', 'vec2f'),
        (r'\bvec3\b', 'vec3f'),
        (r'\bvec4\b', 'vec4f'),
        (r'\bmat2\b', 'mat2x2f'),
        (r'\bmat3\b', 'mat3x3f'),
        (r'\bmat4\b', 'mat4x4f'),

        # Function declarations (preserve return type context)
        (r'\bf32\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvec2f\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvec3f\s+(\w+)\s*\(', r'fn \1('),
        (r'\bvec4f\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+vec2f\s+(\w+)\s*=', r'const \1 ='),
        (r'\bconst\s+vec3f\s+(\w+)\s*=', r'const \1 ='),
        (r'\bconst\s+vec4f\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(in: VertexOutput) -> @location(0) vec4f {{
  // Flip Y to match ShaderToy convention (origin at bottom-left)
  let flipped = vec2f(in.position.x, uniforms.resolution.y - in.position.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 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

#pragma once

#include "gpu/effect.h"
#include "gpu/wgpu_resource.h"

class {class_name} : public Effect {{
 public:
  {class_name}(const GpuContext& ctx,
               const std::vector<std::string>& inputs,
               const std::vector<std::string>& outputs,
               float start_time, float end_time);

  void render(WGPUCommandEncoder encoder,
              const UniformsSequenceParams& params,
              NodeRegistry& nodes) override;

 private:
  RenderPipeline pipeline_;
  BindGroup bind_group_;
}};
"""
    else:
        return f"""// This file is part of the 64k demo project.
// {effect_name} effect - ShaderToy conversion (scene)
// Generated by convert_shadertoy.py

#pragma once

#include "gpu/effect.h"
#include "gpu/wgpu_resource.h"

class {class_name} : public Effect {{
 public:
  {class_name}(const GpuContext& ctx,
               const std::vector<std::string>& inputs,
               const std::vector<std::string>& outputs,
               float start_time, float end_time);

  void render(WGPUCommandEncoder encoder,
              const UniformsSequenceParams& params,
              NodeRegistry& nodes) override;

 private:
  RenderPipeline pipeline_;
  BindGroup bind_group_;
}};
"""

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 "effects/{snake_name}_effect.h"
#include "effects/shaders.h"
#include "gpu/gpu.h"
#include "gpu/post_process_helper.h"
#include "util/fatal_error.h"

{class_name}::{class_name}(const GpuContext& ctx,
                            const std::vector<std::string>& inputs,
                            const std::vector<std::string>& outputs,
                            float start_time, float end_time)
    : Effect(ctx, inputs, outputs, start_time, end_time) {{
  HEADLESS_RETURN_IF_NULL(ctx_.device);
  create_linear_sampler();
  pipeline_.set(create_post_process_pipeline(
      ctx_.device, WGPUTextureFormat_RGBA8Unorm, {snake_name}_shader_wgsl));
}}

void {class_name}::render(WGPUCommandEncoder encoder,
                           const UniformsSequenceParams& params,
                           NodeRegistry& nodes) {{
  WGPUTextureView input_view = nodes.get_view(input_nodes_[0]);
  WGPUTextureView output_view = nodes.get_view(output_nodes_[0]);

  // uniforms_buffer_ auto-updated by base class dispatch_render()
  pp_update_bind_group(ctx_.device, pipeline_.get(), bind_group_.get_address(),
                       input_view, uniforms_buffer_.get(), {{nullptr, 0}});

  WGPURenderPassColorAttachment color_attachment = {{}};
  gpu_init_color_attachment(color_attachment, output_view);

  WGPURenderPassDescriptor pass_desc = {{}};
  pass_desc.colorAttachmentCount = 1;
  pass_desc.colorAttachments = &color_attachment;

  WGPURenderPassEncoder pass =
      wgpuCommandEncoderBeginRenderPass(encoder, &pass_desc);
  wgpuRenderPassEncoderSetPipeline(pass, pipeline_.get());
  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_.get(), 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, 3, 1, 0, 0);
  wgpuRenderPassEncoderEnd(pass);
  wgpuRenderPassEncoderRelease(pass);
}}
"""
    else:
        return f"""// This file is part of the 64k demo project.
// {effect_name} effect - ShaderToy conversion (scene)
// Generated by convert_shadertoy.py

#include "effects/{snake_name}_effect.h"
#include "effects/shaders.h"
#include "gpu/gpu.h"
#include "gpu/post_process_helper.h"
#include "util/fatal_error.h"

{class_name}::{class_name}(const GpuContext& ctx,
                            const std::vector<std::string>& inputs,
                            const std::vector<std::string>& outputs,
                            float start_time, float end_time)
    : Effect(ctx, inputs, outputs, start_time, end_time) {{
  HEADLESS_RETURN_IF_NULL(ctx_.device);
  create_nearest_sampler();
  create_dummy_scene_texture();
  pipeline_.set(create_post_process_pipeline(
      ctx_.device, WGPUTextureFormat_RGBA8Unorm, {snake_name}_shader_wgsl));
}}

void {class_name}::render(WGPUCommandEncoder encoder,
                           const UniformsSequenceParams& params,
                           NodeRegistry& nodes) {{
  WGPUTextureView output_view = nodes.get_view(output_nodes_[0]);

  // uniforms_buffer_ auto-updated by base class dispatch_render()
  pp_update_bind_group(ctx_.device, pipeline_.get(), bind_group_.get_address(),
                       dummy_texture_view_.get(), uniforms_buffer_.get(),
                       {{nullptr, 0}});

  WGPURenderPassColorAttachment color_attachment = {{}};
  gpu_init_color_attachment(color_attachment, output_view);

  WGPURenderPassDescriptor pass_desc = {{}};
  pass_desc.colorAttachmentCount = 1;
  pass_desc.colorAttachments = &color_attachment;

  WGPURenderPassEncoder pass =
      wgpuCommandEncoderBeginRenderPass(encoder, &pass_desc);
  wgpuRenderPassEncoderSetPipeline(pass, pipeline_.get());
  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_.get(), 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, 3, 1, 0, 0);
  wgpuRenderPassEncoderEnd(pass);
  wgpuRenderPassEncoderRelease(pass);
}}
"""

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 "sequence_uniforms"
#include "render/fullscreen_uv_vs"

@group(0) @binding(2) var<uniform> uniforms: UniformsSequenceParams;"""
    else:
        # Scene effect - no texture input
        bindings = """#include "sequence_uniforms"
#include "render/fullscreen_uv_vs"

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

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

{bindings}

{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 from scratch)")
        print("  --shader-only   Only regenerate .wgsl shader (skip .h/.cc files)")
        print("  --dry-run       Print next steps without writing any files")
        print()
        print("This will generate:")
        print("  src/effects/<effect>_effect.h")
        print("  src/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
    dry_run = '--dry-run' 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" / "effects" / f"{snake_name}_effect.h"
    impl_path = repo_root / "src" / "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()
        if not dry_run:
            shader_path.write_text(generate_shader(effect_name, shadertoy_code, is_post_process))
            print(f"✓ Shader regenerated")
        else:
            print(f"[dry-run] Would write shader")
        return

    effect_type = "post-process" if is_post_process else "scene"
    if dry_run:
        print(f"[dry-run] Would generate effect: {effect_name} ({effect_type})")
        print(f"  Header:     {header_path}")
        print(f"  Impl:       {impl_path}")
        print(f"  Shader:     {shader_path}")
    else:
        print(f"Generating effect: {effect_name}")
        print(f"  Header:     {header_path}")
        print(f"  Impl:       {impl_path}")
        print(f"  Shader:     {shader_path}")
        print()
        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))
        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("2. Add shader declaration to src/effects/shaders.h:")
    print(f"   extern const char* {snake_name}_shader_wgsl;")
    print()
    print("3. Add shader definition to src/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 "effects/{snake_name}_effect.h"')
    print()
    print("5. Add to cmake/DemoSourceLists.cmake COMMON_GPU_EFFECTS list:")
    print(f"   src/effects/{snake_name}_effect.cc")
    print()
    print("6. Add to timeline in workspaces/main/timeline.seq:")
    print(f"   EFFECT + {effect_name}Effect source -> sink 0.0 10.0")
    print()
    print("7. Regenerate timeline.cc:")
    print("   python3 tools/seq_compiler.py workspaces/main/timeline.seq \\")
    print("     --output src/generated/timeline.cc")
    print()
    print("8. Add to src/tests/gpu/test_demo_effects.cc:")
    print(f'   {{"{effect_name}", std::make_shared<{effect_name}Effect>(')
    print(f'       fixture.ctx(), inputs, outputs, 0.0f, 1000.0f)}}')
    print()
    print("9. 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()