import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import {
  CostModel,
  LIBRARY_FILES,
  METHODS,
  SCENARIO_PRESETS,
  buildModesBundle,
  buildPlanBundle,
  buildRunBundle,
  buildStrikeBundle,
  checkModesValue,
  cloneSpec,
  defaultSpec,
  methodById,
  modesKey,
  parseMethodPlan,
  parsePlateResult,
  validateSpec,
  withEdits,
  type BackendInfo,
  type MethodId,
  type MethodPlan,
  type PlainValue,
  type PlateResult,
  type ProblemSpec,
} from "@plate-dyn/core";
import {
  CodeView,
  EquationsView,
  ParameterPanel,
  RunStatus,
  useRunner,
  useSolve,
  type ParameterSection,
  type StagedStatus,
} from "@plate-dyn/components";
import { ResultsView } from "./views/ResultsView";
import { ComparisonView } from "./views/ComparisonView";
import { clearPersistentState, usePersistentState } from "./usePersistentState";

type RightTab = "equations" | "source" | "run";
type MiddleTab = "result" | "compare";

const RIGHT_TABS: { id: RightTab; label: string }[] = [
  { id: "equations", label: "Equations" },
  { id: "source", label: "Source" },
  { id: "run", label: "Run details" },
];

/** What is on the bench: the settings a first run needs. */
const PROBLEM_SECTIONS: ParameterSection[] = [
  "plate",
  "material",
  "suspension",
  "hammer",
  "strike",
];
/** Reasonable defaults, rarely the point. */
const ADVANCED_SECTIONS: ParameterSection[] = ["probe", "mic", "render"];
const ADVANCED_PREFIXES = ["probe", "mic", "render"];

const SUMMARIES: Record<string, string> = Object.fromEntries(
  LIBRARY_FILES.map((f) => [f.name, f.summary]),
);
SUMMARIES["main.m"] = "Generated from the controls: the params struct, the modes and the strike";

/** A kept result and what it came from. */
export interface KeptResult {
  result: PlateResult;
  spec: ProblemSpec;
  stamp: string;
}

const isSpec = (v: unknown): v is ProblemSpec => {
  const s = v as ProblemSpec | null;
  return (
    !!s &&
    typeof s === "object" &&
    !!s.plate &&
    !!s.material &&
    !!s.hole &&
    !!s.thread &&
    !!s.hammer &&
    !!s.strike &&
    !!s.probe &&
    !!s.mic &&
    !!s.render &&
    !!s.rr &&
    !!s.fem &&
    typeof s.render.frame_span === "number"
  );
};
const isRightTab = (v: unknown): v is RightTab => RIGHT_TABS.some((t) => t.id === v);
const isMiddleTab = (v: unknown): v is MiddleTab => v === "result" || v === "compare";
const isMethodId = (v: unknown): v is MethodId => v === "rr" || v === "fem";
const isBool = (v: unknown): v is boolean => typeof v === "boolean";
const isEdits = (v: unknown): v is Record<string, string> =>
  !!v && typeof v === "object" && Object.values(v as object).every((x) => typeof x === "string");

/** Edits to the library files (not main.m), which change the modes and so their cache key. */
function libraryEdits(edits: Record<string, string>): Record<string, string> {
  return Object.fromEntries(Object.entries(edits).filter(([k]) => k !== "main.m"));
}

