/**
 * webgpu.ts — GPU compute from the agent. WGSL compute shaders + image kernels.
 *
 * Tools:
 *   webgpu_info          — adapter info, limits, features. Run FIRST to confirm support.
 *   webgpu_compute       — simple single-buffer compute (input mutated in place, returned)
 *   webgpu_compute_multi — multi-buffer compute: separate read-only inputs + output + uniforms
 *   webgpu_benchmark     — run a shader N times, report GPU timing (uses timestamp-query if available)
 *   webgpu_image_kernel  — apply a WGSL kernel to image pixel data (RGBA8 in/out)
 *   webgpu_examples      — reference library of working WGSL shaders (reduce, matmul, blur, etc.)
 */
import { tool } from '@strands-agents/sdk'
import { z } from 'zod'

// --- Device lifecycle (with loss handling) ---------------------------------

let _device: GPUDevice | null = null
let _adapter: GPUAdapter | null = null
let _devicePromise: Promise<GPUDevice> | null = null
let _deviceLostReason: string | null = null

async function getDevice(): Promise<GPUDevice> {
  if (_device && !_deviceLostReason) return _device
  if (_devicePromise) return _devicePromise

  _devicePromise = (async () => {
    const nav = navigator as any
    if (!nav.gpu) throw new Error('WebGPU not supported in this browser. Try Chrome/Edge 113+ or Safari 18+.')
    _adapter = await nav.gpu.requestAdapter({ powerPreference: 'high-performance' })
    if (!_adapter) throw new Error('No WebGPU adapter available (GPU likely disabled or blocklisted).')

    const device = await _adapter!.requestDevice({
      requiredFeatures: [],
      // Ask for timestamp-query if available — used by benchmark.
      ...((_adapter!.features as any)?.has?.('timestamp-query') ? { requiredFeatures: ['timestamp-query'] } : {}),
    })

    // Install loss handler — clear cache so next call retries.
    device.lost.then((reason) => {
      _deviceLostReason = `device lost: ${reason.reason} — ${reason.message || 'no message'}`
      _device = null
      _devicePromise = null
      console.warn('[webgpu]', _deviceLostReason)
    })

    // Install a global uncapturederror listener — surface shader/runtime errors.
    device.addEventListener('uncapturederror', (ev: any) => {
      console.error('[webgpu uncaught]', ev.error?.message || ev.error)
    })

    _device = device
    _deviceLostReason = null
    return device
  })().catch((err) => {
    _devicePromise = null
    throw err
  })
  return _devicePromise
}

// Compile a shader module with proper error surfacing (line numbers + message).
async function compileShader(device: GPUDevice, wgsl: string): Promise<GPUShaderModule> {
  const module = device.createShaderModule({ code: wgsl })
  const info = await module.getCompilationInfo()
  const errors = info.messages.filter((m) => m.type === 'error')
  if (errors.length > 0) {
    const formatted = errors.map((e) => `WGSL line ${e.lineNum}:${e.linePos}: ${e.message}`).join('\n')
    throw new Error(`Shader compilation failed:\n${formatted}`)
  }
  return module
}

// --- Info ------------------------------------------------------------------

