import { useCallback, useEffect, useRef, useState } from "react";
import { wavBlob } from "./sim/analysis";
import { BeadView } from "./ui/BeadView";
import { Editor } from "./ui/Editor";
import { DecayTimes, Energy, ModeEnergy, Separation, Spectrogram, Spectrum, Waveform } from "./ui/Plots";
import type { FromWorker, RunPayload, RunRequest } from "./worker";

const exampleSources = import.meta.glob("../matlab/examples/*.m", { query: "?raw", import: "default", eager: true }) as Record<string, string>;
const helperSources = import.meta.glob("../matlab/helpers/*.m", { query: "?raw", import: "default", eager: true }) as Record<string, string>;
const examples = Object.entries(exampleSources)
  .map(([path, source]) => ({ name: path.split("/").pop()!.replace(/\.m$/, ""), source }))
  .sort((a, b) => (a.name === "drum" ? -1 : b.name === "drum" ? 1 : a.name.localeCompare(b.name)));
const helpers = Object.entries(helperSources).map(([path, source]) => ({ name: path.split("/").pop()!, source }));

const STORE_KEY = "bead-percussion:script";

// a boolean remembered in this browser (falls back to the default if storage is unavailable)
function usePersistent(name: string, initial: boolean): [boolean, (v: boolean) => void] {
  const key = `bead-percussion:${name}`;
  const [value, setValue] = useState(() => {
    try {
      const s = localStorage.getItem(key);
      return s === null ? initial : s === "1";
    } catch {
      return initial;
    }
  });
  const set = (v: boolean) => {
    setValue(v);
    try {
      localStorage.setItem(key, v ? "1" : "0");
    } catch {
      /* storage unavailable */
    }
  };
  return [value, set];
}
const load = () => {
  try {
    return localStorage.getItem(STORE_KEY);
  } catch {
    return null;
  }
};
const save = (s: string) => {
  try {
    localStorage.setItem(STORE_KEY, s);
  } catch {
    /* storage unavailable */
  }
};

