import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import {
  CostModel,
  FD_METHOD,
  LIBRARY_FILES,
  SCENARIO_PRESETS,
  buildGpuStages,
  buildPlanBundle,
  buildRunBundle,
  cloneSpec,
  defaultSpec,
  methodBlockers,
  parseMethodPlan,
  parsePlateResult,
  validateSpec,
  withEdits,
  type BackendInfo,
  type MethodPlan,
  type ProblemSpec,
} from "@plate-dyn/core";
import {
  CodeView,
  EquationsView,
  ParameterPanel,
  RunStatus,
  useRunner,
  useSolve,
  useWebGpuSupport,
  type ParameterSection,
} from "@plate-dyn/components";
import { ResultsView } from "./views/ResultsView";
import { clearPersistentState, usePersistentState } from "./usePersistentState";

type RightTab = "equations" | "source" | "run";

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

const IDLE = { status: "idle" } as const;
const method = FD_METHOD;

/** What is on the bench: the settings a first run needs. */
const PROBLEM_SECTIONS: ParameterSection[] = ["plate", "material", "mount", "excitation"];
/** Reasonable defaults, rarely the point. */
const ADVANCED_SECTIONS: ParameterSection[] = ["gravity", "damping", "probe", "listener", "output"];
const ADVANCED_PREFIXES = ["gravity", "damping", "probe", "listener", "output"];
const NUMERICS_SECTIONS: ParameterSection[] = ["fd"];

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 and the solver call";

