import { useState, useRef, useCallback, useEffect } from "react";
import { Upload, Download, X, Wand2, ImageIcon, Loader2 } from "lucide-react";
import { cn } from "../lib/utils";

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const CHECKERBOARD: React.CSSProperties = {
  backgroundImage: `
    linear-gradient(45deg, #444 25%, transparent 25%),
    linear-gradient(-45deg, #444 25%, transparent 25%),
    linear-gradient(45deg, transparent 75%, #444 75%),
    linear-gradient(-45deg, transparent 75%, #444 75%)
  `,
  backgroundSize: "16px 16px",
  backgroundPosition: "0 0, 0 8px, 8px -8px, -8px 0px",
  backgroundColor: "#2a2a2a",
};

type BgPreset = { label: string; style: React.CSSProperties };

const BG_PRESETS: BgPreset[] = [
  { label: "Checker", style: CHECKERBOARD },
  { label: "White",   style: { backgroundColor: "#ffffff" } },
  { label: "Black",   style: { backgroundColor: "#000000" } },
  { label: "Red",     style: { backgroundColor: "#e53e3e" } },
  { label: "Green",   style: { backgroundColor: "#38a169" } },
  { label: "Blue",    style: { backgroundColor: "#3182ce" } },
];

// Max normalised distances for each keying space
const MAX_RGB = Math.sqrt(3); // ≈ 1.732
const MAX_HSV = 1.5;          // sqrt(0.5² × 4 + 1² + 1² × 0.25)

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

function rgbToHex(r: number, g: number, b: number): string {
  return `#${[r, g, b].map((v) => v.toString(16).padStart(2, "0")).join("")}`;
}

function hexToRgb(hex: string): [number, number, number] | null {
  const m = /^#?([0-9a-f]{2})([0-9a-f]{2})([0-9a-f]{2})$/i.exec(hex.trim());
  if (!m) return null;
  return [parseInt(m[1], 16), parseInt(m[2], 16), parseInt(m[3], 16)];
}

// ---------------------------------------------------------------------------
// WebGL — shaders
// ---------------------------------------------------------------------------

const VERT = `
attribute vec2 aPos;
varying vec2 vUV;
void main() {
  vUV = aPos * 0.5 + 0.5;
  gl_Position = vec4(aPos, 0.0, 1.0);
}`;

const KEY_FRAG = `
precision mediump float;
uniform sampler2D uImage;
uniform vec3  uKey;
uniform float uTol;
uniform float uFeather;
uniform float uSatMin;
uniform float uSpill;
uniform float uHueMode;
varying vec2 vUV;

vec3 rgbToHsv(vec3 c) {
  float cmax  = max(c.r, max(c.g, c.b));
  float cmin  = min(c.r, min(c.g, c.b));
  float delta = cmax - cmin;
  float h = 0.0;
  float s = (cmax < 0.0001) ? 0.0 : delta / cmax;
  if (delta > 0.0001) {
    if      (cmax == c.r) h = mod((c.g - c.b) / delta, 6.0);
    else if (cmax == c.g) h = (c.b - c.r) / delta + 2.0;
    else                  h = (c.r - c.g) / delta + 4.0;
    h /= 6.0;
  }
  return vec3(h, s, cmax);
}

void main() {
  vec4  c      = texture2D(uImage, vUV);
  vec3  hsv    = rgbToHsv(c.rgb);
  vec3  keyHsv = rgbToHsv(uKey);

  float d;
  if (uHueMode > 0.5) {
    float dh = abs(hsv.x - keyHsv.x);
    if (dh > 0.5) dh = 1.0 - dh;
    float ds = abs(hsv.y - keyHsv.y);
    float dv = abs(hsv.z - keyHsv.z);
    d = sqrt(dh * dh * 4.0 + ds * ds + dv * dv * 0.25);
  } else {
    d = distance(c.rgb, uKey);
  }

  float a = c.a;
  if (hsv.y >= uSatMin) {
    if (d <= uTol) {
      a = 0.0;
    } else if (uFeather > 0.0 && d <= uTol + uFeather) {
      a = (d - uTol) / uFeather * c.a;
    }
  }

  vec3 rgb = c.rgb;
  if (uSpill > 0.0) {
    float edgeDist  = max(0.0, d - uTol);
    float spillZone = max(uFeather + uTol * 0.5, 0.01);
    float spillFact = clamp(1.0 - edgeDist / spillZone, 0.0, 1.0) * uSpill;

    if (uKey.g >= uKey.r && uKey.g >= uKey.b) {
      float excess = rgb.g - max(rgb.r, rgb.b);
      if (excess > 0.0) rgb.g -= excess * spillFact;
    } else if (uKey.b >= uKey.r && uKey.b >= uKey.g) {
      float excess = rgb.b - max(rgb.r, rgb.g);
      if (excess > 0.0) rgb.b -= excess * spillFact;
    } else {
      float excess = rgb.r - max(rgb.g, rgb.b);
      if (excess > 0.0) rgb.r -= excess * spillFact;
    }
  }

  gl_FragColor = vec4(rgb, a);
}`;

