import { useCallback, useRef, useState } from "react";
import compileStanProgram, {
  DEFAULT_WASI_SERVER_URL,
} from "./compileStanProgram";
import runRemoteChains from "./runRemoteChains";
import {
  ChainStatusMap,
  defaultSamplingOpts,
  MultiChainResult,
  SamplingOpts,
} from "./types";

const DEFAULT_STAN_PROGRAM = `data {
  int<lower=0> N;
  array[N] int<lower=0, upper=1> y;
}
parameters {
  real<lower=0, upper=1> theta;
}
model {
  theta ~ beta(1, 1);
  y ~ bernoulli(theta);
}
`;

const DEFAULT_DATA_JSON = `{
  "N": 10,
  "y": [0, 1, 0, 0, 0, 0, 0, 0, 0, 1]
}
`;

const mean = (xs: number[]) => xs.reduce((a, b) => a + b, 0) / xs.length;
const std = (xs: number[]) => {
  const m = mean(xs);
  return Math.sqrt(mean(xs.map((x) => (x - m) ** 2)));
};

function App() {
  const [stanProgram, setStanProgram] = useState(DEFAULT_STAN_PROGRAM);
  const [dataJson, setDataJson] = useState(DEFAULT_DATA_JSON);

  const [wasiServerUrl, setWasiServerUrl] = useState(
    () => localStorage.getItem("wasiServerUrl") ?? DEFAULT_WASI_SERVER_URL,
  );
  const [wasmExecUrl, setWasmExecUrl] = useState(
    () => localStorage.getItem("wasmExecUrl") ?? "https://wasm-exec.fly.dev",
  );
  const [wasmExecKey, setWasmExecKey] = useState(
    () => localStorage.getItem("wasmExecKey") ?? "",
  );

  const [compileStatus, setCompileStatus] = useState("");
  const [wasiWasmUrl, setWasiWasmUrl] = useState<string | undefined>(undefined);
  const [compiling, setCompiling] = useState(false);

  const [samplingOpts, setSamplingOpts] =
    useState<SamplingOpts>(defaultSamplingOpts);
  const [sampling, setSampling] = useState(false);
  const [chainStatus, setChainStatus] = useState<ChainStatusMap>({});
  const [samplingError, setSamplingError] = useState("");
  const [result, setResult] = useState<MultiChainResult | undefined>(undefined);
  const cancelRef = useRef<(() => void) | undefined>(undefined);

  const handleCompile = useCallback(async () => {
    setCompiling(true);
    setResult(undefined);
    setWasiWasmUrl(undefined);
    const { artifactUrl } = await compileStanProgram(
      wasiServerUrl,
      stanProgram,
      setCompileStatus,
    );
    setWasiWasmUrl(artifactUrl);
    setCompiling(false);
  }, [wasiServerUrl, stanProgram]);

  const handleSample = useCallback(async () => {
    if (!wasiWasmUrl) return;
    setSampling(true);
    setSamplingError("");
    setResult(undefined);
    setChainStatus({});

    try {
      // fetch the compiled WASI module from the compile server (browser
      // HTTP cache makes repeat samples cheap)
      const resp = await fetch(wasiWasmUrl);
      if (!resp.ok)
        throw new Error(`failed to fetch ${wasiWasmUrl}: ${resp.statusText}`);
      const moduleBytes = new Uint8Array(await resp.arrayBuffer());

      const run = runRemoteChains(
        wasmExecUrl,
        wasmExecKey,
        moduleBytes,
        dataJson,
        samplingOpts,
        (chainId, status) =>
          setChainStatus((prev) => ({ ...prev, [chainId]: status })),
      );
      cancelRef.current = run.cancel;
      setResult(await run.result);
    } catch (e) {
      setSamplingError(`${e}`);
    } finally {
      cancelRef.current = undefined;
      setSampling(false);
    }
  }, [wasiWasmUrl, wasmExecUrl, wasmExecKey, dataJson, samplingOpts]);

  const handleCancel = useCallback(() => {
    cancelRef.current?.();
    cancelRef.current = undefined;
    setSampling(false);
    setSamplingError("canceled");
  }, []);

  const setOpt = (key: keyof SamplingOpts, value: number) =>
    setSamplingOpts((prev) => ({ ...prev, [key]: value }));

  return (
    <div style={{ fontFamily: "sans-serif", margin: 20, maxWidth: 1100 }}>
      <h2>Stan Remote Sampling</h2>
      <p style={{ color: "#555", fontSize: 14 }}>
        The Stan program is compiled to a pure-WASI module by{" "}
        <code>stan-wasm-wasi</code>; sampling runs one chain per parallel job
        on <code>wasm-exec</code> workers.
      </p>

      <div style={{ display: "flex", gap: 20 }}>
        <div style={{ flex: 2 }}>
          <h3>Stan program</h3>
          <textarea
            value={stanProgram}
            onChange={(e) => setStanProgram(e.target.value)}
            rows={16}
            style={{ width: "100%", fontFamily: "monospace" }}
            spellCheck={false}
          />
        </div>
        <div style={{ flex: 1 }}>
          <h3>data.json</h3>
          <textarea
            value={dataJson}
            onChange={(e) => setDataJson(e.target.value)}
            rows={16}
            style={{ width: "100%", fontFamily: "monospace" }}
            spellCheck={false}
          />
        </div>
      </div>

      <h3>Compile</h3>
      <div style={{ display: "flex", gap: 10, alignItems: "center", flexWrap: "wrap" }}>
        <label>
          Compile server:{" "}
          <input
            value={wasiServerUrl}
            onChange={(e) => {
              setWasiServerUrl(e.target.value);
              localStorage.setItem("wasiServerUrl", e.target.value);
            }}
            style={{ width: 280 }}
          />
        </label>
        <button onClick={handleCompile} disabled={compiling || sampling}>
          Compile
        </button>
        <span>{compileStatus}</span>
      </div>

      <h3>Sample</h3>
      <div
        style={{
          display: "flex",
          gap: 15,
          alignItems: "center",
          flexWrap: "wrap",
          fontSize: 14,
          marginBottom: 10,
        }}
      >
        <label>
          wasm-exec server:{" "}
          <input
            value={wasmExecUrl}
            onChange={(e) => {
              setWasmExecUrl(e.target.value);
              localStorage.setItem("wasmExecUrl", e.target.value);
            }}
            style={{ width: 200 }}
          />
        </label>
        <label>
          Client key (ck_...):{" "}
          <input
            value={wasmExecKey}
            onChange={(e) => {
              setWasmExecKey(e.target.value);
              localStorage.setItem("wasmExecKey", e.target.value);
            }}
            style={{ width: 220 }}
          />
        </label>
      </div>

      <div style={{ display: "flex", gap: 15, alignItems: "center", flexWrap: "wrap" }}>
        <label>
          Chains (parallel jobs):{" "}
          <input
            type="number"
            min={1}
            max={16}
            value={samplingOpts.num_chains}
            onChange={(e) => setOpt("num_chains", parseInt(e.target.value) || 1)}
            style={{ width: 60 }}
          />
        </label>
        <label>
          Warmup:{" "}
          <input
            type="number"
            min={0}
            value={samplingOpts.num_warmup}
            onChange={(e) => setOpt("num_warmup", parseInt(e.target.value) || 0)}
            style={{ width: 80 }}
          />
        </label>
        <label>
          Samples:{" "}
          <input
            type="number"
            min={1}
            value={samplingOpts.num_samples}
            onChange={(e) => setOpt("num_samples", parseInt(e.target.value) || 1)}
            style={{ width: 80 }}
          />
        </label>
        <button
          onClick={handleSample}
          disabled={sampling || compiling || !wasiWasmUrl || !wasmExecKey}
        >
          Sample
        </button>
        {sampling && <button onClick={handleCancel}>Cancel</button>}
      </div>

      {Object.keys(chainStatus).length > 0 && (
        <div style={{ marginTop: 15 }}>
          {Array.from({ length: samplingOpts.num_chains }, (_, i) => i + 1).map(
            (chainId) => {
              const s = chainStatus[chainId];
              let text = "submitted / queued...";
              let pct = 0;
              if (s === "done") {
                text = "done";
                pct = 100;
              } else if (s === "error") {
                text = "error";
              } else if (s) {
                pct = Math.round((s.iteration / s.totalIterations) * 100);
                text = `${s.iteration} / ${s.totalIterations} ${s.warmup ? "(warmup)" : "(sampling)"}`;
              }
              return (
                <div
                  key={chainId}
                  style={{ display: "flex", alignItems: "center", gap: 10, marginBottom: 4 }}
                >
                  <span style={{ width: 70 }}>Chain {chainId}</span>
                  <div style={{ width: 300, height: 14, background: "#eee", border: "1px solid #ccc" }}>
                    <div
                      style={{
                        width: `${pct}%`,
                        height: "100%",
                        background: s === "error" ? "#c33" : "#4a90d9",
                      }}
                    />
                  </div>
                  <span style={{ fontSize: 13, color: "#555" }}>{text}</span>
                </div>
              );
            },
          )}
        </div>
      )}

      {samplingError && (
        <pre style={{ color: "#c33", whiteSpace: "pre-wrap" }}>{samplingError}</pre>
      )}

      {result && (
        <div style={{ marginTop: 15 }}>
          <h3>Results</h3>
          <p style={{ fontSize: 14 }}>
            {result.numChains} chains, {result.draws[0]?.length ?? 0} total draws,{" "}
            {result.computeTimeSec.toFixed(2)} sec
          </p>
          <table style={{ borderCollapse: "collapse" }}>
            <thead>
              <tr>
                {["parameter", "mean", "std"].map((h) => (
                  <th
                    key={h}
                    style={{ border: "1px solid #ccc", padding: "4px 12px", textAlign: "left" }}
                  >
                    {h}
                  </th>
                ))}
              </tr>
            </thead>
            <tbody>
              {result.paramNames.map((name, p) => (
                <tr key={name}>
                  <td style={{ border: "1px solid #ccc", padding: "4px 12px" }}>{name}</td>
                  <td style={{ border: "1px solid #ccc", padding: "4px 12px" }}>
                    {mean(result.draws[p]).toFixed(4)}
                  </td>
                  <td style={{ border: "1px solid #ccc", padding: "4px 12px" }}>
                    {std(result.draws[p]).toFixed(4)}
                  </td>
                </tr>
              ))}
            </tbody>
          </table>
          {result.consoleText.trim() && (
            <details style={{ marginTop: 10 }}>
              <summary>Console output</summary>
              <pre style={{ fontSize: 12, background: "#f6f6f6", padding: 10 }}>
                {result.consoleText}
              </pre>
            </details>
          )}
        </div>
      )}
    </div>
  );
}

export default App;