import { useEffect, useRef, useState, useSyncExternalStore } from "react";
import { usePythonStatus } from "@elvis/pyodide";

import { python } from "./python";
import { Timeline } from "./Timeline";
import type { Graph, Mode, RunSettings, State } from "./types";

// The example from the course slides, with sin(x1 - x2) added
const DEFAULT_FUNCTION = "exp(x1) + x1*x2 + sin(x1-x2)";
// After typing stops, wait this long before rerunning
const TYPING_DELAY_MS = 500;
const ANIMATED_BY_DEFAULT = !window.matchMedia("(prefers-reduced-motion: reduce)").matches;

function parseNumber(text: string): number | null {
  const value = Number(text);
  return text.trim() !== "" && Number.isFinite(value) ? value : null;
}

export default function useAutodiff() {
  const status = usePythonStatus(python);
  const [expression, setExpression] = useState(DEFAULT_FUNCTION);
  const [functionError, setFunctionError] = useState<string | null>(null);
  // The last function that parsed, kept while the box holds one that doesn't
  const [graph, setGraph] = useState<Graph | null>(null);
  // As typed, by name, so a value survives while its variable leaves the function
  const [inputTexts, setInputTexts] = useState<Record<string, string>>({});
  const [mode, setMode] = useState<Mode>("reverse");
  const [seedChoice, setSeedChoice] = useState<string | null>(null);
  const [runError, setRunError] = useState<string | null>(null);
  // Delay before the next rerun: TYPING_DELAY_MS after a keystroke, else 0
  const runDelay = useRef(0);
  const [animated, setAnimatedState] = useState(ANIMATED_BY_DEFAULT);
  const [timeline] = useState(() => new Timeline(ANIMATED_BY_DEFAULT));
  const snapshot = useSyncExternalStore(timeline.subscribe, timeline.getSnapshot);

  useEffect(() => timeline.pause, [timeline]);

  useEffect(() => {
    let current = true;
    const timer = setTimeout(async () => {
      const result = await python.call("handle_set_function", { expression });
      if (!current) {
        return;
      }
      setFunctionError(result.ok ? null : result.error.message);
      if (result.ok) {
        timeline.load(null, []);
        setGraph(result.data);
      }
    }, TYPING_DELAY_MS);
    return () => {
      current = false;
      clearTimeout(timer);
    };
  }, [expression, timeline]);

  const inputs = (graph?.inputs ?? []).map((name, i) => {
    const text = inputTexts[name] ?? String(i + 2);
    return { name, text, value: parseNumber(text) };
  });
  const seed = seedChoice !== null && graph?.inputs.includes(seedChoice) ? seedChoice : (graph?.inputs[0] ?? "");

  const settings: RunSettings | null =
    graph && inputs.every((input) => input.value !== null)
      ? {
          inputs: Object.fromEntries(inputs.map((input) => [input.name, input.value!])),
          ...(mode === "forward" ? { mode, seed } : { mode }),
        }
      : null;
  const settingsKey = settings && JSON.stringify(settings);

  useEffect(() => {
    if (!settingsKey) {
      return;
    }
    let current = true;
    const timer = setTimeout(async () => {
      const next: RunSettings = JSON.parse(settingsKey);
      const result = await python.call("handle_set_run", next);
      if (!current) {
        return;
      }
      // Every state up front, as the playhead can land on any of them
      const states: State[] = [];
      let error = result.ok ? null : result.error.message;
      for (let k = 0; result.ok && k <= result.data.steps && !error; k++) {
        const state = await python.call("handle_state", { iteration: k });
        if (!current) {
          return;
        }
        if (state.ok && state.data) {
          states.push(state.data);
        } else {
          error = state.ok ? "The run has no state to show." : state.error.message;
        }
      }
      setRunError(error);
      timeline.load(next, error ? [] : states);
    }, runDelay.current);
    runDelay.current = 0;
    return () => {
      current = false;
      clearTimeout(timer);
    };
  }, [graph, settingsKey, timeline]);

  function setInput(name: string, text: string) {
    runDelay.current = TYPING_DELAY_MS;
    setInputTexts((texts) => ({ ...texts, [name]: text }));
  }

  function setAnimated(on: boolean) {
    timeline.setAnimated(on);
    setAnimatedState(on);
  }

  return {
    status,
    expression,
    setExpression,
    functionError,
    graph,
    inputs,
    setInput,
    mode,
    setMode,
    seed,
    setSeed: setSeedChoice,
    run: snapshot.run,
    runError,
    animated,
    setAnimated,
    timeline,
    snapshot,
  };
}