const HALO_FRAG = `
precision mediump float;
uniform sampler2D uKeyed;
uniform float uHaloStr;
uniform float uHaloRadius;
uniform vec2  uTexelSize;
varying vec2 vUV;

void main() {
  vec4 c = texture2D(uKeyed, vUV);

  if (uHaloStr <= 0.0) {
    gl_FragColor = c;
    return;
  }

  vec2 r  = uTexelSize * uHaloRadius;
  vec2 rd = r * 0.7071;

  float minA = c.a;
  minA = min(minA, texture2D(uKeyed, vUV + vec2( r.x,   0.0 )).a);
  minA = min(minA, texture2D(uKeyed, vUV + vec2(-r.x,   0.0 )).a);
  minA = min(minA, texture2D(uKeyed, vUV + vec2(  0.0,  r.y )).a);
  minA = min(minA, texture2D(uKeyed, vUV + vec2(  0.0, -r.y )).a);
  minA = min(minA, texture2D(uKeyed, vUV + vec2( rd.x,  rd.y)).a);
  minA = min(minA, texture2D(uKeyed, vUV + vec2(-rd.x,  rd.y)).a);
  minA = min(minA, texture2D(uKeyed, vUV + vec2( rd.x, -rd.y)).a);
  minA = min(minA, texture2D(uKeyed, vUV + vec2(-rd.x, -rd.y)).a);

  gl_FragColor = vec4(c.rgb, mix(c.a, minA, uHaloStr));
}`;

// ---------------------------------------------------------------------------
// WebGL — types + init
// ---------------------------------------------------------------------------

type GlState = {
  gl:      WebGLRenderingContext;
  kProg:   WebGLProgram;
  srcTex:  WebGLTexture;
  fbo:     WebGLFramebuffer;
  fboTex:  WebGLTexture;
  uKey:     WebGLUniformLocation;
  uTol:     WebGLUniformLocation;
  uFeather: WebGLUniformLocation;
  uSatMin:  WebGLUniformLocation;
  uSpill:   WebGLUniformLocation;
  uHueMode: WebGLUniformLocation;
  hProg:      WebGLProgram;
  uHaloStr:    WebGLUniformLocation;
  uHaloRadius: WebGLUniformLocation;
  uTexelSize:  WebGLUniformLocation;
};

function compileShader(gl: WebGLRenderingContext, type: number, src: string): WebGLShader {
  const s = gl.createShader(type)!;
  gl.shaderSource(s, src);
  gl.compileShader(s);
  return s;
}

function makeProgram(gl: WebGLRenderingContext, vs: WebGLShader, fs: WebGLShader): WebGLProgram {
  const p = gl.createProgram()!;
  gl.attachShader(p, vs);
  gl.attachShader(p, fs);
  gl.bindAttribLocation(p, 0, "aPos");
  gl.linkProgram(p);
  return p;
}

function makeTexture(gl: WebGLRenderingContext): WebGLTexture {
  const t = gl.createTexture()!;
  gl.bindTexture(gl.TEXTURE_2D, t);
  gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
  gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
  gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR);
  gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.LINEAR);
  return t;
}