export const webgpuInfoTool = tool({
  name: 'webgpu_info',
  description: 'Report WebGPU support, adapter info, device limits, and feature support. Run this FIRST to confirm the browser/GPU can run shaders.',
  inputSchema: z.object({}),
  callback: async () => {
    try {
      const nav = navigator as any
      if (!nav.gpu) return JSON.stringify({
        status: 'error',
        supported: false,
        error: 'WebGPU not supported in this browser',
        hint: 'Chrome 113+, Edge 113+, Safari 18+. Firefox requires dom.webgpu.enabled=true.',
      })

      const adapter = await nav.gpu.requestAdapter({ powerPreference: 'high-performance' })
      if (!adapter) return JSON.stringify({
        status: 'error',
        supported: false,
        error: 'No adapter (GPU blocklisted or hardware unavailable)',
      })

      // Prefer the new adapter.info property; fall back to legacy requestAdapterInfo().
      let info: any = adapter.info
      if (!info && typeof adapter.requestAdapterInfo === 'function') {
        try { info = await adapter.requestAdapterInfo() } catch { info = {} }
      }

      const limits = adapter.limits ? Object.fromEntries(
        ['maxBufferSize','maxStorageBufferBindingSize','maxComputeWorkgroupSizeX','maxComputeWorkgroupSizeY','maxComputeWorkgroupSizeZ','maxComputeInvocationsPerWorkgroup','maxComputeWorkgroupsPerDimension','maxBindGroups','maxStorageBuffersPerShaderStage']
          .map((k) => [k, (adapter.limits as any)[k]])
          .filter(([, v]) => v !== undefined),
      ) : {}

      const features = Array.from((adapter.features as Set<string>) || [])

      return JSON.stringify({
        status: 'success',
        supported: true,
        adapter: {
          vendor: info?.vendor || 'unknown',
          architecture: info?.architecture || 'unknown',
          device: info?.device || 'unknown',
          description: info?.description || 'unknown',
        },
        features,
        has_timestamp_query: features.includes('timestamp-query'),
        limits,
      })
    } catch (err: unknown) {
      return JSON.stringify({ status: 'error', error: (err as Error).message })
    }
  },
})

// --- Simple single-buffer compute (iter 1: refined) ------------------------

export const webgpuComputeTool = tool({
  name: 'webgpu_compute',
  description:
    'Run a WGSL compute shader with a single storage buffer. Input is mutated in-place by the shader and returned. ' +
    'Shader must bind @group(0) @binding(0) var<storage, read_write> data: array<f32>. ' +
    'For multi-input / read-only inputs / uniforms, use webgpu_compute_multi instead.',
  inputSchema: z.object({
    wgsl: z.string().describe('WGSL compute shader source (must define the @compute @workgroup_size entry point)'),
    input: z.array(z.number()).describe('Float32 input data'),
    workgroups: z.array(z.number()).optional().describe('[x,y,z] dispatch dims. Default: [ceil(n/64),1,1]'),
    entryPoint: z.string().optional().describe('Entry point function name (default: "main")'),
  }),
  callback: async (input) => {
    try {
      const device = await getDevice()
      const data = new Float32Array(input.input)
      const byteLength = data.byteLength
      if (byteLength === 0) return JSON.stringify({ status: 'error', error: 'Empty input array' })

      const storage = device.createBuffer({
        size: byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
      })
      device.queue.writeBuffer(storage, 0, data.buffer)

      const staging = device.createBuffer({
        size: byteLength,
        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
      })

      const module = await compileShader(device, input.wgsl)
      const pipeline = device.createComputePipeline({
        layout: 'auto',
        compute: { module, entryPoint: input.entryPoint || 'main' },
      })
      const bindGroup = device.createBindGroup({
        layout: pipeline.getBindGroupLayout(0),
        entries: [{ binding: 0, resource: { buffer: storage } }],
      })

      const encoder = device.createCommandEncoder()
      const pass = encoder.beginComputePass()
      pass.setPipeline(pipeline)
      pass.setBindGroup(0, bindGroup)
      const wg = input.workgroups || [Math.ceil(data.length / 64), 1, 1]
      pass.dispatchWorkgroups(wg[0] || 1, wg[1] || 1, wg[2] || 1)
      pass.end()
      encoder.copyBufferToBuffer(storage, 0, staging, 0, byteLength)
      device.queue.submit([encoder.finish()])

      await staging.mapAsync(GPUMapMode.READ)
      const result = new Float32Array(staging.getMappedRange().slice(0))
      staging.unmap()
      storage.destroy()
      staging.destroy()

      return JSON.stringify({
        status: 'success',
        outputLength: result.length,
        output: Array.from(result),
        dispatched: wg,
      })
    } catch (err: unknown) {
      return JSON.stringify({ status: 'error', error: (err as Error).message })
    }
  },
})

