import { LayerExtension, type Layer } from "@deck.gl/core";

const flowPathUniforms = {
  name: "flowPath",
  fs: `layout(std140) uniform flowPathUniforms { float phase; } flowPath;`,
  uniformTypes: { phase: "f32" }
} as const;

export class FlowPathExtension extends LayerExtension {
  static extensionName = "FlowPathExtension";
  static defaultProps = { flowPhase: { type: "number", value: 0 } };

  getShaders() {
    return {
      modules: [flowPathUniforms],
      inject: {
        "fs:DECKGL_FILTER_COLOR": `
          float flowStripe = fract((geometry.uv.y * 0.16) - flowPath.phase);
          float flowPulse = smoothstep(0.0, 0.18, flowStripe) * (1.0 - smoothstep(0.55, 0.82, flowStripe));
          color.a *= 0.22 + flowPulse * 0.78;
        `
      }
    };
  }

  draw(this: Layer) {
    const state = this.state as unknown as {
      model?: { shaderInputs?: { setProps: (props: { flowPath: { phase: number } }) => void } };
    };
    state.model?.shaderInputs?.setProps({
      flowPath: { phase: Number((this.props as Record<string, unknown>).flowPhase ?? 0) }
    });
  }
}