function initGl(canvas: HTMLCanvasElement): GlState | null {
  const gl = canvas.getContext("webgl", {
    premultipliedAlpha: false,
    preserveDrawingBuffer: true,
  }) as WebGLRenderingContext | null;
  if (!gl) return null;

  const vs    = compileShader(gl, gl.VERTEX_SHADER,   VERT);
  const kFrag = compileShader(gl, gl.FRAGMENT_SHADER, KEY_FRAG);
  const hFrag = compileShader(gl, gl.FRAGMENT_SHADER, HALO_FRAG);
  const kProg = makeProgram(gl, vs, kFrag);
  const hProg = makeProgram(gl, vs, hFrag);

  const buf = gl.createBuffer()!;
  gl.bindBuffer(gl.ARRAY_BUFFER, buf);
  gl.bufferData(gl.ARRAY_BUFFER, new Float32Array([-1,-1, 1,-1, -1,1, 1,1]), gl.STATIC_DRAW);
  gl.enableVertexAttribArray(0);
  gl.vertexAttribPointer(0, 2, gl.FLOAT, false, 0, 0);

  const srcTex = makeTexture(gl);
  const fboTex = makeTexture(gl);

  const fbo = gl.createFramebuffer()!;
  gl.bindFramebuffer(gl.FRAMEBUFFER, fbo);
  gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, fboTex, 0);
  gl.bindFramebuffer(gl.FRAMEBUFFER, null);

  return {
    gl, kProg, hProg, srcTex, fbo, fboTex,
    uKey:        gl.getUniformLocation(kProg, "uKey")!,
    uTol:        gl.getUniformLocation(kProg, "uTol")!,
    uFeather:    gl.getUniformLocation(kProg, "uFeather")!,
    uSatMin:     gl.getUniformLocation(kProg, "uSatMin")!,
    uSpill:      gl.getUniformLocation(kProg, "uSpill")!,
    uHueMode:    gl.getUniformLocation(kProg, "uHueMode")!,
    uHaloStr:    gl.getUniformLocation(hProg, "uHaloStr")!,
    uHaloRadius: gl.getUniformLocation(hProg, "uHaloRadius")!,
    uTexelSize:  gl.getUniformLocation(hProg, "uTexelSize")!,
  };
}

// ---------------------------------------------------------------------------
// AI remove tab
// ---------------------------------------------------------------------------

type AiStatus = "idle" | "loading" | "done" | "error";

function AiRemoveTab({ imageFile, imageSrc, onClear }: {
  imageFile: File;
  imageSrc: string;
  onClear: () => void;
}) {
  const [status, setStatus]       = useState<AiStatus>("idle");
  const [resultUrl, setResultUrl] = useState<string | null>(null);
  const [bgPreset, setBgPreset]   = useState(0);
  const [progress, setProgress]   = useState("");

  const handleRemove = async () => {
    setStatus("loading");
    setProgress("Loading AI model…");
    try {
      const { removeBackground } = await import("@imgly/background-removal");
      setProgress("Removing background…");
      const blob = await removeBackground(imageSrc, {
        progress: (_key: string, current: number, total: number) => {
          if (total > 0) {
            setProgress(`Downloading model… ${Math.round((current / total) * 100)}%`);
          }
        },
      });
      const url = URL.createObjectURL(blob);
      setResultUrl((prev) => { if (prev) URL.revokeObjectURL(prev); return url; });
      setStatus("done");
    } catch (err) {
      console.error(err);
      setStatus("error");
    }
  };

  const handleDownload = () => {
    if (!resultUrl) return;
    const a = document.createElement("a");
    a.download = imageFile.name.replace(/\.[^.]+$/, "") + "_nobg.png";
    a.href = resultUrl;
    a.click();
  };

  return (
    <div className="space-y-4">
      {/* Toolbar */}
      <div className="flex items-center gap-3 flex-wrap">
        <span className="text-xs text-text-tertiary font-mono truncate max-w-[200px]">
          {imageFile.name}
        </span>
        <button
          onClick={onClear}
          className={cn(
            "flex items-center gap-1.5 px-3 py-1.5 rounded-md border border-border text-xs text-text-tertiary",
            "hover:text-destructive hover:border-destructive transition-colors",
          )}
        >
          <X className="w-3.5 h-3.5" /> Clear
        </button>
        {status !== "done" ? (
          <button
            onClick={handleRemove}
            disabled={status === "loading"}
            className={cn(
              "flex items-center gap-1.5 px-4 py-1.5 rounded-md text-xs font-semibold transition-colors",
              status === "loading"
                ? "bg-raised border border-border text-text-tertiary cursor-not-allowed"
                : "bg-primary text-white hover:bg-primary/90",
            )}
          >
            {status === "loading" ? (
              <><Loader2 className="w-3.5 h-3.5 animate-spin" /> {progress}</>
            ) : (
              <><Wand2 className="w-3.5 h-3.5" /> Remove Background</>
            )}
          </button>
        ) : (
          <button
            onClick={handleDownload}
            className="flex items-center gap-1.5 px-4 py-1.5 rounded-md text-xs font-semibold bg-primary text-white hover:bg-primary/90 transition-colors"
          >
            <Download className="w-3.5 h-3.5" /> Download PNG
          </button>
        )}
      </div>

      {status === "error" && (
        <p className="text-xs text-destructive">
          Something went wrong. Check the console for details.
        </p>
      )}

      {/* Side-by-side */}
      <div className="grid grid-cols-2 gap-4">
        <div className="space-y-2">
          <p className="text-xs font-semibold text-text-tertiary uppercase tracking-wider">Original</p>
          <div className="bg-card border border-border rounded-lg overflow-hidden">
            <img src={imageSrc} className="w-full block" alt="Original" />
          </div>
        </div>

        <div className="space-y-2">
          <div className="flex items-center gap-2 flex-wrap">
            <p className="text-xs font-semibold text-text-tertiary uppercase tracking-wider flex-1">
              Result — transparent background
            </p>
            <div className="flex items-center gap-1">
              {BG_PRESETS.map((preset, i) => (
                <button
                  key={preset.label}
                  title={preset.label}
                  onClick={() => setBgPreset(i)}
                  className={cn(
                    "w-5 h-5 rounded border transition-all",
                    i === bgPreset
                      ? "ring-2 ring-primary ring-offset-1 ring-offset-background border-transparent"
                      : "border-border hover:border-primary/50",
                  )}
                  style={preset.style}
                />
              ))}
            </div>
          </div>
          <div className="bg-card border border-border rounded-lg overflow-hidden">
            {resultUrl ? (
              <div style={BG_PRESETS[bgPreset].style}>
                <img src={resultUrl} className="w-full block" alt="Result" />
              </div>
            ) : (
              <div className="aspect-square flex flex-col items-center justify-center text-text-tertiary gap-3 p-8">
                {status === "loading" ? (
                  <><Loader2 className="w-10 h-10 opacity-40 animate-spin" /><span className="text-xs opacity-40 text-center">{progress}</span></>
                ) : (
                  <><ImageIcon className="w-10 h-10 opacity-20" /><span className="text-xs opacity-40 text-center">Click Remove Background to process</span></>
                )}
              </div>
            )}
          </div>
        </div>
      </div>
    </div>
  );
}