export default function App() {
  const [script, setScript] = useState(() => load() ?? examples[0].source);
  const [running, setRunning] = useState(false);
  const [stage, setStage] = useState("");
  const [progress, setProgress] = useState<number | null>(null);
  const [log, setLog] = useState("");
  const [error, setError] = useState<string | null>(null);
  const [payload, setPayload] = useState<RunPayload | null>(null);
  const [fmax, setFmax] = useState(10000);
  const [cursor, setCursor] = useState<number | null>(null);
  const [viewTime, setViewTime] = useState<number | null>(null);
  const [dtCheck, setDtCheck] = useState(false);
  const [linearRef, setLinearRef] = useState(false);
  const [twin, setTwin] = useState(false);
  const [modesOn, setModesOn] = usePersistent("modes", false);
  const [gpuOn, setGpuOn] = usePersistent("gpu", false);
  const gpuAvailable = typeof navigator !== "undefined" && "gpu" in navigator;
  const [playing, setPlaying] = useState<"main" | "linear" | "twin" | null>(null);
  const [runId, setRunId] = useState(0); // remounts panels with per-run state
  const worker = useRef<Worker | null>(null);
  const audioCtx = useRef<AudioContext | null>(null);
  const source = useRef<AudioBufferSourceNode | null>(null);
  const startedAt = useRef(0);

  useEffect(() => save(script), [script]);

  const stopAudio = useCallback(() => {
    source.current?.stop();
    source.current = null;
    setCursor(null);
    setPlaying(null);
  }, []);

  const run = useCallback(
    () => {
      worker.current?.terminate();
      stopAudio();
      const w = new Worker(new URL("./worker.ts", import.meta.url), { type: "module" });
      worker.current = w;
      setRunning(true);
      setError(null);
      setLog("");
      setProgress(null);
      setStage("Starting");
      w.onmessage = (ev: MessageEvent<FromWorker>) => {
        const m = ev.data;
        if (m.type === "log") setLog(l => l + m.text);
        else if (m.type === "stage") {
          setStage(m.text);
          setProgress(null);
        } else if (m.type === "progress") setProgress(m.t / m.duration);
        else if (m.type === "error") {
          setError(m.message);
          setRunning(false);
          w.terminate();
        } else if (m.type === "result") {
          setPayload(m.payload);
          setRunId(i => i + 1);
          const fsHalf = m.payload.result.sampleRate / 2;
          const top = m.payload.modeFreqs ? m.payload.modeFreqs[m.payload.modeFreqs.length - 1] * 1.15 : 10000;
          setFmax([1000, 2000, 5000, 10000, 20000].find(f => f >= top && f <= fsHalf) ?? Math.min(20000, fsHalf));
          setRunning(false);
          w.terminate();
        }
      };
      w.onerror = e => {
        setError(e.message || "Worker error");
        setRunning(false);
      };
      const req: RunRequest = { type: "run", script, dtCheck, linear: linearRef, twin, modes: modesOn, gpu: gpuOn && gpuAvailable };
      w.postMessage(req);
    },
    [script, stopAudio, dtCheck, linearRef, twin, modesOn, gpuOn, gpuAvailable]
  );

  const stop = () => {
    worker.current?.terminate();
    worker.current = null;
    setRunning(false);
    setStage("Stopped");
  };

  const play = (which: "main" | "linear" | "twin") => {
    if (!payload) return;
    stopAudio();
    const { sampleRate } = payload.result;
    const audio =
      which === "linear" && payload.linear
        ? payload.linear.audio
        : which === "twin" && payload.twin
          ? payload.twin.audio
          : payload.result.audio;
    const ctx = audioCtx.current ?? new AudioContext();
    audioCtx.current = ctx;
    // both signals share the main run's gain, so the A/B comparison keeps levels
    let peak = 0;
    for (const v of payload.result.audio) peak = Math.max(peak, Math.abs(v));
    const buf = ctx.createBuffer(1, audio.length, sampleRate);
    const ch = buf.getChannelData(0);
    const g = peak > 0 ? 0.9 / peak : 1;
    for (let i = 0; i < audio.length; i++) ch[i] = Math.max(-1, Math.min(1, g * audio[i]));
    const src = ctx.createBufferSource();
    src.buffer = buf;
    src.connect(ctx.destination);
    src.onended = () => {
      if (source.current === src) {
        source.current = null;
        setCursor(null);
        setPlaying(null);
      }
    };
    src.start();
    source.current = src;
    setPlaying(which);
    startedAt.current = ctx.currentTime;
    const T = audio.length / sampleRate;
    const tick = () => {
      if (source.current !== src) return;
      const t = ctx.currentTime - startedAt.current;
      setCursor(Math.min(t, T));
      if (t < T) requestAnimationFrame(tick);
    };
    requestAnimationFrame(tick);
  };

  const download = () => {
    if (!payload) return;
    const url = URL.createObjectURL(wavBlob(payload.result.audio, payload.result.sampleRate));
    const a = document.createElement("a");
    a.href = url;
    a.download = "bead-percussion.wav";
    a.click();
    setTimeout(() => URL.revokeObjectURL(url), 1000);
  };

  const r = payload?.result;
  return (
    <div className="app">
      <header>
        <h1>bead-percussion</h1>
        <span className="sub">beads and springs, set up in numbl, heard through one bead</span>
      </header>
      <div className="columns">
        <section className="left">
          <div className="controls">
            <select
              value=""
              onChange={e => {
                const ex = examples.find(x => x.name === e.target.value);
                if (ex && (script === ex.source || confirm(`Replace the script with the "${ex.name}" example?`))) setScript(ex.source);
              }}
            >
              <option value="" disabled>
                Load example…
              </option>
              {examples.map(ex => (
                <option key={ex.name} value={ex.name}>
                  {ex.name}
                </option>
              ))}
            </select>
            <span className="grow" />
            {running ? (
              <button onClick={stop}>Stop</button>
            ) : (
              <button className="primary" onClick={run} title="Ctrl/Cmd+Enter">
                Run
              </button>
            )}
          </div>
          <div className="controls options">
            <label title="Compute the linear modes (eigenvalues) and track the energy in each; can take seconds for large or curved nets">
              <input type="checkbox" checked={modesOn} onChange={e => setModesOn(e.target.checked)} /> linear modes
            </label>
            <label
              title={
                gpuAvailable
                  ? "Integrate on the GPU with WebGPU (single precision). Faster for large nets; the perturbed twin still runs on the CPU."
                  : "WebGPU is not available in this browser"
              }
            >
              <input type="checkbox" disabled={!gpuAvailable} checked={gpuOn && gpuAvailable} onChange={e => setGpuOn(e.target.checked)} /> WebGPU
            </label>
            <label title="Run again at dt/2 and compare the listener signals">
              <input type="checkbox" checked={dtCheck} onChange={e => setDtCheck(e.target.checked)} /> dt/2 check
            </label>
            <label title="Run again with every hit scaled by 1e-3 (then scaled back up): what the system would do if it were linear">
              <input type="checkbox" checked={linearRef} onChange={e => setLinearRef(e.target.checked)} /> small-strike reference
            </label>
            <label title="Run again with every hit changed by one part in 1e9 and track how fast the two runs separate (exponential growth means chaos)">
              <input type="checkbox" checked={twin} onChange={e => setTwin(e.target.checked)} /> perturbed twin
            </label>
          </div>
          <Editor value={script} onChange={setScript} onRun={() => !running && run()} />
          {running && (
            <div className="status">
              {stage}
              {progress !== null && ` ${(100 * progress).toFixed(0)}%`}
              <div className="bar">
                <div style={{ width: `${100 * (progress ?? 0)}%` }} />
              </div>
            </div>
          )}
          {error && <pre className="error">{error}</pre>}
          {log && <pre className="log">{log}</pre>}
          {payload && <Report p={payload} />}
          <details className="helpers">
            <summary>Helper functions (bp_*.m)</summary>
            {helpers.map(h => (
              <div key={h.name}>
                <h4>{h.name}</h4>
                <pre>{h.source}</pre>
              </div>
            ))}
          </details>
          <details className="helpers">
            <summary>sys fields</summary>
            <pre>{SYS_DOC}</pre>
          </details>
        </section>
        <section className="right">
          {!payload && <div className="empty">Run a script to hear and see the result.</div>}
          {payload && r && (
            <>
              <div className="controls">
                <button className="primary" onClick={() => (playing === "main" ? stopAudio() : play("main"))}>
                  {playing === "main" ? "Stop" : "Play"}
                </button>
                {payload.linear && (
                  <button onClick={() => (playing === "linear" ? stopAudio() : play("linear"))} title="The same strike scaled down 1000x, then scaled back up">
                    {playing === "linear" ? "Stop" : "Play small-strike reference"}
                  </button>
                )}
                {payload.twin && (
                  <button onClick={() => (playing === "twin" ? stopAudio() : play("twin"))} title="The run with the strike changed by one part in 1e9">
                    {playing === "twin" ? "Stop" : "Play twin"}
                  </button>
                )}
                <button onClick={download}>Download WAV</button>
                <span className="grow" />
                <label>
                  freq range{" "}
                  <select value={fmax} onChange={e => setFmax(+e.target.value)}>
                    {[1000, 2000, 5000, 10000, 20000]
                      .filter(f => f <= r.sampleRate / 2)
                      .map(f => (
                        <option key={f} value={f}>
                          {f / 1000} kHz
                        </option>
                      ))}
                  </select>
                </label>
              </div>
              <h3>Listener signal</h3>
              <Waveform p={payload} cursor={cursor ?? viewTime} />
              {payload.twin && (
                <>
                  <h3>Separation from the perturbed twin (chaos test)</h3>
                  <Separation p={payload} />
                </>
              )}
              <h3>Spectrogram</h3>
              <Spectrogram p={payload} fmax={fmax} />
              <h3>Spectrum, two windows (detuning shows as peaks shifting)</h3>
              <Spectrum key={`spectrum-${runId}`} p={payload} fmax={fmax} />
              <h3>Energy by linear mode (the cascade)</h3>
              <ModeEnergy key={`modes-${runId}`} p={payload} />
              {payload.modeFreqs && payload.modeDecay && (
                <>
                  <h3>Decay time by mode</h3>
                  <DecayTimes p={payload} />
                </>
              )}
              <h3>Energy</h3>
              <Energy p={payload} />
              <h3>Beads</h3>
              <BeadView key={`beads-${runId}`} p={payload} onTime={setViewTime} />
            </>
          )}
        </section>
      </div>
    </div>
  );
}