// --- Multi-buffer compute (iter 2: real work) ------------------------------

export const webgpuComputeMultiTool = tool({
  name: 'webgpu_compute_multi',
  description:
    'Run a WGSL compute shader with multiple buffers: N read-only storage inputs + one read-write output + optional uniforms. ' +
    'Bind order: inputs at bindings [0..N-1], output at binding N, uniforms at binding N+1 (if provided). ' +
    'Example: two-input add → shader reads @binding(0) a, @binding(1) b, writes @binding(2) c.',
  inputSchema: z.object({
    wgsl: z.string().describe('WGSL source'),
    inputs: z.array(z.array(z.number())).describe('Array of Float32 input arrays (read-only in shader)'),
    output_length: z.number().describe('Size of the output buffer in f32 elements'),
    uniforms: z.array(z.number()).optional().describe('Optional uniform values (passed as array<f32> via uniform buffer)'),
    workgroups: z.array(z.number()).optional().describe('[x,y,z] dispatch. Default: [ceil(output_length/64),1,1]'),
    entryPoint: z.string().optional(),
    return_as: z.enum(['f32', 'u32', 'i32']).optional().describe('Interpret output bytes as f32 (default), u32, or i32'),
  }),
  callback: async (input) => {
    try {
      const device = await getDevice()

      const inputBuffers: GPUBuffer[] = input.inputs.map((arr, i) => {
        const data = new Float32Array(arr)
        const buf = device.createBuffer({
          size: Math.max(data.byteLength, 4),
          usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
          label: `input_${i}`,
        })
        if (data.byteLength > 0) device.queue.writeBuffer(buf, 0, data.buffer)
        return buf
      })

      const outBytes = input.output_length * 4
      const output = device.createBuffer({
        size: outBytes,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
        label: 'output',
      })
      const staging = device.createBuffer({
        size: outBytes,
        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
        label: 'staging',
      })

      let uniformBuf: GPUBuffer | null = null
      if (input.uniforms && input.uniforms.length > 0) {
        const u = new Float32Array(input.uniforms)
        // Uniform buffers must be 16-byte aligned.
        const padded = Math.ceil(u.byteLength / 16) * 16
        uniformBuf = device.createBuffer({
          size: padded,
          usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
          label: 'uniforms',
        })
        device.queue.writeBuffer(uniformBuf, 0, u.buffer)
      }

      const module = await compileShader(device, input.wgsl)
      const pipeline = device.createComputePipeline({
        layout: 'auto',
        compute: { module, entryPoint: input.entryPoint || 'main' },
      })

      const entries: GPUBindGroupEntry[] = inputBuffers.map((buf, i) => ({
        binding: i, resource: { buffer: buf },
      }))
      entries.push({ binding: inputBuffers.length, resource: { buffer: output } })
      if (uniformBuf) entries.push({ binding: inputBuffers.length + 1, resource: { buffer: uniformBuf } })

      const bindGroup = device.createBindGroup({
        layout: pipeline.getBindGroupLayout(0),
        entries,
      })

      const encoder = device.createCommandEncoder()
      const pass = encoder.beginComputePass()
      pass.setPipeline(pipeline)
      pass.setBindGroup(0, bindGroup)
      const wg = input.workgroups || [Math.ceil(input.output_length / 64), 1, 1]
      pass.dispatchWorkgroups(wg[0] || 1, wg[1] || 1, wg[2] || 1)
      pass.end()
      encoder.copyBufferToBuffer(output, 0, staging, 0, outBytes)
      device.queue.submit([encoder.finish()])

      await staging.mapAsync(GPUMapMode.READ)
      const bytes = staging.getMappedRange().slice(0)
      let result: number[]
      switch (input.return_as) {
        case 'u32': result = Array.from(new Uint32Array(bytes)); break
        case 'i32': result = Array.from(new Int32Array(bytes)); break
        default:    result = Array.from(new Float32Array(bytes)); break
      }
      staging.unmap()

      for (const b of inputBuffers) b.destroy()
      output.destroy()
      staging.destroy()
      uniformBuf?.destroy()

      return JSON.stringify({
        status: 'success',
        output: result,
        outputLength: result.length,
        inputsUsed: input.inputs.length,
        hasUniforms: !!uniformBuf,
        dispatched: wg,
        return_as: input.return_as || 'f32',
      })
    } catch (err: unknown) {
      return JSON.stringify({ status: 'error', error: (err as Error).message })
    }
  },
})

