/* global React, ReactDOM, PRESETS, sample, sampleWithTemp, applyTemperature, tailProb, fmtPct, pickTailToken, TAIL_TOKENS */
const { useState, useEffect, useRef, useMemo, useCallback } = React;

/* =====================================================================
   SECTION 1 — THE FAIR DIE
   A small animated 6-sided die that "speaks" by emitting uniformly
   random tokens. The output looks like nonsense, which is the point.
   ===================================================================== */

function FairDieDemo() {
  const [tokens, setTokens] = useState([]);
  const [rolling, setRolling] = useState(false);
  const intervalRef = useRef(null);

  // Stop any running roll on unmount.
  useEffect(() => () => clearInterval(intervalRef.current), []);

  const roll = useCallback(() => {
    if (rolling) return;
    setTokens([]);
    setRolling(true);
    let i = 0;
    intervalRef.current = setInterval(() => {
      const t = TAIL_TOKENS[Math.floor(Math.random() * TAIL_TOKENS.length)];
      setTokens(prev => [...prev, t]);
      i++;
      if (i >= 8) {
        clearInterval(intervalRef.current);
        setRolling(false);
      }
    }, 180);
  }, [rolling]);

  return (
    <div className="fair-die-stage">
      <div>
        <span className="stamp">Standard</span>
        <h3>Six sides, equal weight</h3>
        <DieFace rolling={rolling} weighted={false} />
        <p style={{ marginTop: 24, fontSize: 14, color: "var(--ink-soft)" }}>
          Every token in the vocabulary has the same chance of being picked.
          The output is grammatically empty.
        </p>
        <button className="btn" onClick={roll} disabled={rolling} style={{ marginTop: 16 }}>
          {rolling ? "Rolling…" : "Roll the fair die"}
        </button>
        <div className="output-box" style={{ marginTop: 24 }}>
          <span style={{ color: "var(--muted)" }}>The capital of France is </span>
          {tokens.map((t, i) => <span key={i} className="token bad">{t}</span>)}
          {!tokens.length && <span style={{ color: "var(--muted)", fontStyle: "italic" }}>(press the button)</span>}
        </div>
      </div>
      <div>
        <span className="stamp weight">Loaded</span>
        <h3>100,000 sides, weighted</h3>
        <DieFace rolling={false} weighted={true} />
        <p style={{ marginTop: 24, fontSize: 14, color: "var(--ink-soft)" }}>
          Same mechanism — random sampling — but the weights are tuned
          by trillions of sentences. The die almost always lands on a
          fluent, plausible token.
        </p>
        <button className="btn ghost" disabled style={{ marginTop: 16 }}>
          See it in action below ↓
        </button>
        <div className="output-box" style={{ marginTop: 24 }}>
          <span style={{ color: "var(--muted)" }}>The capital of France is </span>
          <span className="token good">Paris</span>
        </div>
      </div>
    </div>
  );
}