const isSpec = (v: unknown): v is ProblemSpec => {
  const s = v as ProblemSpec | null;
  return (
    !!s &&
    typeof s === "object" &&
    !!s.plate &&
    !!s.fd &&
    !!s.mount?.type &&
    !!s.gravity &&
    !!s.listener?.model &&
    !!s.output
  );
};
const isRightTab = (v: unknown): v is RightTab => RIGHT_TABS.some((t) => t.id === v);
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");

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 [edits, setEdits] = usePersistentState<Record<string, string>>(
    "edits",
    () => ({}),
    isEdits,
  );
  const [rightTab, setRightTab] = usePersistentState<RightTab>(
    "rightTab",
    () => "equations",
    isRightTab,
  );
  const [advancedOpen, setAdvancedOpen] = usePersistentState<boolean>(
    "advancedOpen",
    () => false,
    isBool,
  );
  // The time loop runs as its WebGPU port by default: a run of seconds is hundreds of thousands of steps.
  const [fdGpu, setFdGpu] = usePersistentState<boolean>("fdGpu", () => true, isBool);
  const gpuSupport = useWebGpuSupport();
  const gpuOn = fdGpu && gpuSupport.status === "available";
  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 blockers = methodBlockers(method, spec);
  const mainEdited = "main.m" in edits;

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

  // Source bundles: what the code view shows is what runs. With the GPU option on, the script is
  // written as the three stages the host then runs; an edited main script runs as written, in numbl.
  const bundles = useMemo(() => {
    const single = buildRunBundle(spec, method, { gpu: gpuOn });
    const stages = gpuOn && !mainEdited ? (buildGpuStages(spec, method) ?? undefined) : undefined;
    return { display: single, run: withEdits(single, edits), gpu: stages };
  }, [spec, edits, gpuOn, mainEdited]);

  const solve = useSolve(parsePlateResult);
  const plan = useSolve(parseMethodPlan);
  const cost = useRef(new CostModel());

  // Keep the plan (step, cost, statics, pendulum frequencies) current
  useEffect(() => {
    if (hasErrors) return;
    const t = window.setTimeout(() => {
      const pb = buildPlanBundle(spec, method);
      if (pb) void plan.run({ bundle: withEdits(pb, edits) });
    }, 300);
    return () => window.clearTimeout(t);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [spec, edits, hasErrors]);
  const methodPlan: MethodPlan | null = plan.last;
  const planError = plan.state.status === "error" ? plan.state.run.error : null;

  const running = solve.state.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, gpuOn }), [spec, edits, gpuOn]);
  const [ranWith, setRanWith] = useState<string | null>(null);

  const run = useCallback(async () => {
    if (running || hasErrors || blockers.length) return;
    solve.reset();
    setRanWith(stamp);
    const r = await solve.run({ bundle: bundles.run, gpu: bundles.gpu });
    if (r?.ok) {
      const res = parsePlateResult(r.value);
      cost.current.observe(
        method,
        Number(res.info.extra.cell_steps),
        r.wallMs,
        r.compute?.loop === "webgpu",
      );
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [running, hasErrors, blockers.length, bundles, solve, stamp]);

  const cancel = () => solve.cancel();

  const resetAll = () => {
    if (!window.confirm("Reset every setting and source edit to the defaults?")) return;
    clearPersistentState();
    setSpec(defaultSpec());
    setEdits({});
    setRightTab("equations");
    setAdvancedOpen(false);
    setFdGpu(true);
    solve.reset();
    setRanWith(null);
  };

  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 currentState = ranWith !== null ? solve.state : IDLE;
  const currentResult = ranWith !== null ? solve.last : null;
  const stale = ranWith !== null && !running && ranWith !== stamp && currentResult !== null;
  const estimateMs = methodPlan
    ? cost.current.estimateMs(method, methodPlan.cell_steps, gpuOn)
    : 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 equationGroups = method.equationGroups.filter(
    (g) => !g.mounts || g.mounts.includes(spec.mount.type),
  );

  return (
    <div className="app">
      <header className="app-header">
        <div>
          <h1>Plate Dynamics Lab</h1>
          <div className="sub">
            A metal plate on a thread, on pads, in a vise or in the hand: how it swings, rings and
            sounds, solved by MATLAB code and a WebGPU time loop 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>
        </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 || blockers.length > 0}
                title="Ctrl+Enter or Cmd+Enter"
              >
                Run
              </button>
            )}
            <div className="action-meta">
              {running && <span className="pl-badge pl-badge-info">running</span>}
              {!running && currentState.status === "done" && (
                <span className="pl-badge pl-badge-ok">finished</span>
              )}
              {!running && currentState.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`}
                </span>
              )}
              <span style={{ flex: 1 }} />
              <button
                className="pl-btn"
                onClick={resetAll}
                title="Forget every setting and source edit saved in this browser"
              >
                Reset to defaults
              </button>
            </div>
            {hasErrors && !running && (
              <div className="pl-issue">
                Some settings are out of range; the groups that need attention are marked below.
              </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. 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>
          )}
          {blockers.map((b) => (
            <div className="pl-note" key={b}>
              {b}
            </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={["mount", "excitation"]}
            />
            {methodPlan && !hasErrors && <PlanSummary plan={methodPlan} spec={spec} />}
            {planError && (
              <div className="pl-note">The setup refuses this problem: {planError}</div>
            )}
            <div className="panel-head">Grid and time</div>
            <ParameterPanel
              spec={spec}
              onChange={setSpec}
              issues={issues}
              fdDtMax={methodPlan?.dt_max}
              sections={NUMERICS_SECTIONS}
              open={NUMERICS_SECTIONS}
              fdGpu={{ enabled: fdGpu, onChange: setFdGpu, support: gpuSupport }}
            />
            {methodPlan && (
              <div className="pl-small pl-muted plan-line">
                {methodPlan.nsteps.toLocaleString()} steps of {(methodPlan.dt * 1e6).toPrecision(3)}{" "}
                µs on {methodPlan.nx} × {methodPlan.ny} nodes (
                {methodPlan.nunknowns.toLocaleString()} unknowns),{" "}
                {methodPlan.cell_steps.toLocaleString()} cell-steps
                {gpuOn ? ", time loop on the GPU" : ", in numbl"}.
              </div>
            )}
            <details
              className="advanced"
              open={advancedOpen || advancedIssue}
              onToggle={(e) => setAdvancedOpen((e.currentTarget as HTMLDetailsElement).open)}
            >
              <summary>
                More settings: gravity, damping, probe, listener, output
                {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">
          {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>
          )}
          {currentState.status === "error" && (
            <div className="pl-error">{currentState.run.error}</div>
          )}
          {currentResult && <ResultsView result={currentResult} spec={spec} method={method} />}
          {!currentResult && <EmptyState status={currentState.status} />}
        </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, how the {spec.mount.type} mount enters, and how the scheme
                discretizes it. The code that does it is under Source.
              </div>
              {equationGroups.map((g) => (
                <div key={g.title} className="eq-group">
                  <h3>{g.title}</h3>
                  <EquationsView equations={g.equations} />
                </div>
              ))}
            </div>
          )}
          {rightTab === "source" && (
            <CodeView
              files={bundles.display.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={currentState}
                result={currentResult}
                method={method}
                backend={backend}
                estimateMs={estimateMs}
                onCancel={cancel}
              />
            </div>
          )}
        </section>
      </div>
    </div>
  );
}

/** The statics and the rigid-body modes the plan found, before any run. */
function PlanSummary({ plan, spec }: { plan: MethodPlan; spec: ProblemSpec }) {
  const m = spec.mount;
  const lines: string[] = [];
  lines.push(`mass ${(plan.mass * 1000).toFixed(0)} g`);
  if (plan.tension.length === 1) lines.push(`thread tension ${plan.tension[0].toFixed(2)} N`);
  if (plan.tension.length === 2)
    lines.push(`thread tensions ${plan.tension[0].toFixed(2)} and ${plan.tension[1].toFixed(2)} N`);
  if (plan.sag > 0) lines.push(`sag under gravity ${(plan.sag * 1000).toPrecision(3)} mm`);
  if (plan.rigid_hz.length > 0) {
    const hz = plan.rigid_hz.map(
      (h, k) =>
        `${Number.isFinite(h) ? (h < 0.05 ? `a ${(1 / Math.max(h, 1e-9) / 60).toFixed(0)} min period` : `${h.toFixed(2)} Hz`) : "unstable"} (${plan.rigid_mode_names[k] ?? ""})`,
    );
    lines.push(`rigid-body modes at ${hz.join(", ")}`);
  }
  if (m.type === "thread" || m.type === "threads")
    lines.push(
      `hangs turned ${(Math.atan2(plan.gx, -plan.gy) * (180 / Math.PI)).toFixed(1)}° in its plane`,
    );
  return (
    <div className="pl-small pl-muted plan-line">
      Statics: {lines.join("; ")}.
      {plan.notes
        .filter((n) => /unstable|slack|widened|no load|clear of/.test(n))
        .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 }: { status: 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</h2>
      <ol>
        <li>
          Pick a scenario at the top, or set the plate, material, mount and strike on the left.
        </li>
        <li>
          Press <strong>Run</strong>, or Ctrl+Enter.
        </li>
        <li>
          The plate in its mount, its displacement and stress, the swing and the support forces, the
          probe traces 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>
  );
}