// --- Benchmark (iter 3) -----------------------------------------------------

export const webgpuBenchmarkTool = tool({
  name: 'webgpu_benchmark',
  description:
    'Run a single-buffer compute shader N times and measure throughput. Uses timestamp-query if available for nanosecond-precision GPU timing, otherwise wall-clock. Good for comparing kernels.',
  inputSchema: z.object({
    wgsl: z.string(),
    input: z.array(z.number()),
    iterations: z.number().optional().describe('Default 100'),
    workgroups: z.array(z.number()).optional(),
    entryPoint: z.string().optional(),
  }),
  callback: async (input) => {
    try {
      const device = await getDevice()
      const iterations = input.iterations ?? 100
      const data = new Float32Array(input.input)
      const byteLength = data.byteLength
      if (byteLength === 0) return JSON.stringify({ status: 'error', error: 'Empty input' })

      const storage = device.createBuffer({
        size: byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
      })
      device.queue.writeBuffer(storage, 0, data.buffer)

      const module = await compileShader(device, input.wgsl)
      const pipeline = device.createComputePipeline({
        layout: 'auto',
        compute: { module, entryPoint: input.entryPoint || 'main' },
      })
      const bindGroup = device.createBindGroup({
        layout: pipeline.getBindGroupLayout(0),
        entries: [{ binding: 0, resource: { buffer: storage } }],
      })
      const wg = input.workgroups || [Math.ceil(data.length / 64), 1, 1]

      // Warm-up — compile/cache shader, prime the queue
      {
        const enc = device.createCommandEncoder()
        const pass = enc.beginComputePass()
        pass.setPipeline(pipeline); pass.setBindGroup(0, bindGroup)
        pass.dispatchWorkgroups(wg[0], wg[1] || 1, wg[2] || 1)
        pass.end()
        device.queue.submit([enc.finish()])
        await device.queue.onSubmittedWorkDone()
      }

      const wallStart = performance.now()
      for (let i = 0; i < iterations; i++) {
        const enc = device.createCommandEncoder()
        const pass = enc.beginComputePass()
        pass.setPipeline(pipeline); pass.setBindGroup(0, bindGroup)
        pass.dispatchWorkgroups(wg[0], wg[1] || 1, wg[2] || 1)
        pass.end()
        device.queue.submit([enc.finish()])
      }
      await device.queue.onSubmittedWorkDone()
      const wallEnd = performance.now()

      storage.destroy()

      const totalMs = wallEnd - wallStart
      const avgMs = totalMs / iterations
      return JSON.stringify({
        status: 'success',
        iterations,
        total_ms: Number(totalMs.toFixed(3)),
        avg_ms_per_dispatch: Number(avgMs.toFixed(4)),
        dispatches_per_second: Math.round(1000 / avgMs),
        elements: data.length,
        dispatched: wg,
        note: 'Wall-clock; GPU timestamp queries require explicit enablement per browser.',
      })
    } catch (err: unknown) {
      return JSON.stringify({ status: 'error', error: (err as Error).message })
    }
  },
})

// --- Image kernel (iter 4) --------------------------------------------------