export function App() {
  // Everything the user chooses is remembered in this browser; "Reset to defaults" forgets it all.
  const [spec, setSpec] = usePersistentState<ProblemSpec>("spec-v1", defaultSpec, isSpec);
  const [methodId, setMethodId] = usePersistentState<MethodId>("method", () => "rr", isMethodId);
  const [edits, setEdits] = usePersistentState<Record<string, string>>(
    "edits",
    () => ({}),
    isEdits,
  );
  const [rightTab, setRightTab] = usePersistentState<RightTab>(
    "rightTab",
    () => "equations",
    isRightTab,
  );
  const [middleTab, setMiddleTab] = usePersistentState<MiddleTab>(
    "middleTab",
    () => "result",
    isMiddleTab,
  );
  const [advancedOpen, setAdvancedOpen] = usePersistentState<boolean>(
    "advancedOpen",
    () => false,
    isBool,
  );
  const method = methodById(methodId);
  const runner = useRunner();
  const [backend, setBackend] = useState<BackendInfo | null>(null);
  useEffect(() => {
    void runner.backend.then(setBackend);
  }, [runner]);

  const issues = useMemo(() => validateSpec(spec), [spec]);
  const hasErrors = issues.some((i) => i.severity === "error");
  const mainEdited = "main.m" in edits;

  const scenarioId = useMemo(
    () => SCENARIO_PRESETS.find((s) => JSON.stringify(s.spec) === JSON.stringify(spec))?.id ?? "",
    [spec],
  );

  // The source on screen is what runs: main.m as one script, or its two stages.
  const displayBundle = useMemo(() => buildRunBundle(spec, method), [spec, method]);

  const plan = useSolve(parseMethodPlan);
  const cost = useRef(new CostModel());
  useEffect(() => {
    if (hasErrors) return;
    const t = window.setTimeout(() => {
      void plan.run({ bundle: withEdits(buildPlanBundle(spec, method), edits) });
    }, 300);
    return () => window.clearTimeout(t);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [spec, edits, method, hasErrors]);
  const methodPlan: MethodPlan | null =
    plan.last && plan.last.method === method.id ? plan.last : null;
  const planError = plan.state.status === "error" ? plan.state.run.error : null;

  // ---- the staged run: modes (kept), then the strike
  const [status, setStatus] = useState<StagedStatus>({ status: "idle" });
  const [results, setResults] = useState<Partial<Record<MethodId, KeptResult>>>({});
  const modesCache = useRef(new Map<string, PlainValue>());
  const runSeq = useRef(0);
  const running = status.status === "running";

  /** Identifies the inputs a result came from, so a result from other settings can say so. */
  const stamp = useMemo(
    () => JSON.stringify({ spec, edits, method: method.id }),
    [spec, edits, method.id],
  );
  const currentModesKey = modesKey(spec, method.id) + "|" + JSON.stringify(libraryEdits(edits));
  const modesKept = !mainEdited && modesCache.current.has(currentModesKey);

  const runMethod = useCallback(
    async (m: typeof method, s: ProblemSpec, ed: Record<string, string>, st: string) => {
      const seq = ++runSeq.current;
      const output: string[] = [];
      const startedAt = performance.now();
      const edited = "main.m" in ed;
      const say = (text: string) => {
        output.push(text);
        if (seq === runSeq.current)
          setStatus((prev) =>
            prev.status === "running" ? { ...prev, output: [...output] } : prev,
          );
      };
      const fail = (error: string, backend: BackendInfo, cancelled?: boolean) =>
        setStatus({
          status: "error",
          error,
          output,
          wallMs: performance.now() - startedAt,
          backend,
          cancelled,
        });
      let value: PlainValue;
      let modesMs: number | null = null;
      let strikeMs = 0;
      let backendInfo: BackendInfo;
      if (edited) {
        setStatus({ status: "running", stage: "script", startedAt, output: [] });
        const r = await runner.run({ bundle: withEdits(buildRunBundle(s, m), ed) }, say);
        if (seq !== runSeq.current) return;
        if (!r.ok) return fail(r.error, r.backend, r.cancelled);
        value = r.value;
        modesMs = r.wallMs;
        backendInfo = r.backend;
      } else {
        const key = modesKey(s, m.id) + "|" + JSON.stringify(libraryEdits(ed));
        let modes = modesCache.current.get(key);
        if (!modes) {
          setStatus({ status: "running", stage: "modes", startedAt, output: [] });
          const r = await runner.run({ bundle: withEdits(buildModesBundle(s, m), ed) }, say);
          if (seq !== runSeq.current) return;
          if (!r.ok) return fail(r.error, r.backend, r.cancelled);
          try {
            checkModesValue(r.value);
          } catch (e) {
            return fail(
              `the modes stage returned an unexpected shape: ${(e as Error).message}`,
              r.backend,
            );
          }
          modes = r.value;
          // keep a handful of the most recent, so switching between a few plates stays instant
          if (modesCache.current.size >= 6)
            modesCache.current.delete(modesCache.current.keys().next().value as string);
          modesCache.current.set(key, modes);
          modesMs = r.wallMs;
          if (methodPlan) cost.current.observe(m, "modes", methodPlan.work, r.wallMs);
        }
        setStatus({ status: "running", stage: "strike", startedAt, output: [...output] });
        const r2 = await runner.run(
          { bundle: withEdits(buildStrikeBundle(s, m), ed), inputs: { modes } },
          say,
        );
        if (seq !== runSeq.current) return;
        if (!r2.ok) return fail(r2.error, r2.backend, r2.cancelled);
        value = r2.value;
        strikeMs = r2.wallMs;
        backendInfo = r2.backend;
      }
      let result: PlateResult;
      try {
        result = parsePlateResult(value);
      } catch (e) {
        return fail(
          `the script ran, but its result has an unexpected shape: ${(e as Error).message}`,
          backendInfo,
        );
      }
      cost.current.observe(
        m,
        "strike",
        result.info.nsamples * result.info.nmodes,
        strikeMs || modesMs || 0,
      );
      setResults((prev) => ({ ...prev, [m.id]: { result, spec: s, stamp: st } }));
      setStatus({
        status: "done",
        output,
        wallMs: performance.now() - startedAt,
        modesMs,
        strikeMs,
        backend: backendInfo,
      });
    },
    [runner, methodPlan],
  );

  const run = useCallback(async () => {
    if (running || hasErrors) return;
    await runMethod(method, cloneSpec(spec), { ...edits }, stamp);
  }, [running, hasErrors, runMethod, method, spec, edits, stamp]);

  const runBoth = useCallback(async () => {
    if (running || hasErrors) return;
    const s = cloneSpec(spec);
    const ed = { ...edits };
    for (const m of METHODS) {
      const st = JSON.stringify({ spec: s, edits: ed, method: m.id });
      await runMethod(m, s, ed, st);
    }
    setMiddleTab("compare");
  }, [running, hasErrors, runMethod, spec, edits, setMiddleTab]);

  const cancel = () => {
    runSeq.current++;
    runner.cancel();
    setStatus({ status: "idle" });
  };

  const resetAll = () => {
    if (!window.confirm("Reset every setting and source edit to the defaults?")) return;
    clearPersistentState();
    setSpec(defaultSpec());
    setMethodId("rr");
    setEdits({});
    setRightTab("equations");
    setMiddleTab("result");
    setAdvancedOpen(false);
    setResults({});
    setStatus({ status: "idle" });
    modesCache.current.clear();
  };

  useEffect(() => {
    const onKey = (e: KeyboardEvent) => {
      if ((e.ctrlKey || e.metaKey) && e.key === "Enter") {
        e.preventDefault();
        void run();
      }
    };
    window.addEventListener("keydown", onKey);
    return () => window.removeEventListener("keydown", onKey);
  }, [run]);

  const current = results[method.id] ?? null;
  const stale = current !== null && !running && current.stamp !== stamp;
  const estimate = methodPlan
    ? {
        modesMs: modesKept ? null : cost.current.modesMs(method, methodPlan.work),
        strikeMs: cost.current.strikeMs(
          method,
          methodPlan.nsamples,
          Math.min(methodPlan.modes_cap, methodPlan.modes_expected + 6) + 3,
        ),
      }
    : undefined;
  const estimateMs = estimate ? (estimate.modesMs ?? 0) + estimate.strikeMs : undefined;
  const slow = estimateMs !== undefined && estimateMs > 5000;
  const advancedIssue = issues.some(
    (i) =>
      i.severity === "error" &&
      ADVANCED_PREFIXES.some((p) => i.path === p || i.path.startsWith(p + ".")),
  );
  const numericsSection: ParameterSection = method.id === "rr" ? "rr" : "fem";

  return (
    <div className="app">
      <header className="app-header">
        <div>
          <h1>Plate Dynamics Lab 2</h1>
          <div className="sub">
            A small metal plate on a thread, struck by a hammer: its modes, the contact, the sound,
            by two independent methods, solved by MATLAB code running in your browser.
          </div>
        </div>
        <div className="method-picker">
          <span className="picker-label">Scenario</span>
          <select
            className="pl-select"
            style={{ width: "auto" }}
            value={scenarioId}
            onChange={(e) => {
              const s = SCENARIO_PRESETS.find((x) => x.id === e.target.value);
              if (s) setSpec(cloneSpec(s.spec));
            }}
          >
            <option value="">custom settings</option>
            {SCENARIO_PRESETS.map((s) => (
              <option key={s.id} value={s.id}>
                {s.name}
              </option>
            ))}
          </select>
          <span className="picker-label">Method</span>
          <select
            className="pl-select"
            style={{ width: "auto" }}
            value={method.id}
            onChange={(e) => setMethodId(e.target.value as MethodId)}
            title={method.summary}
          >
            {METHODS.map((m) => (
              <option key={m.id} value={m.id}>
                {m.name}
              </option>
            ))}
          </select>
        </div>
      </header>
      <div className="app-body">
        <aside className="col">
          <div className="action-bar">
            {running ? (
              <button className="pl-btn pl-btn-danger run-btn" onClick={cancel}>
                Cancel
              </button>
            ) : (
              <button
                className={`pl-btn pl-btn-primary run-btn${stale ? " run-btn-stale" : ""}`}
                onClick={() => void run()}
                disabled={hasErrors}
                title="Ctrl+Enter or Cmd+Enter"
              >
                Run {method.shortName}
              </button>
            )}
            <div className="action-meta">
              {running && (
                <span className="pl-badge pl-badge-info">
                  {status.status === "running" && status.stage === "modes"
                    ? "solving the modes"
                    : status.status === "running" && status.stage === "strike"
                      ? "striking"
                      : "running"}
                </span>
              )}
              {!running && status.status === "done" && (
                <span className="pl-badge pl-badge-ok">finished</span>
              )}
              {!running && status.status === "error" && (
                <span className="pl-badge pl-badge-error">error</span>
              )}
              {slow && (
                <span className="pl-badge pl-badge-warn">
                  about {(estimateMs! / 1000).toFixed(0)} s
                </span>
              )}
              {!slow && estimateMs !== undefined && (
                <span className="pl-muted">
                  about{" "}
                  {estimateMs < 1000
                    ? `${estimateMs.toFixed(0)} ms`
                    : `${(estimateMs / 1000).toFixed(1)} s`}
                  {modesKept ? ", modes kept" : ""}
                </span>
              )}
              <span style={{ flex: 1 }} />
              <button
                className="pl-btn"
                onClick={() => void runBoth()}
                disabled={running || hasErrors}
                title="Run both methods on these settings, then open the comparison"
              >
                Run both
              </button>
              <button
                className="pl-btn"
                onClick={resetAll}
                title="Forget every setting and source edit saved in this browser"
              >
                Reset
              </button>
            </div>
            {hasErrors && !running && (
              <div className="pl-issue">
                Some settings are out of range; the groups that need attention are marked below.
              </div>
            )}
          </div>
          <div className="method-note">
            <strong>{method.name}</strong>
            <div className="pl-muted">{method.summary}</div>
          </div>
          {scenarioId && (
            <div className="method-note">
              <strong>{SCENARIO_PRESETS.find((s) => s.id === scenarioId)?.name}</strong>
              <div className="pl-muted">
                {SCENARIO_PRESETS.find((s) => s.id === scenarioId)?.note}
              </div>
            </div>
          )}
          {mainEdited && (
            <div className="banner">
              <span>
                The generated script was edited, so the controls no longer feed the run, which runs
                the script as one piece. Keep editing the source, or restore it to reconnect the
                controls.
              </span>
              <button
                className="pl-btn"
                onClick={() =>
                  setEdits((e) =>
                    Object.fromEntries(Object.entries(e).filter(([k]) => k !== "main.m")),
                  )
                }
              >
                Restore
              </button>
            </div>
          )}
          <div className={mainEdited ? "pl-disabled" : undefined}>
            <div className="panel-head">On the bench</div>
            <ParameterPanel
              spec={spec}
              onChange={setSpec}
              issues={issues}
              sections={PROBLEM_SECTIONS}
              open={["hammer", "strike"]}
            />
            {methodPlan && !hasErrors && <PlanSummary plan={methodPlan} />}
            {planError && (
              <div className="pl-note">The setup refuses this problem: {planError}</div>
            )}
            <div className="panel-head">Numerics</div>
            <ParameterPanel
              spec={spec}
              onChange={setSpec}
              issues={issues}
              sections={[numericsSection]}
              open={[numericsSection]}
            />
            {methodPlan && (
              <div className="pl-small pl-muted plan-line">
                {methodPlan.size_note}; about {methodPlan.modes_expected} modes below{" "}
                {(methodPlan.f_max / 1000).toFixed(1)} kHz; {methodPlan.nsamples.toLocaleString()}{" "}
                samples. {modesKept ? "The modes of this plate are kept from an earlier run." : ""}
              </div>
            )}
            <details
              className="advanced"
              open={advancedOpen || advancedIssue}
              onToggle={(e) => setAdvancedOpen((e.currentTarget as HTMLDetailsElement).open)}
            >
              <summary>
                More settings: probe, microphone, render
                {advancedIssue && (
                  <span className="pl-badge pl-badge-error" style={{ marginLeft: 6 }}>
                    check
                  </span>
                )}
              </summary>
              <ParameterPanel
                spec={spec}
                onChange={setSpec}
                issues={issues}
                sections={ADVANCED_SECTIONS}
                open={ADVANCED_SECTIONS}
              />
            </details>
          </div>
          <div className="pl-small pl-muted">
            Settings and source edits are saved in this browser. Ctrl+Enter runs.
          </div>
        </aside>

        <main className="col">
          <div className="tabs" role="tablist">
            <button
              className={middleTab === "result" ? "active" : ""}
              onClick={() => setMiddleTab("result")}
              role="tab"
              aria-selected={middleTab === "result"}
            >
              {method.shortName}
            </button>
            <button
              className={middleTab === "compare" ? "active" : ""}
              onClick={() => setMiddleTab("compare")}
              role="tab"
              aria-selected={middleTab === "compare"}
            >
              Comparison
            </button>
          </div>
          {middleTab === "result" && (
            <>
              {stale && (
                <div className="banner">
                  <span>
                    The settings changed after this run, so what you see below is from the earlier
                    ones.
                  </span>
                  <button className="pl-btn" onClick={() => void run()}>
                    Run again
                  </button>
                </div>
              )}
              {status.status === "error" && <div className="pl-error">{status.error}</div>}
              {current && (
                <ResultsView result={current.result} spec={current.spec} method={method} />
              )}
              {!current && <EmptyState status={status.status} method={method.shortName} />}
            </>
          )}
          {middleTab === "compare" && <ComparisonView results={results} current={spec} />}
        </main>

        <section className="col-right">
          <div className="tabs" role="tablist">
            {RIGHT_TABS.map((t) => (
              <button
                key={t.id}
                className={rightTab === t.id ? "active" : ""}
                onClick={() => setRightTab(t.id)}
                role="tab"
                aria-selected={rightTab === t.id}
              >
                {t.label}
              </button>
            ))}
          </div>
          {rightTab === "equations" && (
            <div className="tab-body">
              <div className="pl-small pl-muted tab-intro">
                What the lab solves and how the {method.shortName} method discretizes it. The code
                that does it is under Source.
              </div>
              {method.equationGroups.map((g) => (
                <div key={g.title} className="eq-group">
                  <h3>{g.title}</h3>
                  <EquationsView equations={g.equations} />
                </div>
              ))}
            </div>
          )}
          {rightTab === "source" && (
            <CodeView
              files={displayBundle.files}
              edits={edits}
              onEditsChange={setEdits}
              order={["main.m", ...method.readingOrder]}
              summaries={SUMMARIES}
              height="100%"
            />
          )}
          {rightTab === "run" && (
            <div className="tab-body" style={{ display: "flex", flexDirection: "column", gap: 12 }}>
              <RunStatus
                state={status}
                result={current?.result}
                method={method}
                backend={backend}
                estimate={estimate}
                onCancel={cancel}
              />
            </div>
          )}
        </section>
      </div>
    </div>
  );
}

/** The estimates the plan found, before any run. */
function PlanSummary({ plan }: { plan: MethodPlan }) {
  const lines: string[] = [];
  lines.push(`mass ${(plan.plate_mass * 1000).toPrecision(3)} g`);
  lines.push(`pendulum ${plan.pendulum_hz.toFixed(2)} Hz`);
  lines.push(
    `Hertz contact ${(plan.hertz_duration * 1e6).toPrecision(3)} µs, ${plan.hertz_peak_force.toPrecision(3)} N peak, ${(plan.hertz_impulse * 1e3).toPrecision(3)} mN s on a rigid target`,
  );
  lines.push(`coincidence ${(plan.critical_frequency / 1000).toPrecision(3)} kHz`);
  return (
    <div className="pl-small pl-muted plan-line">
      Statics: {lines.join("; ")}.
      {plan.notes.map((n) => (
        <div key={n} className="pl-issue pl-issue-warning">
          {n}
        </div>
      ))}
    </div>
  );
}

/** What the middle column says before anything has been solved. */
function EmptyState({ status, method }: { status: string; method: string }) {
  if (status === "running") return <div className="empty">Solving…</div>;
  if (status === "error")
    return (
      <div className="empty">
        The run failed. The message is above, and the full log is under Run details.
      </div>
    );
  return (
    <div className="empty">
      <h2>Nothing solved yet by {method}</h2>
      <ol>
        <li>
          Pick a scenario and a method at the top, or set the plate, the hammer and the strike on
          the left.
        </li>
        <li>
          Press <strong>Run</strong>, or Ctrl+Enter. <strong>Run both</strong> solves the other
          method too and opens the comparison.
        </li>
        <li>
          The plate on its thread, its modes, the contact, the fields, the probe traces, the
          spectrum and the sound appear here.
        </li>
      </ol>
      <p className="pl-muted">
        The equations are on the right; Source shows the MATLAB that runs, and you can edit it.
      </p>
    </div>
  );
}