function Report({ p }: { p: RunPayload }) {
  const r = p.result;
  const nPinned = p.sys.pinned.reduce((a, b) => a + b, 0);
  const rows: [string, string][] = [
    ["beads", `${p.sys.n} (${nPinned} pinned)`],
    ["springs, bends, hinges, dashpots", `${p.sys.si.length}, ${p.sys.nBends}, ${p.sys.nHinges}, ${p.sys.nDashpots}`],
    ["method", `${p.sys.method} (${p.sys.engine})`],
    ["dt", `${r.dt.toExponential(3)} s (${r.substeps} per audio sample)`],
    ["highest linear frequency", `${(r.omegaMax / (2 * Math.PI)).toFixed(0)} Hz`],
    ["frequency error", r.freqErrors.map(e => `${e.cents.toFixed(e.cents < 0.1 ? 3 : 2)}¢ at ${e.hz / 1000} kHz`).join(", ")],
    ["energy drift after excitation", r.energyDrift.toExponential(2)],
    ["run time", `${(r.elapsedMs / 1000).toFixed(1)} s (${(r.elapsedMs / 1000 / p.sys.duration).toFixed(1)}x real time)`],
  ];
  if (p.modeFreqs) rows.push(["lowest modes", Array.from(p.modeFreqs.slice(0, 5)).map(f => f.toFixed(1)).join(", ") + " Hz"]);
  return (
    <div className="report">
      {r.aborted && <div className="error">{r.aborted}</div>}
      {[...p.notes, ...r.warnings].map((w, i) => (
        <div key={i} className="warn">
          {w}
        </div>
      ))}
      <table>
        <tbody>
          {rows.map(([k, v]) => (
            <tr key={k}>
              <td>{k}</td>
              <td>{v}</td>
            </tr>
          ))}
        </tbody>
      </table>
      {p.twin && (
        <div className="caption">
          Perturbed twin (hits changed by 1 part in 10^9):{" "}
          {p.twin.rate !== null
            ? `separation grows exponentially, e-folding time ${(1000 / p.twin.rate).toPrecision(2)} ms (rate ${p.twin.rate.toPrecision(3)} /s). Chaotic.`
            : `separation reached ${p.twin.sep.reduce((a, b) => Math.max(a, b), 0).toExponential(1)} of the motion's size; no exponential growth to 1e-2 within the run.`}
        </div>
      )}
      {p.linear && (
        <>
          <div className="caption">
            Listener signal vs the small-strike reference (how nonlinear the run is): relative RMS difference over [0, t]
          </div>
          <table>
            <tbody>
              {p.linear.windows.map(w => (
                <tr key={w.t}>
                  <td>t = {w.t} s</td>
                  <td>{w.rel.toExponential(2)}</td>
                </tr>
              ))}
            </tbody>
          </table>
        </>
      )}
      {p.check && (
        <>
          <div className="caption">
            Listener signal, dt vs dt/2 ({p.check.dt.toExponential(2)} s): relative RMS difference over [0, t]
          </div>
          <table>
            <tbody>
              {p.check.windows.map(w => (
                <tr key={w.t}>
                  <td>t = {w.t} s</td>
                  <td>{w.rel.toExponential(2)}</td>
                </tr>
              ))}
            </tbody>
          </table>
        </>
      )}
    </div>
  );
}