// A small isometric die illustration. We don't try to render 100k faces
// for the loaded die — instead we show a die "leaning" toward one side,
// which reads instantly as a center-of-gravity shift.
function DieFace({ rolling, weighted }) {
  const ref = useRef(null);
  useEffect(() => {
    if (!rolling || !ref.current) return;
    let raf;
    const start = performance.now();
    const tick = (t) => {
      const dt = (t - start) / 1000;
      ref.current.style.transform = `rotateX(${dt * 320}deg) rotateY(${dt * 280}deg)`;
      raf = requestAnimationFrame(tick);
    };
    raf = requestAnimationFrame(tick);
    return () => {
      cancelAnimationFrame(raf);
      if (ref.current) ref.current.style.transform = "";
    };
  }, [rolling]);

  return (
    <div style={{ width: 140, height: 140, perspective: 800, margin: "8px auto 0" }}>
      <div
        ref={ref}
        style={{
          width: "100%",
          height: "100%",
          transformStyle: "preserve-3d",
          transform: weighted ? "rotateX(-22deg) rotateY(34deg)" : "rotateX(-22deg) rotateY(34deg)",
          transition: rolling ? "none" : "transform 0.4s ease",
          position: "relative",
        }}
      >
        {/* faces */}
        {["front","back","right","left","top","bottom"].map((face, i) => {
          const transforms = {
            front:  "translateZ(60px)",
            back:   "rotateY(180deg) translateZ(60px)",
            right:  "rotateY(90deg) translateZ(60px)",
            left:   "rotateY(-90deg) translateZ(60px)",
            top:    "rotateX(90deg) translateZ(60px)",
            bottom: "rotateX(-90deg) translateZ(60px)",
          };
          const isWeighted = weighted && face === "front";
          return (
            <div key={face} style={{
              position: "absolute",
              width: 120,
              height: 120,
              left: 10, top: 10,
              background: isWeighted ? "var(--vermilion)" : "#FAF6EC",
              border: "1.5px solid var(--ink)",
              transform: transforms[face],
              display: "flex",
              alignItems: "center",
              justifyContent: "center",
              boxShadow: "inset 0 0 0 0.5px rgba(0,0,0,0.1)",
              color: isWeighted ? "var(--paper)" : "var(--ink)",
              fontFamily: "var(--serif-display)",
              fontSize: face === "front" && weighted ? 18 : 36,
              fontStyle: face === "front" && weighted ? "italic" : "normal",
              textAlign: "center",
              padding: 8,
              lineHeight: 1,
            }}>
              {!weighted && (
                <DiePips n={i + 1} />
              )}
              {weighted && face === "front" && "Paris"}
              {weighted && face !== "front" && (
                <span style={{ fontSize: 10, fontFamily: "var(--mono)", color: "var(--muted)", letterSpacing: "0.1em", textTransform: "uppercase" }}>
                  +99,999 sides
                </span>
              )}
            </div>
          );
        })}
        {weighted && (
          // The "weight" shadow — a hint that this die's center of mass is off.
          <div style={{
            position: "absolute",
            width: 28, height: 28,
            left: 56, top: 56,
            transform: "translateZ(58px)",
            borderRadius: "50%",
            background: "radial-gradient(circle, var(--vermilion-deep) 30%, transparent 70%)",
            opacity: 0.5,
            pointerEvents: "none",
          }} />
        )}
      </div>
    </div>
  );
}

function DiePips({ n }) {
  // Standard d6 pip layout.
  const dot = (cx, cy) => <circle key={`${cx}-${cy}`} cx={cx} cy={cy} r="6" fill="var(--ink)" />;
  const positions = {
    1: [[50,50]],
    2: [[25,25],[75,75]],
    3: [[25,25],[50,50],[75,75]],
    4: [[25,25],[25,75],[75,25],[75,75]],
    5: [[25,25],[25,75],[50,50],[75,25],[75,75]],
    6: [[25,25],[25,50],[25,75],[75,25],[75,50],[75,75]],
  };
  return (
    <svg viewBox="0 0 100 100" width="100%" height="100%">
      {positions[n].map(([x,y]) => dot(x,y))}
    </svg>
  );
}

/* =====================================================================
   SECTION 2 — THE LOADED DIE
   The hero: a probability strip you can actually roll, with a tally
   of outcomes building up over time. Plus a "your own prompt" lab
   that asks Claude to estimate a distribution.
   ===================================================================== */