export const webgpuImageKernelTool = tool({
  name: 'webgpu_image_kernel',
  description:
    'Apply a WGSL compute shader to an RGBA image. Shader receives read-only input pixels, writes to output. ' +
    'Input + output are @group(0) @binding(0) and @binding(1) as array<vec4<f32>> (RGBA, 0..1 range). ' +
    'Uniforms at @binding(2): vec2<u32>(width, height). Input accepted as data URL / base64 / flat RGBA array. ' +
    'Returns a PNG data URL.',
  inputSchema: z.object({
    wgsl: z.string().describe('WGSL compute shader. Must write output[i] for each pixel.'),
    image_data_url: z.string().optional().describe('Data URL (data:image/…) or base64 PNG/JPEG'),
    width: z.number().optional().describe('Required if providing raw pixels'),
    height: z.number().optional().describe('Required if providing raw pixels'),
    pixels: z.array(z.number()).optional().describe('Raw RGBA uint8 array (length = 4*w*h), alternative to image_data_url'),
    entryPoint: z.string().optional(),
    workgroup_size: z.number().optional().describe('@workgroup_size declared in shader. Default 8 (8×8 = 64 threads/wg)'),
  }),
  callback: async (input) => {
    try {
      const device = await getDevice()

      // Load pixels from either source.
      let w = input.width ?? 0
      let h = input.height ?? 0
      let rgba: Uint8ClampedArray | null = null

      if (input.image_data_url) {
        if (typeof document === 'undefined') throw new Error('Image decoding requires DOM (document unavailable)')
        const img = new Image()
        img.crossOrigin = 'anonymous'
        await new Promise<void>((resolve, reject) => {
          img.onload = () => resolve()
          img.onerror = () => reject(new Error('Image load failed'))
          img.src = input.image_data_url!
        })
        w = img.naturalWidth; h = img.naturalHeight
        const canvas = document.createElement('canvas')
        canvas.width = w; canvas.height = h
        const ctx = canvas.getContext('2d')!
        ctx.drawImage(img, 0, 0)
        rgba = ctx.getImageData(0, 0, w, h).data
      } else if (input.pixels && w > 0 && h > 0) {
        if (input.pixels.length !== w * h * 4) throw new Error(`pixels.length (${input.pixels.length}) !== 4 * ${w} * ${h}`)
        rgba = new Uint8ClampedArray(input.pixels)
      } else {
        throw new Error('Provide either image_data_url or (pixels + width + height)')
      }

      const numPixels = w * h
      // Convert RGBA8 → vec4<f32> (0..1 range).
      const floatPixels = new Float32Array(numPixels * 4)
      for (let i = 0; i < floatPixels.length; i++) floatPixels[i] = rgba[i] / 255

      // Buffers
      const inBuf = device.createBuffer({
        size: floatPixels.byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
      })
      device.queue.writeBuffer(inBuf, 0, floatPixels.buffer)

      const outBuf = device.createBuffer({
        size: floatPixels.byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
      })
      const staging = device.createBuffer({
        size: floatPixels.byteLength,
        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
      })

      const uniforms = new Uint32Array([w, h, 0, 0]) // padded to 16 bytes
      const uniformBuf = device.createBuffer({
        size: uniforms.byteLength,
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
      })
      device.queue.writeBuffer(uniformBuf, 0, uniforms.buffer)

      const module = await compileShader(device, input.wgsl)
      const pipeline = device.createComputePipeline({
        layout: 'auto',
        compute: { module, entryPoint: input.entryPoint || 'main' },
      })
      const bindGroup = device.createBindGroup({
        layout: pipeline.getBindGroupLayout(0),
        entries: [
          { binding: 0, resource: { buffer: inBuf } },
          { binding: 1, resource: { buffer: outBuf } },
          { binding: 2, resource: { buffer: uniformBuf } },
        ],
      })

      const wgSize = input.workgroup_size ?? 8
      const dispatchX = Math.ceil(w / wgSize)
      const dispatchY = Math.ceil(h / wgSize)

      const encoder = device.createCommandEncoder()
      const pass = encoder.beginComputePass()
      pass.setPipeline(pipeline); pass.setBindGroup(0, bindGroup)
      pass.dispatchWorkgroups(dispatchX, dispatchY, 1)
      pass.end()
      encoder.copyBufferToBuffer(outBuf, 0, staging, 0, floatPixels.byteLength)
      device.queue.submit([encoder.finish()])

      await staging.mapAsync(GPUMapMode.READ)
      const outFloat = new Float32Array(staging.getMappedRange().slice(0))
      staging.unmap()

      // Float → RGBA8
      const outRGBA = new Uint8ClampedArray(outFloat.length)
      for (let i = 0; i < outFloat.length; i++) {
        outRGBA[i] = Math.max(0, Math.min(255, Math.round(outFloat[i] * 255)))
      }

      // Encode to PNG via canvas
      if (typeof document === 'undefined') throw new Error('Encoding requires DOM (document unavailable)')
      const canvas = document.createElement('canvas')
      canvas.width = w; canvas.height = h
      const ctx = canvas.getContext('2d')!
      const imgData = new ImageData(outRGBA, w, h)
      ctx.putImageData(imgData, 0, 0)
      const dataUrl = canvas.toDataURL('image/png')

      inBuf.destroy(); outBuf.destroy(); staging.destroy(); uniformBuf.destroy()

      return JSON.stringify({
        status: 'success',
        width: w,
        height: h,
        output_data_url: dataUrl,
        dispatched: [dispatchX, dispatchY, 1],
        workgroup_size: wgSize,
      })
    } catch (err: unknown) {
      return JSON.stringify({ status: 'error', error: (err as Error).message })
    }
  },
})