const SYS_DOC = `Required
  sys.pos       n-by-3 bead positions (m)
  sys.mass      scalar or n-by-1 (kg)
  sys.listen    struct: beads (index or indices, averaged), dir (default [0 0 1]),
                quantity: 'acceleration' (default), 'velocity' or 'displacement'
                (sound in air follows acceleration more closely; displacement
                favors low frequencies by 1/f^2)
Elements
  sys.springs   rows [i j k L0] or [i j k L0 k3]; force k*e + k3*e^3, e = length - L0
  sys.bends     rows [i j k kb] or [i j k kb cb], j the middle bead: bending,
                energy kb/2 |x_i - 2x_j + x_k|^2 (straight, evenly spaced rest);
                cb damps (v_i - 2v_j + v_k). Triplets come from bp_hexdisk, bp_chain
  sys.hinges    rows [i j k l kh] or [i j k l kh ch]: bending across the edge j-k
                shared by triangles (i,j,k) and (l,k,j); energy kh/2 (theta-theta0)^2,
                theta0 the angle at sys.pos (so any rest shape, e.g. a dome, is
                unstressed). Rows come from bp_hinges(tri)
  sys.dashpots  rows [i j c] or [i j c ca]; force c*(v_j - v_i) on bead i, plus
                ca*(d.(v_j - v_i))*d along the direction d between the beads
  sys.friction  scalar or n-by-1; force -c0*v
                (if neither friction nor dashpots is given: every mode falls 60 dB in 2.5 s)
                sys = bp_decay(sys, [f1 T1; f2 T2]) sets friction, dashpots and the
                damping columns of bends and hinges from decay times (T60, s) at
                two frequencies (Hz); call it after the elements are set
  sys.pinned    bead indices (or logical mask) held fixed; for a soft hold
                (cord, felt) use bp_anchor
Excitation
  sys.hits      struct or struct array, each with
                  beads, t (default 0), and either
                  velocity [vx vy vz]: added instantly at t, or
                  force [fx fy fz] + duration: raised-cosine pulse per bead
Solver (optional)
  sys.duration  seconds (default 2)
  sys.method    'verlet' (default) or 'yoshida4'
  sys.safety    fraction of the stability limit for automatic dt (default 0.5)
  sys.dt        fixed step (rounded to divide the audio sample interval)
  sys.sampleRate  audio rate (default 48000)
  sys.frameRate   energy/mode/animation records per second (default 500)
  sys.modes     true/false: compute linear modes, overriding the page's
                "linear modes" toggle (off by default)
Bead indices are 1-based.`;