function LoadedDieDemo() {
  const [presetId, setPresetId] = useState("paris");
  const preset = PRESETS.find(p => p.id === presetId);
  const [needleLeft, setNeedleLeft] = useState(0);
  const [rolling, setRolling] = useState(false);
  const [lastHit, setLastHit] = useState(null); // index into dist, or -1 for tail
  const [tally, setTally] = useState({}); // { tokenOrIndex: count }
  const [totalRolls, setTotalRolls] = useState(0);
  const tickerRef = useRef(null);

  // Reset tally when preset changes.
  useEffect(() => {
    setTally({});
    setTotalRolls(0);
    setLastHit(null);
    setNeedleLeft(0);
  }, [presetId]);

  const segments = useMemo(() => {
    // Cumulative left% for each segment, plus a "tail" segment at the end.
    let acc = 0;
    const named = preset.dist.map(d => {
      const seg = { ...d, left: acc, width: d.p * 100 };
      acc += d.p * 100;
      return seg;
    });
    const tailW = Math.max(0, 100 - acc);
    return { named, tail: { left: acc, width: tailW, p: tailW / 100 } };
  }, [preset]);

  // Center of each segment (in % from left), used as targets when rolling.
  const targetFor = (idx) => {
    if (idx === -1) return segments.tail.left + segments.tail.width / 2;
    const s = segments.named[idx];
    return s.left + s.width / 2;
  };

  const roll = useCallback(() => {
    if (rolling) return;
    setRolling(true);
    setLastHit(null);
    const target = sample(preset.dist);
    // "Bouncing" pre-roll — needle ricochets across the strip a few times
    // before settling. Easing makes the settle feel weighty.
    const steps = 16;
    let i = 0;
    tickerRef.current = setInterval(() => {
      i++;
      if (i < steps) {
        setNeedleLeft(Math.random() * 100);
      } else {
        clearInterval(tickerRef.current);
        setNeedleLeft(targetFor(target));
        setTimeout(() => {
          setLastHit(target);
          setTally(prev => {
            const key = target === -1 ? "__tail" : preset.dist[target].token;
            return { ...prev, [key]: (prev[key] || 0) + 1 };
          });
          setTotalRolls(t => t + 1);
          setRolling(false);
        }, 620);
      }
    }, 70);
  }, [preset, rolling, segments]);

  const rollMany = useCallback((n) => {
    if (rolling) return;
    // Fast-forward: skip animation, just bump the tally.
    const additions = {};
    for (let k = 0; k < n; k++) {
      const t = sample(preset.dist);
      const key = t === -1 ? "__tail" : preset.dist[t].token;
      additions[key] = (additions[key] || 0) + 1;
    }
    setTally(prev => {
      const next = { ...prev };
      for (const k in additions) next[k] = (next[k] || 0) + additions[k];
      return next;
    });
    setTotalRolls(t => t + n);
  }, [preset, rolling]);

  return (
    <div>
      <div style={{ marginBottom: 28, display: "flex", justifyContent: "space-between", alignItems: "baseline", flexWrap: "wrap", gap: 16 }}>
        <div>
          <div className="mono" style={{ fontSize: 10, letterSpacing: "0.18em", textTransform: "uppercase", color: "var(--muted)", marginBottom: 6 }}>
            Prompt
          </div>
          <div style={{ fontFamily: "var(--serif-display)", fontSize: 28, lineHeight: 1.2 }}>
            <span style={{ color: "var(--ink)" }}>“{preset.prompt}</span>
            <span style={{ color: "var(--vermilion)" }}> ___</span>
            <span style={{ color: "var(--ink)" }}>”</span>
          </div>
        </div>
        <div style={{ display: "flex", gap: 6 }}>
          {PRESETS.map(p => (
            <button
              key={p.id}
              onClick={() => setPresetId(p.id)}
              style={{
                fontFamily: "var(--mono)",
                fontSize: 11,
                padding: "5px 10px",
                background: p.id === presetId ? "var(--ink)" : "transparent",
                color: p.id === presetId ? "var(--paper)" : "var(--ink-soft)",
                border: "1px solid " + (p.id === presetId ? "var(--ink)" : "var(--hairline)"),
                borderRadius: 100,
                cursor: "pointer",
                letterSpacing: "0.08em",
              }}
            >
              {p.id}
            </button>
          ))}
        </div>
      </div>

      <div className={"prob-strip " + (rolling ? "rolling" : "")}>
        {segments.named.map((s, i) => (
          <div
            key={s.token}
            className={`seg s${Math.min(i, 3)}`}
            style={{ left: s.left + "%", width: s.width + "%" }}
            title={`${s.token} — ${fmtPct(s.p)}`}
          >
            {s.width > 4 ? s.token : ""}
          </div>
        ))}
        <div
          className="seg tail"
          style={{ left: segments.tail.left + "%", width: segments.tail.width + "%" }}
        >
          {segments.tail.width > 1.5 ? "…99,996 other tokens" : ""}
        </div>
        <div className="needle" style={{ left: needleLeft + "%" }} />
      </div>
      <div className="prob-axis">
        <span>0%</span><span>25%</span><span>50%</span><span>75%</span><span>100%</span>
      </div>

      {/* Legend — show each named token with its weight and whether the last roll hit it. */}
      <div className="prob-legend">
        {preset.dist.map((d, i) => {
          const hit = lastHit === i;
          const count = tally[d.token] || 0;
          return (
            <div key={d.token} className={"row " + (hit ? "hit" : "")}>
              <div className="sw" style={{
                background: ["rgba(180,53,28,0.85)","rgba(180,53,28,0.32)","rgba(180,53,28,0.16)","rgba(180,53,28,0.07)"][Math.min(i,3)]
              }} />
              <div className="tok">{d.token}</div>
              <div className="pct">{fmtPct(d.p)} · {count} hit{count===1?"":"s"}</div>
            </div>
          );
        })}
        <div className={"row " + (lastHit === -1 ? "hit" : "")} style={{ gridColumn: "1 / -1" }}>
          <div className="sw" style={{ background: "rgba(26,24,21,0.15)" }} />
          <div className="tok muted">the long tail — every other token in the 100,000-side vocabulary, combined</div>
          <div className="pct">{fmtPct(segments.tail.p)} · {tally.__tail || 0} hit{(tally.__tail||0)===1?"":"s"}</div>
        </div>
      </div>

      {/* Roll controls + outcome strip */}
      <div className="roll-log">
        <div className="stat-block">
          <div className="big">{totalRolls}</div>
          <div className="lbl">Rolls cast</div>
          <div style={{ marginTop: 16, display: "flex", gap: 8, flexWrap: "wrap" }}>
            <button className="btn" onClick={roll} disabled={rolling}>
              {rolling ? "Rolling…" : "Roll once"}
            </button>
            <button className="btn ghost" onClick={() => rollMany(99)} disabled={rolling}>
              + 99 rolls
            </button>
          </div>
          <button
            onClick={() => { setTally({}); setTotalRolls(0); setLastHit(null); }}
            style={{
              marginTop: 12, background: "transparent", border: 0, padding: 0,
              fontFamily: "var(--mono)", fontSize: 10, letterSpacing: "0.16em",
              textTransform: "uppercase", color: "var(--muted)", cursor: "pointer",
              borderBottom: "1px solid var(--hairline)",
            }}
          >
            Reset tally
          </button>
        </div>
        <div className="breakdown">
          <div className="mono" style={{ fontSize: 10, letterSpacing: "0.16em", textTransform: "uppercase", color: "var(--muted)", marginBottom: 6 }}>
            Observed frequency after {totalRolls || 0} rolls
          </div>
          {preset.dist.map((d, i) => {
            const count = tally[d.token] || 0;
            const pct = totalRolls ? count / totalRolls : 0;
            return (
              <div key={d.token} className="bar">
                <div style={{ color: i === 0 ? "var(--ink)" : "var(--ink-soft)" }}>{d.token}</div>
                <div className="track">
                  <div className={"fill " + (i === 0 ? "" : "dim")} style={{ width: (pct * 100) + "%" }} />
                </div>
                <div style={{ color: "var(--muted)", textAlign: "right" }}>{(pct * 100).toFixed(1)}%</div>
              </div>
            );
          })}
          <div className="bar">
            <div style={{ color: "var(--muted)" }}>tail</div>
            <div className="track">
              <div className="fill tail" style={{ width: ((totalRolls ? (tally.__tail || 0) / totalRolls : 0) * 100) + "%" }} />
            </div>
            <div style={{ color: "var(--muted)", textAlign: "right" }}>{(totalRolls ? ((tally.__tail || 0) / totalRolls * 100) : 0).toFixed(1)}%</div>
          </div>
        </div>
      </div>

      {(() => {
        // Offline build replaces the Claude-backed lab with a pre-baked
        // library of additional dice; both render in the same `.lab` shell.
        const Offline = window.OfflineMoreLab;
        if (window.__OFFLINE && Offline) return <Offline />;
        return <YourOwnPromptLab />;
      })()}
    </div>
  );
}

