import type { CoreNode } from '../../CoreNode.js';
import { calcFactoredRadiusArray } from '../../lib/utils.js';
import type { Vec4 } from '../../renderers/webgl/internal/ShaderUtils.js';
import type { WebGlShaderType } from '../../renderers/webgl/WebGlShaderNode.js';
import {
  RoundedWithBorderAndShadowTemplate,
  type RoundedWithBorderAndShadowProps,
} from '../templates/RoundedWithBorderAndShadowTemplate.js';

export const RoundedWithBorderAndShadow: WebGlShaderType<RoundedWithBorderAndShadowProps> =
  {
    props: RoundedWithBorderAndShadowTemplate.props,
    update(node: CoreNode) {
      const props = this.props!;
      this.uniformRGBA('u_borderColor', props['border-color']);
      this.uniform4fa('u_borderWidth', props['border-w'] as Vec4);
      this.uniform1f('u_borderGap', this.props!['border-gap'] as number);
      this.uniform1f('u_borderAlign', this.props!['border-align'] as number);

      this.uniformRGBA('u_shadowColor', props['shadow-color']);
      this.uniform4fa('u_shadow', props['shadow-projection']);

      this.uniform4fa(
        'u_radius',
        calcFactoredRadiusArray(props.radius as Vec4, node.w, node.h),
      );
    },
    vertex: `
    # ifdef GL_FRAGMENT_PRECISION_HIGH
    precision highp float;
    # else
    precision mediump float;
    # endif

    attribute vec2 a_position;
    attribute vec2 a_textureCoords;
    attribute vec4 a_color;
    attribute vec2 a_nodeCoords;

    uniform vec2 u_resolution;
    uniform float u_pixelRatio;
    uniform float u_rtt;
    uniform vec2 u_dimensions;

    uniform vec4 u_shadow;
    uniform vec4 u_radius;
    uniform vec4 u_borderWidth;
    uniform float u_borderGap;
    uniform float u_borderAlign;

    varying vec4 v_color;
    varying vec2 v_textureCoords;
    varying vec2 v_nodeCoords;

    varying vec2 v_innerSize;
    varying vec2 v_outerSize;
    varying vec2 v_outerBorderUv;
    varying vec2 v_innerBorderUv;
    varying vec4 v_innerBorderRadius;
    varying vec4 v_outerBorderRadius;
    varying vec2 v_halfDimensions;

    void main() {
      vec2 screenSpace = vec2(2.0 / u_resolution.x,  -2.0 / u_resolution.y);
      vec2 edge = clamp(a_nodeCoords * 2.0 - vec2(1.0), -1.0, 1.0);
      float borderZero = 1.0 - step(0.001, dot(abs(u_borderWidth), vec4(1.0)));

      vec2 edgeOffset = edge * ((u_shadow.w * 2.0)+ u_shadow.z) + u_shadow.xy;
      vec2 vertexPos = (a_position + edge + edgeOffset) * u_pixelRatio;

      v_innerSize = vec2(0.0);
      v_outerSize = vec2(0.0);

      // Defaults for the zero-border case: the fragment shader is branchless
      // and always evaluates the border SDFs, so every varying must carry a
      // finite value (an unwritten varying is undefined and can poison the
      // final mix() with NaN).
      v_outerBorderUv = vec2(0.0);
      v_innerBorderUv = vec2(0.0);
      v_outerBorderRadius = u_radius;
      v_innerBorderRadius = u_radius;

      if(borderZero == 0.0) {
        vec4 adjustedBorderWidth = u_borderWidth - 1.0 + clamp(u_borderWidth, -1.0, 1.0);

        float borderTop = adjustedBorderWidth.x;
        float borderRight = adjustedBorderWidth.y;
        float borderBottom = adjustedBorderWidth.z;
        float borderLeft = adjustedBorderWidth.w;

        vec2 borderSize = vec2(borderRight + borderLeft, borderTop + borderBottom);
        vec2 extraSize = borderSize * u_borderAlign;
        float gapLeft = step(0.001, borderLeft) * u_borderGap;
        float gapRight = step(0.001, borderRight) * u_borderGap;
        float gapTop = step(0.001, borderTop) * u_borderGap;
        float gapBottom = step(0.001, borderBottom) * u_borderGap;
        vec2 gapSize = vec2(gapLeft + gapRight, gapTop + gapBottom);

        v_outerSize = (u_dimensions + gapSize + extraSize) * 0.5;
        v_innerSize = v_outerSize - borderSize * 0.5;

        // Use sign() to avoid branching
        vec2 borderDiff = vec2(borderRight - borderLeft, borderBottom - borderTop);
        vec2 signDiff = sign(borderDiff);
        borderDiff = abs(borderDiff);

        vec2 gapDiff = vec2(gapRight - gapLeft, gapBottom - gapTop);
        vec2 signGapDiff = sign(gapDiff);
        gapDiff = abs(gapDiff);

        v_outerBorderUv = -signDiff * borderDiff * u_borderAlign * 0.5 - signGapDiff * gapDiff * 0.5;
        v_innerBorderUv = v_outerBorderUv + signDiff * borderDiff * 0.5;

        v_outerBorderRadius = vec4(
          max(0.0, u_radius.x + max(borderTop * u_borderAlign + u_borderGap, borderLeft * u_borderAlign + u_borderGap)),
          max(0.0, u_radius.y + max(borderTop * u_borderAlign + u_borderGap, borderRight * u_borderAlign + u_borderGap)),
          max(0.0, u_radius.z + max(borderBottom * u_borderAlign + u_borderGap, borderRight * u_borderAlign + u_borderGap)),
          max(0.0, u_radius.w + max(borderBottom * u_borderAlign + u_borderGap, borderLeft * u_borderAlign + u_borderGap))
        );

        v_innerBorderRadius = vec4(
          max(0.0, v_outerBorderRadius.x - max(borderTop, borderLeft)),
          max(0.0, v_outerBorderRadius.y - max(borderTop, borderRight)),
          max(0.0, v_outerBorderRadius.z - max(borderBottom, borderRight)),
          max(0.0, v_outerBorderRadius.w - max(borderBottom, borderLeft))
        );

        vec2 edgeOffsetExtra = step(u_dimensions * 0.5, v_outerSize) * edge * (extraSize + u_borderGap);
        edgeOffset += edgeOffsetExtra;

        vertexPos = (a_position + edge + edgeOffset) * u_pixelRatio;
      }

      gl_Position = vec4(vertexPos.x * screenSpace.x - 1.0, -sign(screenSpace.y) * (vertexPos.y * -abs(screenSpace.y)) + 1.0, 0.0, 1.0);

      v_halfDimensions = u_dimensions * 0.5;
      v_color = a_color;
      v_nodeCoords = a_nodeCoords + (screenSpace + edgeOffset) / (u_dimensions);
      v_textureCoords = a_textureCoords + (screenSpace + edgeOffset) / (u_dimensions);
    }
  `,
    fragment: `
    # ifdef GL_FRAGMENT_PRECISION_HIGH
    precision highp float;
    # else
    precision mediump float;
    # endif

    uniform vec2 u_resolution;
    uniform float u_pixelRatio;
    uniform float u_alpha;
    uniform vec2 u_dimensions;
    uniform sampler2D u_texture;
    uniform float u_rtt;

    uniform vec4 u_radius;
    uniform vec4 u_borderWidth;
    uniform vec4 u_borderColor;
    uniform vec4 u_shadowColor;
    uniform vec4 u_shadow;
    uniform float u_borderGap;

    varying vec4 v_color;
    varying vec2 v_textureCoords;
    varying vec2 v_nodeCoords;

    varying vec2 v_innerSize;
    varying vec2 v_outerSize;
    varying vec2 v_outerBorderUv;
    varying vec2 v_innerBorderUv;
    varying vec4 v_innerBorderRadius;
    varying vec4 v_outerBorderRadius;
    varying vec2 v_halfDimensions;

    // Branchless quadrant radius select (r: x TL, y TR, z BR, w BL) --
    // Mali 400-class fragment pipelines serialize any branch, ternaries included.
    float quadRadius(vec2 p, vec4 r) {
      vec2 stepVal = step(vec2(0.0), p);
      return mix(mix(r.x, r.y, stepVal.x), mix(r.w, r.z, stepVal.x), stepVal.y);
    }

    float roundedBox(vec2 p, vec2 s, vec4 r) {
      float rad = quadRadius(p, r);
      vec2 q = abs(p) - s + rad;
      return (min(max(q.x, q.y), 0.0) + length(max(q, 0.0))) - rad;
    }

    float shadowBox(vec2 p, vec2 s, vec4 r) {
      float rad = quadRadius(p, r);
      vec2 q = abs(p) - s + rad;
      float dist = min(max(q.x, q.y), 0.0) + length(max(q, 0.0)) - rad;
      return 1.0 - smoothstep(-u_shadow.w, u_shadow.w + u_shadow.z, dist);
    }

    void main() {
      vec4 color = texture2D(u_texture, v_textureCoords) * v_color;
      vec2 boxUv = v_nodeCoords.xy * u_dimensions - v_halfDimensions;
      float borderZero = 1.0 - step(0.001, dot(abs(u_borderWidth), vec4(1.0)));
      float edgeWidth = 1.0 / u_pixelRatio;

      float nodeDist = roundedBox(boxUv, v_halfDimensions - edgeWidth, u_radius);
      float nodeAlpha = 1.0 - smoothstep(-0.5 * edgeWidth, 0.5 * edgeWidth, nodeDist);

      // Shadow geometry select, branchless: when the border grows the quad
      // beyond the node (outer size exceeds half dimensions on either axis)
      // the shadow hugs the outer border box, otherwise the node box. This
      // replaces a varying-driven if/else (worst offender: diverges per
      // fragment) with a step()/mix() parameter select into ONE shadowBox call.
      // When borderZero, v_outerSize is vec2(0.0) so outerSel is 0 and the
      // node-box parameters win, matching the old zero-border path.
      float outerSel = max(
        step(v_halfDimensions.x, v_outerSize.x),
        step(v_halfDimensions.y, v_outerSize.y)
      );
      float shadowAlpha = shadowBox(
        boxUv + v_outerBorderUv * outerSel - u_shadow.xy,
        mix(v_halfDimensions, v_outerSize, outerSel) + u_shadow.w - edgeWidth,
        mix(u_radius, v_outerBorderRadius, outerSel) + u_shadow.z
      );

      float outerDist = roundedBox(boxUv + v_outerBorderUv, v_outerSize - edgeWidth, v_outerBorderRadius);
      float innerDist = roundedBox(boxUv + v_innerBorderUv, v_innerSize - edgeWidth, v_innerBorderRadius);
      float outerAlpha = 1.0 - smoothstep(-0.5 * edgeWidth, 0.5 * edgeWidth, outerDist);
      float innerAlpha = 1.0 - smoothstep(-0.5 * edgeWidth, 0.5 * edgeWidth, innerDist);

      float borderDist = max(-innerDist, outerDist);
      float borderAlpha = 1.0 - smoothstep(-0.5 * edgeWidth, 0.5 * edgeWidth, borderDist);

      // Branchless path select -- Mali 400 serializes uniform branches, so all
      // three composites (no gap / gap / zero border) are computed and
      // mix()-selected instead of if/else'd.
      vec4 shadowBase = u_shadowColor * shadowAlpha;

      vec4 resNoGap = mix(shadowBase, u_borderColor, outerAlpha * u_borderColor.a);
      resNoGap = mix(resNoGap, color, innerAlpha);

      vec4 resFill = mix(shadowBase, color, nodeAlpha);
      vec4 resGap = mix(resFill, u_borderColor, borderAlpha * u_borderColor.a);

      float hasGap = step(0.0001, abs(u_borderGap));
      vec4 resultColor = mix(resNoGap, resGap, hasGap);
      resultColor = mix(resultColor, resFill, borderZero);
      gl_FragColor = resultColor * u_alpha;
    }
  `,
  };