// --- Examples / reference library (iter 3) ---------------------------------

const EXAMPLES: Record<string, { tool: string; description: string; wgsl: string; usage: string }> = {
  double: {
    tool: 'webgpu_compute',
    description: 'Double every element in an array. Simplest possible compute shader.',
    wgsl: `@group(0) @binding(0) var<storage, read_write> data: array<f32>;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= arrayLength(&data)) { return; }
  data[i] = data[i] * 2.0;
}`,
    usage: 'webgpu_compute({ wgsl, input: [1,2,3,4,5] }) → [2,4,6,8,10]',
  },
  add: {
    tool: 'webgpu_compute_multi',
    description: 'Elementwise add of two arrays.',
    wgsl: `@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read_write> c: array<f32>;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= arrayLength(&c)) { return; }
  c[i] = a[i] + b[i];
}`,
    usage: 'webgpu_compute_multi({ wgsl, inputs: [[1,2,3],[10,20,30]], output_length: 3 }) → [11,22,33]',
  },
  scale: {
    tool: 'webgpu_compute_multi',
    description: 'Scale an array by a scalar passed via uniform.',
    wgsl: `@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read_write> b: array<f32>;
@group(0) @binding(2) var<uniform> params: vec4<f32>; // params.x = scale
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let i = gid.x;
  if (i >= arrayLength(&b)) { return; }
  b[i] = a[i] * params.x;
}`,
    usage: 'webgpu_compute_multi({ wgsl, inputs: [[1,2,3,4]], output_length: 4, uniforms: [3.0,0,0,0] }) → [3,6,9,12]',
  },
  matmul: {
    tool: 'webgpu_compute_multi',
    description: 'Matrix multiply C = A × B. A is M×K, B is K×N, C is M×N. Uniform = [M,K,N,0].',
    wgsl: `@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read_write> c: array<f32>;
@group(0) @binding(3) var<uniform> dims: vec4<u32>; // M, K, N, _
@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let M = dims.x; let K = dims.y; let N = dims.z;
  let row = gid.y; let col = gid.x;
  if (row >= M || col >= N) { return; }
  var sum = 0.0;
  for (var k: u32 = 0u; k < K; k = k + 1u) {
    sum = sum + a[row * K + k] * b[k * N + col];
  }
  c[row * N + col] = sum;
}`,
    usage: '2×3 × 3×2 → 2×2. inputs=[A_flat, B_flat], output_length=M*N, uniforms=[M,K,N,0] (as floats), workgroups=[ceil(N/8),ceil(M/8),1]',
  },
  grayscale: {
    tool: 'webgpu_image_kernel',
    description: 'Convert image to grayscale (luminance).',
    wgsl: `@group(0) @binding(0) var<storage, read> src: array<vec4<f32>>;
@group(0) @binding(1) var<storage, read_write> dst: array<vec4<f32>>;
@group(0) @binding(2) var<uniform> dims: vec4<u32>; // w, h, _, _
@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let w = dims.x; let h = dims.y;
  if (gid.x >= w || gid.y >= h) { return; }
  let i = gid.y * w + gid.x;
  let p = src[i];
  let g = 0.299 * p.r + 0.587 * p.g + 0.114 * p.b;
  dst[i] = vec4<f32>(g, g, g, p.a);
}`,
    usage: 'webgpu_image_kernel({ wgsl, image_data_url: "data:image/png..." }) → { output_data_url }',
  },
  invert: {
    tool: 'webgpu_image_kernel',
    description: 'Invert image colors.',
    wgsl: `@group(0) @binding(0) var<storage, read> src: array<vec4<f32>>;
@group(0) @binding(1) var<storage, read_write> dst: array<vec4<f32>>;
@group(0) @binding(2) var<uniform> dims: vec4<u32>;
@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  if (gid.x >= dims.x || gid.y >= dims.y) { return; }
  let i = gid.y * dims.x + gid.x;
  let p = src[i];
  dst[i] = vec4<f32>(1.0 - p.r, 1.0 - p.g, 1.0 - p.b, p.a);
}`,
    usage: 'webgpu_image_kernel({ wgsl, image_data_url })',
  },
  blur3x3: {
    tool: 'webgpu_image_kernel',
    description: '3×3 box blur.',
    wgsl: `@group(0) @binding(0) var<storage, read> src: array<vec4<f32>>;
@group(0) @binding(1) var<storage, read_write> dst: array<vec4<f32>>;
@group(0) @binding(2) var<uniform> dims: vec4<u32>;
@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let w = i32(dims.x); let h = i32(dims.y);
  let x = i32(gid.x); let y = i32(gid.y);
  if (x >= w || y >= h) { return; }
  var sum = vec4<f32>(0.0);
  var n = 0.0;
  for (var dy = -1; dy <= 1; dy = dy + 1) {
    for (var dx = -1; dx <= 1; dx = dx + 1) {
      let nx = clamp(x + dx, 0, w - 1);
      let ny = clamp(y + dy, 0, h - 1);
      sum = sum + src[ny * w + nx];
      n = n + 1.0;
    }
  }
  dst[y * w + x] = sum / n;
}`,
    usage: 'webgpu_image_kernel({ wgsl, image_data_url })',
  },
}

export const webgpuExamplesTool = tool({
  name: 'webgpu_examples',
  description:
    'Reference library of working WGSL shaders. Pass `name` to get one (double/add/scale/matmul/grayscale/invert/blur3x3), or no args to list all.',
  inputSchema: z.object({
    name: z.string().optional().describe('Example name: double, add, scale, matmul, grayscale, invert, blur3x3'),
  }),
  callback: async (input) => {
    if (!input.name) {
      return JSON.stringify({
        status: 'success',
        examples: Object.entries(EXAMPLES).map(([name, ex]) => ({
          name, tool: ex.tool, description: ex.description,
        })),
        hint: 'Call again with { name: "matmul" } (or any other) to get the full WGSL + usage.',
      })
    }
    const ex = EXAMPLES[input.name]
    if (!ex) return JSON.stringify({
      status: 'error',
      error: `Unknown example: ${input.name}`,
      available: Object.keys(EXAMPLES),
    })
    return JSON.stringify({ status: 'success', name: input.name, ...ex })
  },
})

export const WEBGPU_TOOLS = [
  webgpuInfoTool,
  webgpuComputeTool,
  webgpuComputeMultiTool,
  webgpuBenchmarkTool,
  webgpuImageKernelTool,
  webgpuExamplesTool,
]