/* ---- "Try your own prompt" — asks Claude to estimate the top-K tokens. ---- */

function YourOwnPromptLab() {
  const [prompt, setPrompt] = useState("The first man to walk on the moon was");
  const [working, setWorking] = useState(false);
  const [error, setError] = useState(null);
  const [dist, setDist] = useState(null); // [{token, p}]
  const [lastHit, setLastHit] = useState(null);
  const [needleLeft, setNeedleLeft] = useState(0);
  const [rolling, setRolling] = useState(false);

  const estimate = useCallback(async () => {
    if (!prompt.trim()) return;
    setWorking(true);
    setError(null);
    setDist(null);
    setLastHit(null);
    try {
      const sys = `You are a probabilistic-language-model simulator. Given a partial sentence, list the 4 most likely next single tokens (words or short fragments) along with rough probability estimates, plus one comically unlikely token for contrast. Probabilities should reflect realistic language-model behavior: a confident next-token prediction often has p > 0.5; an unlikely one is < 0.001. Respond ONLY with a JSON array of objects, no prose, no markdown, no code fences. Each object: {"token": string, "p": number}.`;
      const raw = await window.claude.complete({
        messages: [
          { role: "user", content: `${sys}\n\nPartial sentence: "${prompt.trim()}"` },
        ],
      });
      // Tolerate code-fence wrappers if Claude includes them anyway.
      const cleaned = raw.replace(/```(?:json)?/g, "").trim();
      const start = cleaned.indexOf("[");
      const end = cleaned.lastIndexOf("]");
      if (start === -1 || end === -1) throw new Error("Could not find a JSON array in the response.");
      const parsed = JSON.parse(cleaned.slice(start, end + 1));
      if (!Array.isArray(parsed) || !parsed.length) throw new Error("Empty distribution.");
      // Normalize + sort. Clamp p to leave some tail mass.
      const cleanedDist = parsed
        .filter(x => x && typeof x.token === "string" && typeof x.p === "number")
        .map(x => ({ token: x.token.replace(/^["']|["']$/g, ""), p: Math.max(0, Math.min(1, x.p)) }))
        .sort((a, b) => b.p - a.p)
        .slice(0, 5);
      const sum = cleanedDist.reduce((a, b) => a + b.p, 0);
      // If the model returned probabilities that sum to > 0.999, scale them
      // down a touch so there's room for a visible long tail.
      if (sum > 0.985) {
        const factor = 0.985 / sum;
        cleanedDist.forEach(x => x.p *= factor);
      }
      setDist(cleanedDist);
    } catch (err) {
      console.error(err);
      setError(err.message || "Something went wrong.");
    } finally {
      setWorking(false);
    }
  }, [prompt]);

  const segs = useMemo(() => {
    if (!dist) return null;
    let acc = 0;
    const named = dist.map(d => {
      const seg = { ...d, left: acc, width: d.p * 100 };
      acc += d.p * 100;
      return seg;
    });
    return { named, tail: { left: acc, width: Math.max(0, 100 - acc) } };
  }, [dist]);

  const rollOnce = useCallback(() => {
    if (!dist || rolling) return;
    setRolling(true);
    setLastHit(null);
    const target = sample(dist);
    let i = 0;
    const id = setInterval(() => {
      i++;
      if (i < 14) {
        setNeedleLeft(Math.random() * 100);
      } else {
        clearInterval(id);
        const center = target === -1
          ? (segs.tail.left + segs.tail.width / 2)
          : (segs.named[target].left + segs.named[target].width / 2);
        setNeedleLeft(center);
        setTimeout(() => {
          setLastHit(target);
          setRolling(false);
        }, 600);
      }
    }, 70);
  }, [dist, segs, rolling]);

  return (
    <div className="lab">
      <div className="lab-head">
        <h3>Lab · load your own die</h3>
        <span className="mono" style={{ fontSize: 10, letterSpacing: "0.14em", color: "var(--muted)" }}>
          powered by Claude
        </span>
      </div>
      <div>
        <label className="label">Partial sentence</label>
        <div className="input-row">
          <input
            type="text"
            value={prompt}
            onChange={e => setPrompt(e.target.value)}
            onKeyDown={e => { if (e.key === "Enter") estimate(); }}
            placeholder="e.g. The sky on a clear afternoon is"
          />
          <button className="btn" onClick={estimate} disabled={working || !prompt.trim()}>
            {working ? "Estimating…" : "Load the die"}
          </button>
        </div>
        <div className="preset-row">
          {[
            "The first man to walk on the moon was",
            "After it rains, the air smells",
            "She unlocked the door and stepped",
            "The doctor said the prognosis was",
          ].map(p => (
            <button key={p} onClick={() => setPrompt(p)}>{p}</button>
          ))}
        </div>
        {error && <div className="lab-status" style={{ color: "var(--vermilion)" }}>⚠ {error}</div>}
      </div>

      {dist && segs && (
        <div className="lab-result">
          <div className="prob-strip" style={{ height: 64 }}>
            {segs.named.map((s, i) => (
              <div
                key={s.token + i}
                className={`seg s${Math.min(i, 3)}`}
                style={{ left: s.left + "%", width: s.width + "%" }}
              >
                {s.width > 6 ? s.token : ""}
              </div>
            ))}
            <div
              className="seg tail"
              style={{ left: segs.tail.left + "%", width: segs.tail.width + "%" }}
            >
              {segs.tail.width > 4 ? "…tail" : ""}
            </div>
            <div className="needle" style={{ left: needleLeft + "%" }} />
          </div>
          <div style={{ display: "flex", gap: 14, marginTop: 16, alignItems: "center", flexWrap: "wrap" }}>
            <button className="btn" onClick={rollOnce} disabled={rolling}>
              {rolling ? "Rolling…" : "Roll this die"}
            </button>
            {lastHit !== null && (
              <span className="mono" style={{ fontSize: 13 }}>
                landed on <b style={{ color: "var(--vermilion)" }}>
                  {lastHit === -1 ? "(tail — an off-script token)" : dist[lastHit].token}
                </b>
              </span>
            )}
          </div>
          <div className="prob-legend" style={{ marginTop: 20 }}>
            {dist.map((d, i) => (
              <div key={d.token + i} className={"row " + (lastHit === i ? "hit" : "")}>
                <div className="sw" style={{
                  background: ["rgba(180,53,28,0.85)","rgba(180,53,28,0.32)","rgba(180,53,28,0.16)","rgba(180,53,28,0.07)","rgba(180,53,28,0.04)"][Math.min(i,4)]
                }} />
                <div className="tok">{d.token}</div>
                <div className="pct">{fmtPct(d.p)}</div>
              </div>
            ))}
          </div>
        </div>
      )}
    </div>
  );
}

Object.assign(window, { FairDieDemo, LoadedDieDemo });
