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";
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 [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 };
      w.postMessage(req);
    },
    [script, stopAudio, dtCheck, linearRef, twin]
  );

  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="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, dashpots", `${p.sys.si.length}, ${p.sys.nDashpots}`],
    ["method", p.sys.method],
    ["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])
Elements
  sys.springs   rows [i j k L0] or [i j k L0 k3]; force k*e + k3*e^3, e = length - L0
  sys.dashpots  rows [i j c]; force c*(v_j - v_i) on bead i
  sys.friction  scalar or n-by-1; force -c0*v
                [sys.friction, sys.dashpots] = bp_decay(pos, mass, sys.springs, [f1 T1; f2 T2])
                sets both from decay times (T60, s) at two frequencies (Hz)
  sys.pinned    bead indices (or logical mask) held fixed
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     compute linear transverse modes (default true; flat systems)
Bead indices are 1-based.`;