// ---------------------------------------------------------------------------
// BackgroundRemoval component
// ---------------------------------------------------------------------------

type Mode = "chroma" | "ai";

export function BackgroundRemoval() {
  const [mode, setMode] = useState<Mode>("chroma");

  const [imageSrc,    setImageSrc]    = useState<string | null>(null);
  const [imageFile,   setImageFile]   = useState<File | null>(null);
  const [imageReady,  setImageReady]  = useState(false);
  const [keyColor,    setKeyColor]    = useState<[number, number, number] | null>(null);
  const [hexInput,    setHexInput]    = useState("");
  const [tolerance,   setTolerance]   = useState(30);
  const [feather,     setFeather]     = useState(10);
  const [satMin,      setSatMin]      = useState(0);
  const [spillStr,    setSpillStr]    = useState(0);
  const [hueMode,     setHueMode]     = useState(false);
  const [haloStr,     setHaloStr]     = useState(0);
  const [haloRadius,  setHaloRadius]  = useState(2);
  const [dragOver,    setDragOver]    = useState(false);
  const [bgPreset,    setBgPreset]    = useState(0);

  const sourceCanvasRef = useRef<HTMLCanvasElement>(null);
  const glCanvasRef     = useRef<HTMLCanvasElement>(null);
  const glRef           = useRef<GlState | null>(null);
  const fileInputRef    = useRef<HTMLInputElement>(null);

  useEffect(() => {
    if (keyColor) setHexInput(rgbToHex(...keyColor));
  }, [keyColor]);

  useEffect(() => {
    if (!imageSrc) return;
    setImageReady(false);
    setKeyColor(null);
    setHexInput("");

    const img = new Image();
    img.onload = () => {
      const src = sourceCanvasRef.current;
      if (!src) return;
      src.width  = img.naturalWidth;
      src.height = img.naturalHeight;
      src.getContext("2d")!.drawImage(img, 0, 0);

      const glCanvas = glCanvasRef.current;
      if (!glCanvas) return;
      glCanvas.width  = img.naturalWidth;
      glCanvas.height = img.naturalHeight;

      if (!glRef.current) {
        glRef.current = initGl(glCanvas);
        if (!glRef.current) {
          console.error("BackgroundRemoval: WebGL not supported");
          return;
        }
      }

      const { gl, srcTex, fboTex } = glRef.current;
      gl.bindTexture(gl.TEXTURE_2D, srcTex);
      gl.pixelStorei(gl.UNPACK_PREMULTIPLY_ALPHA_WEBGL, false);
      gl.pixelStorei(gl.UNPACK_COLORSPACE_CONVERSION_WEBGL, gl.NONE);
      gl.pixelStorei(gl.UNPACK_FLIP_Y_WEBGL, true);
      gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, img);
      gl.pixelStorei(gl.UNPACK_FLIP_Y_WEBGL, false);

      gl.bindTexture(gl.TEXTURE_2D, fboTex);
      gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, img.naturalWidth, img.naturalHeight, 0, gl.RGBA, gl.UNSIGNED_BYTE, null);
      gl.viewport(0, 0, img.naturalWidth, img.naturalHeight);
      setImageReady(true);
    };
    img.src = imageSrc;
  }, [imageSrc]);

  useEffect(() => {
    const state = glRef.current;
    if (!state || !imageReady || !keyColor) return;

    const w = glCanvasRef.current!.width;
    const h = glCanvasRef.current!.height;
    const { gl } = state;
    const MAX = hueMode ? MAX_HSV : MAX_RGB;

    gl.bindFramebuffer(gl.FRAMEBUFFER, state.fbo);
    gl.viewport(0, 0, w, h);
    gl.useProgram(state.kProg);
    gl.bindTexture(gl.TEXTURE_2D, state.srcTex);
    gl.uniform3f(state.uKey,     keyColor[0] / 255, keyColor[1] / 255, keyColor[2] / 255);
    gl.uniform1f(state.uTol,     (tolerance / 100) * MAX);
    gl.uniform1f(state.uFeather, (feather   / 100) * MAX);
    gl.uniform1f(state.uSatMin,  satMin   / 100);
    gl.uniform1f(state.uSpill,   spillStr / 100);
    gl.uniform1f(state.uHueMode, hueMode ? 1.0 : 0.0);
    gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);

    gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    gl.viewport(0, 0, w, h);
    gl.useProgram(state.hProg);
    gl.bindTexture(gl.TEXTURE_2D, state.fboTex);
    gl.uniform1f(state.uHaloStr,    haloStr / 100);
    gl.uniform1f(state.uHaloRadius, haloRadius);
    gl.uniform2f(state.uTexelSize,  1 / w, 1 / h);
    gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);
  }, [keyColor, tolerance, feather, satMin, spillStr, hueMode, haloStr, haloRadius, imageReady]);

  const loadFile = useCallback((file: File) => {
    if (imageSrc) URL.revokeObjectURL(imageSrc);
    setImageSrc(URL.createObjectURL(file));
    setImageFile(file);
  }, [imageSrc]);

  const handleDrop = (e: React.DragEvent) => {
    e.preventDefault();
    setDragOver(false);
    const file = e.dataTransfer.files[0];
    if (file?.type.startsWith("image/")) loadFile(file);
  };

  const handleCanvasClick = (e: React.MouseEvent<HTMLCanvasElement>) => {
    const canvas = sourceCanvasRef.current;
    if (!canvas || !imageReady) return;
    const rect   = canvas.getBoundingClientRect();
    const scaleX = canvas.width  / rect.width;
    const scaleY = canvas.height / rect.height;
    const x = Math.floor((e.clientX - rect.left) * scaleX);
    const y = Math.floor((e.clientY - rect.top)  * scaleY);
    const px = canvas.getContext("2d")!.getImageData(x, y, 1, 1).data;
    setKeyColor([px[0], px[1], px[2]]);
  };

  const handleHexInput = (v: string) => {
    setHexInput(v);
    const parsed = hexToRgb(v);
    if (parsed) setKeyColor(parsed);
  };

  const handleDownload = () => {
    const glCanvas = glCanvasRef.current;
    if (!glCanvas || !keyColor || !imageFile) return;
    glCanvas.toBlob((blob) => {
      if (!blob) return;
      const url = URL.createObjectURL(blob);
      const a   = document.createElement("a");
      a.download = imageFile.name.replace(/\.[^.]+$/, "") + "_transparent.png";
      a.href = url;
      a.click();
      URL.revokeObjectURL(url);
    }, "image/png");
  };

  const clearAll = useCallback(() => {
    if (imageSrc) URL.revokeObjectURL(imageSrc);
    setImageSrc(null);
    setImageFile(null);
    setImageReady(false);
    setKeyColor(null);
    setHexInput("");
    glRef.current = null;
    if (fileInputRef.current) fileInputRef.current.value = "";
  }, [imageSrc]);

  // ── Upload screen ──────────────────────────────────────────────────────────
  if (!imageSrc) {
    return (
      <div className="max-w-2xl mx-auto py-8 space-y-4">
        {/* Page header + mode toggle */}
        <div className="flex items-start justify-between gap-4 flex-wrap">
          <div>
            <h2 className="text-lg font-semibold text-foreground mb-1">
              Background Removal
            </h2>
            <p className="text-sm text-text-secondary">
              Upload an image then remove its background — either by selecting a
              key color (chroma key) or using the AI model.
            </p>
          </div>
          <ModeToggle mode={mode} onChange={setMode} />
        </div>

        <div
          onDrop={handleDrop}
          onDragOver={(e) => { e.preventDefault(); setDragOver(true); }}
          onDragLeave={() => setDragOver(false)}
          onClick={() => fileInputRef.current?.click()}
          className={cn(
            "border-2 border-dashed rounded-xl p-16 flex flex-col items-center gap-4 cursor-pointer transition-all select-none",
            dragOver
              ? "border-primary bg-primary/5"
              : "border-border hover:border-primary/50 hover:bg-primary/3",
          )}
        >
          <Upload className={cn("w-10 h-10 transition-colors", dragOver ? "text-primary" : "text-text-tertiary")} />
          <div className="text-center">
            <p className="text-sm font-medium text-text-secondary">
              {dragOver ? "Drop to upload" : "Drop an image here"}
            </p>
            <p className="text-xs text-text-tertiary mt-1">
              or click to browse · PNG, JPEG, WebP · max 15 MB
            </p>
          </div>
        </div>

        <input
          ref={fileInputRef}
          type="file"
          accept="image/png,image/jpeg,image/webp"
          onChange={(e) => { const f = e.target.files?.[0]; if (f) loadFile(f); }}
          className="hidden"
        />
      </div>
    );
  }

  // ── Editor screen ──────────────────────────────────────────────────────────
  const hasResult = imageReady && keyColor !== null;

  return (
    <div className="space-y-4 pb-8">
      {/* Page header + mode toggle */}
      <div className="flex items-center gap-3 flex-wrap">
        <h2 className="text-lg font-semibold text-foreground flex-1">
          Background Removal
        </h2>
        <ModeToggle mode={mode} onChange={setMode} />
      </div>

      {/* AI tab */}
      {mode === "ai" && (
        <AiRemoveTab imageFile={imageFile!} imageSrc={imageSrc} onClear={clearAll} />
      )}

      {/* Chroma key tab */}
      {mode === "chroma" && (
        <>
          {/* Toolbar */}
          <div className="flex items-center gap-3 flex-wrap">
            <span className="text-xs text-text-tertiary font-mono truncate max-w-[200px]">
              {imageFile?.name}
            </span>
            <button
              onClick={clearAll}
              className={cn(
                "flex items-center gap-1.5 px-3 py-1.5 rounded-md border border-border text-xs text-text-tertiary",
                "hover:text-destructive hover:border-destructive transition-colors",
              )}
            >
              <X className="w-3.5 h-3.5" /> Clear
            </button>
            <button
              onClick={handleDownload}
              disabled={!hasResult}
              className={cn(
                "flex items-center gap-1.5 px-4 py-1.5 rounded-md text-xs font-semibold transition-colors",
                hasResult
                  ? "bg-primary text-white hover:bg-primary/90"
                  : "bg-raised border border-border text-text-tertiary cursor-not-allowed",
              )}
            >
              <Download className="w-3.5 h-3.5" /> Download PNG
            </button>
          </div>

          {/* Controls */}
          <div className="bg-card border border-border rounded-xl p-4 space-y-4">
            {/* Row 1 — Key color + mode */}
            <div className="flex flex-wrap gap-6 items-center">
              <div className="space-y-1.5 shrink-0">
                <p className="text-xs font-medium text-text-secondary">Key Color</p>
                <div className="flex items-center gap-2">
                  <div
                    className="w-7 h-7 rounded-md border border-border shadow-inner shrink-0"
                    style={
                      keyColor
                        ? { backgroundColor: rgbToHex(...keyColor) }
                        : {
                            backgroundImage:
                              "linear-gradient(45deg,#444 25%,transparent 25%),linear-gradient(-45deg,#444 25%,transparent 25%),linear-gradient(45deg,transparent 75%,#444 75%),linear-gradient(-45deg,transparent 75%,#444 75%)",
                            backgroundSize: "8px 8px",
                            backgroundPosition: "0 0,0 4px,4px -4px,-4px 0",
                            backgroundColor: "#2a2a2a",
                          }
                    }
                  />
                  <input
                    type="text"
                    value={hexInput}
                    onChange={(e) => handleHexInput(e.target.value)}
                    placeholder="#rrggbb"
                    maxLength={7}
                    spellCheck={false}
                    className={cn(
                      "w-[76px] text-xs font-mono bg-transparent border rounded px-1.5 py-0.5",
                      "text-text-secondary focus:outline-none transition-colors",
                      hexInput && !hexToRgb(hexInput)
                        ? "border-destructive focus:border-destructive"
                        : "border-border focus:border-primary",
                    )}
                  />
                </div>
              </div>

              <div className="space-y-1.5 shrink-0">
                <p className="text-xs font-medium text-text-secondary">Keying Mode</p>
                <div className="flex rounded-md border border-border overflow-hidden text-xs font-medium">
                  <button
                    onClick={() => setHueMode(false)}
                    className={cn(
                      "px-3 py-1.5 transition-colors",
                      !hueMode ? "bg-primary text-white" : "text-text-secondary hover:text-foreground hover:bg-raised",
                    )}
                  >
                    RGB
                  </button>
                  <button
                    onClick={() => setHueMode(true)}
                    className={cn(
                      "px-3 py-1.5 transition-colors border-l border-border",
                      hueMode ? "bg-primary text-white" : "text-text-secondary hover:text-foreground hover:bg-raised",
                    )}
                  >
                    HSV
                  </button>
                </div>
              </div>

              {!keyColor && (
                <span className="text-xs text-text-tertiary flex items-center gap-1.5 ml-auto">
                  <Wand2 className="w-3.5 h-3.5" /> Click the image to pick a key color
                </span>
              )}
            </div>

            {/* Row 2 — Matte */}
            <div>
              <p className="text-[10px] font-semibold uppercase tracking-wider text-text-tertiary mb-2">Matte</p>
              <div className="flex flex-wrap gap-6">
                <div className="space-y-1.5 flex-1 min-w-[130px]">
                  <div className="flex justify-between">
                    <p className="text-xs font-medium text-text-secondary">Tolerance</p>
                    <span className="text-xs font-mono text-text-tertiary">{tolerance}%</span>
                  </div>
                  <input type="range" min="0" max="100" value={tolerance} onChange={(e) => setTolerance(Number(e.target.value))} className="w-full accent-primary" />
                </div>

                <div className="space-y-1.5 flex-1 min-w-[130px]">
                  <div className="flex justify-between">
                    <p className="text-xs font-medium text-text-secondary">Sat. Gate</p>
                    <span className="text-xs font-mono text-text-tertiary">{satMin}%</span>
                  </div>
                  <input type="range" min="0" max="100" value={satMin} onChange={(e) => setSatMin(Number(e.target.value))} className="w-full accent-primary" />
                  <p className="text-[10px] text-text-tertiary leading-tight">Skip pixels below this saturation — preserves neutral tones</p>
                </div>

                <div className="space-y-1.5 flex-1 min-w-[130px]">
                  <div className="flex justify-between">
                    <p className="text-xs font-medium text-text-secondary">Edge Feather</p>
                    <span className="text-xs font-mono text-text-tertiary">{feather}%</span>
                  </div>
                  <input type="range" min="0" max="50" value={feather} onChange={(e) => setFeather(Number(e.target.value))} className="w-full accent-primary" />
                </div>
              </div>
            </div>

            {/* Row 3 — Cleanup */}
            <div>
              <p className="text-[10px] font-semibold uppercase tracking-wider text-text-tertiary mb-2">Cleanup</p>
              <div className="flex flex-wrap gap-6">
                <div className="space-y-1.5 flex-1 min-w-[130px]">
                  <div className="flex justify-between">
                    <p className="text-xs font-medium text-text-secondary">Despill</p>
                    <span className="text-xs font-mono text-text-tertiary">{spillStr}%</span>
                  </div>
                  <input type="range" min="0" max="100" value={spillStr} onChange={(e) => setSpillStr(Number(e.target.value))} className="w-full accent-primary" />
                  <p className="text-[10px] text-text-tertiary leading-tight">Suppress key-color fringing on subject edges</p>
                </div>

                <div className="space-y-1.5 flex-1 min-w-[130px]">
                  <div className="flex justify-between">
                    <p className="text-xs font-medium text-text-secondary">Halo Remove</p>
                    <span className="text-xs font-mono text-text-tertiary">{haloStr}%</span>
                  </div>
                  <input type="range" min="0" max="100" value={haloStr} onChange={(e) => setHaloStr(Number(e.target.value))} className="w-full accent-primary" />
                  <p className="text-[10px] text-text-tertiary leading-tight">Erode the matte inward to eliminate bright rim pixels</p>
                </div>

                <div className="space-y-1.5 flex-1 min-w-[130px]">
                  <div className="flex justify-between">
                    <p className="text-xs font-medium text-text-secondary">Halo Radius</p>
                    <span className="text-xs font-mono text-text-tertiary">{haloRadius} px</span>
                  </div>
                  <input type="range" min="1" max="8" step="1" value={haloRadius} onChange={(e) => setHaloRadius(Number(e.target.value))} className="w-full accent-primary" disabled={haloStr === 0} />
                  <p className="text-[10px] text-text-tertiary leading-tight">How far to look for transparent neighbours</p>
                </div>
              </div>
            </div>
          </div>

          {/* Side-by-side view */}
          <div className="grid grid-cols-2 gap-4">
            <div className="space-y-2">
              <p className="text-xs font-semibold text-text-tertiary uppercase tracking-wider">
                Original — click to pick key color
              </p>
              <div className="bg-card border border-border rounded-lg overflow-hidden">
                <canvas ref={sourceCanvasRef} className="w-full cursor-crosshair block" onClick={handleCanvasClick} title="Click a pixel to set it as the chroma key color" />
              </div>
            </div>

            <div className="space-y-2">
              <div className="flex items-center gap-2 flex-wrap">
                <p className="text-xs font-semibold text-text-tertiary uppercase tracking-wider flex-1">
                  Result — transparent background
                </p>
                <div className="flex items-center gap-1">
                  {BG_PRESETS.map((preset, i) => (
                    <button
                      key={preset.label}
                      title={preset.label}
                      onClick={() => setBgPreset(i)}
                      className={cn(
                        "w-5 h-5 rounded border transition-all",
                        i === bgPreset
                          ? "ring-2 ring-primary ring-offset-1 ring-offset-background border-transparent"
                          : "border-border hover:border-primary/50",
                      )}
                      style={preset.style}
                    />
                  ))}
                </div>
              </div>
              <div className="bg-card border border-border rounded-lg overflow-hidden">
                <div style={BG_PRESETS[bgPreset].style} className={cn("w-full", !hasResult && "hidden")}>
                  <canvas ref={glCanvasRef} className="w-full block" />
                </div>
                {!hasResult && (
                  <div className="aspect-square flex flex-col items-center justify-center text-text-tertiary gap-3 p-8">
                    <ImageIcon className="w-10 h-10 opacity-20" />
                    <span className="text-xs opacity-40 text-center leading-relaxed">
                      {imageReady ? "Click the image on the left to pick a key color" : "Loading image…"}
                    </span>
                  </div>
                )}
              </div>
            </div>
          </div>
        </>
      )}
    </div>
  );
}

// ---------------------------------------------------------------------------
// ModeToggle
// ---------------------------------------------------------------------------

function ModeToggle({ mode, onChange }: { mode: Mode; onChange: (m: Mode) => void }) {
  return (
    <div className="flex rounded-md border border-border overflow-hidden text-xs font-medium shrink-0">
      <button
        onClick={() => onChange("chroma")}
        className={cn(
          "px-3 py-1.5 transition-colors",
          mode === "chroma" ? "bg-primary text-white" : "text-text-secondary hover:text-foreground hover:bg-raised",
        )}
      >
        Chroma Key
      </button>
      <button
        onClick={() => onChange("ai")}
        className={cn(
          "px-3 py-1.5 transition-colors border-l border-border",
          mode === "ai" ? "bg-primary text-white" : "text-text-secondary hover:text-foreground hover:bg-raised",
        )}
      >
        AI Remove
      </button>
    </div>
  );
}
