// App — orchestrates the pipeline.
// State: prompt (string), generated tokens (appended), current "step" (which stage is active),
// progress within that stage (0..1). On step=5 completion, append sampled token and loop.

const STAGE_DURATIONS = [1.4, 1.8, 3.2, 2.6, 2.2]; // seconds at 1× — attention + mlp are the "heavy" ones

function App() {
  const [prompt, setPrompt] = React.useState('The machine dreams in');
  const [generated, setGenerated] = React.useState([]); // array of {text, id}
  const [stage, setStage] = React.useState(0);           // 0..4
  const [progress, setProgress] = React.useState(0);     // 0..1
  const [playing, setPlaying] = React.useState(true);
  const [speed, setSpeed] = React.useState(1.0);
  const [temperature, setTemperature] = React.useState(0.9);
  const [helpOpen, setHelpOpen] = React.useState(false);
  const [helpStage, setHelpStage] = React.useState(0);
  const [autoScroll, setAutoScroll] = React.useState(() => {
    // Default ON. Respect explicit off-preference if the user has toggled it.
    try {
      const v = localStorage.getItem('tlv-autoscroll');
      return v === null ? true : v === '1';
    } catch { return true; }
  });
  React.useEffect(() => {
    try { localStorage.setItem('tlv-autoscroll', autoScroll ? '1' : '0'); } catch {}
  }, [autoScroll]);
  const scrollContainerRef = React.useRef(null);
  const wasPlayingBeforeHelpRef = React.useRef(false);

  // Auto-pause while help is open; restore on close
  React.useEffect(() => {
    if (helpOpen) {
      wasPlayingBeforeHelpRef.current = playingRef.current;
      setPlaying(false);
    } else if (wasPlayingBeforeHelpRef.current) {
      setPlaying(true);
      wasPlayingBeforeHelpRef.current = false;
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [helpOpen]);

  // Tokenize prompt + generated
  const promptTokens = React.useMemo(() => toyTokenize(prompt), [prompt]);
  const allTokens = React.useMemo(() => [...promptTokens, ...generated], [promptTokens, generated]);
  const activeIdx = allTokens.length - 1; // we're always generating after the last token

  // Animation loop
  const lastTimeRef = React.useRef(performance.now());
  const stageRef = React.useRef(stage);
  const progressRef = React.useRef(progress);
  const speedRef = React.useRef(speed);
  const playingRef = React.useRef(playing);
  const pendingTokenRef = React.useRef(null);
  stageRef.current = stage; progressRef.current = progress;
  speedRef.current = speed; playingRef.current = playing;

  React.useEffect(() => {
    let raf;
    const tick = (now) => {
      const dt = Math.min(0.1, (now - lastTimeRef.current) / 1000);
      lastTimeRef.current = now;
      if (playingRef.current) {
        const s = stageRef.current;
        const dur = STAGE_DURATIONS[s] || 1;
        let p = progressRef.current + (dt * speedRef.current) / dur;
        if (p >= 1) {
          // advance
          if (s < 4) {
            setStage(s + 1); setProgress(0);
          } else {
            // Completed full cycle — commit sampled token
            if (pendingTokenRef.current) {
              const tok = pendingTokenRef.current;
              pendingTokenRef.current = null;
              let h = 2166136261;
              for (let i = 0; i < tok.length; i++) { h ^= tok.charCodeAt(i); h = Math.imul(h, 16777619); }
              const id = Math.abs(h) % 50000;
              setGenerated((g) => {
                const next = [...g, { text: tok, id }];
                return next.length > 24 ? next.slice(-24) : next;
              });
            }
            setStage(0); setProgress(0);
          }
        } else {
          setProgress(p);
        }
      } else {
        // keep lastTime fresh while paused so unpause doesn't jump
        // (already updated above)
      }
      raf = requestAnimationFrame(tick);
    };
    lastTimeRef.current = performance.now();
    raf = requestAnimationFrame(tick);
    return () => cancelAnimationFrame(raf);
  }, []);

  const handleSample = React.useCallback((tok) => {
    pendingTokenRef.current = tok;
  }, []);

  const handleStep = () => {
    setPlaying(false);
    setProgress(1 - 0.001);
    setTimeout(() => {
      setProgress(0);
      setStage((s) => (s + 1) % 5);
    }, 50);
  };

  const handleReset = () => {
    setGenerated([]); setStage(0); setProgress(0); setPlaying(true);
  };

  // Page-level keyboard shortcuts — disabled when help overlay is open (it owns keys then)
  // and when user is typing in an input.
  React.useEffect(() => {
    if (helpOpen) return; // tour has the keys
    const onKey = (e) => {
      const t = e.target;
      if (t && (t.tagName === 'INPUT' || t.tagName === 'TEXTAREA' || t.isContentEditable)) return;
      // Ignore when any modifier is down — we don't want to hijack browser shortcuts
      if (e.metaKey || e.ctrlKey || e.altKey) return;

      if (e.key === ' ') {
        e.preventDefault();
        setPlaying((p) => !p);
      } else if (e.key === 'ArrowRight' || e.key === '.') {
        e.preventDefault();
        // step forward one stage
        setPlaying(false);
        setProgress(1 - 0.001);
        setTimeout(() => { setProgress(0); setStage((s) => (s + 1) % 5); }, 50);
      } else if (e.key === 'ArrowLeft' || e.key === ',') {
        e.preventDefault();
        // step back one stage
        setPlaying(false);
        setStage((s) => (s - 1 + 5) % 5);
        setProgress(0);
      } else if (e.key === 'r' || e.key === 'R') {
        e.preventDefault();
        handleReset();
      } else if (e.key === 'f' || e.key === 'F') {
        e.preventDefault();
        setAutoScroll((v) => !v);
      } else if (e.key === '?' || (e.shiftKey && e.key === '/')) {
        e.preventDefault();
        setHelpStage(-1); setHelpOpen(true);
      } else if (e.key >= '1' && e.key <= '5') {
        e.preventDefault();
        const idx = parseInt(e.key, 10) - 1;
        setStage(idx); setProgress(0);
      }
    };
    window.addEventListener('keydown', onKey);
    return () => window.removeEventListener('keydown', onKey);
  }, [helpOpen]);

  const explainStage = React.useCallback((i) => {
    // Pipeline stages 0..4 → STAGE_INFO indices (residual-stream at idx 3 is an interstitial, so shift mlp/logits)
    const MAP = [0, 1, 2, 4, 5];
    setHelpStage(MAP[i] ?? i);
    setHelpOpen(true);
  }, []);

  // Auto-scroll to active stage when enabled
  React.useEffect(() => {
    if (!autoScroll) return;
    const container = scrollContainerRef.current;
    if (!container) return;
    const key = ['stage-tokenize','stage-embed','stage-attention','stage-mlp','stage-logits'][stage];
    const el = container.querySelector(`[data-help-target="${key}"]`);
    if (!el) return;
    const cr = container.getBoundingClientRect();
    const er = el.getBoundingClientRect();
    const target = container.scrollTop + (er.top - cr.top) - 12;
    container.scrollTo({ top: target, behavior: 'smooth' });
  }, [stage, autoScroll]);

  // Layout
  return (
    <div style={{ height: '100vh', display: 'grid', gridTemplateRows: 'auto 1fr', overflow: 'hidden' }}>
      <div data-app-header>
      <Controls
        prompt={prompt}
        onPromptChange={(v) => { setPrompt(v); setGenerated([]); setStage(0); setProgress(0); }}
        playing={playing}
        onPlayToggle={() => setPlaying(p => !p)}
        onStep={handleStep}
        onReset={handleReset}
        speed={speed}
        onSpeedChange={setSpeed}
        temperature={temperature}
        onTempChange={setTemperature}
        tokenCount={allTokens.length}
        stepIdx={stage}
        progress={progress}
        onJumpStage={(i) => { setStage(i); setProgress(0); }}
        onHelp={() => { setHelpStage(-1); setHelpOpen(true); setPlaying(false); }}
        autoScroll={autoScroll}
        onAutoScrollToggle={() => setAutoScroll((v) => !v)}
      />
      </div>

      <div style={{
        display: 'grid',
        gridTemplateColumns: 'minmax(0, 1fr) 360px',
        gap: 14,
        padding: 14,
        minHeight: 0, overflow: 'hidden',
      }}>
        {/* left: pipeline stack */}
        <div ref={scrollContainerRef} data-scroll-container style={{
          display: 'grid',
          gridTemplateRows: 'auto auto 1fr auto auto',
          gap: 10,
          minHeight: 0, minWidth: 0, overflow: 'auto',
          paddingRight: 4,
        }}>
          <div data-help-target="stage-tokenize" style={{ minWidth: 0 }}><StageTokenize tokens={allTokens} activeIdx={activeIdx} progress={stage === 0 ? progress : (stage > 0 ? 1 : 0)} active={stage === 0} onExplain={() => explainStage(0)} /></div>
          <div data-help-target="stage-embed" style={{ minWidth: 0 }}><StageEmbed tokens={allTokens} activeIdx={activeIdx} progress={stage === 1 ? progress : (stage > 1 ? 1 : 0)} active={stage === 1} onExplain={() => explainStage(1)} /></div>
          <div data-help-target="stage-attention" style={{ minWidth: 0 }}><StageAttention tokens={allTokens} activeIdx={activeIdx} progress={stage === 2 ? progress : (stage > 2 ? 1 : 0)} active={stage === 2} nHeads={8} onExplain={() => explainStage(2)} /></div>
          <div data-help-target="stage-mlp" style={{ minWidth: 0 }}><StageMLP tokens={allTokens} activeIdx={activeIdx} progress={stage === 3 ? progress : (stage > 3 ? 1 : 0)} active={stage === 3} nLayers={32} onExplain={() => explainStage(3)} /></div>
          <div data-help-target="stage-logits" style={{ minWidth: 0 }}><StageLogits context={allTokens} progress={stage === 4 ? progress : 0} active={stage === 4} paused={!playing} temperature={temperature} onSample={handleSample} onExplain={() => explainStage(4)} /></div>
        </div>

        {/* right: narration + matrix rain backdrop */}
        <div style={{ display: 'grid', gridTemplateRows: '1fr auto', gap: 10, minHeight: 0 }}>
          <Narration stage={stage} tokens={allTokens} activeIdx={activeIdx} nextToken={pendingTokenRef.current} speed={speed} />
          <RainBackdrop
            stage={stage}
            progress={progress}
            tokens={allTokens}
            stageRef={stageRef}
            progressRef={progressRef}
            onExplain={() => { setHelpStage(3); setHelpOpen(true); }}
          />
        </div>
      </div>

      <HelpOverlay open={helpOpen} onClose={() => setHelpOpen(false)} stageIdx={helpStage} setStageIdx={setHelpStage} />
    </div>
  );
}

// Subtle number-rain canvas — NOT recreating branded imagery;
// just streaming decimals in our phosphor green aesthetic. Supports
// two states: the default 160px corner, and an "expanded" state that
// fills the viewport below the header (Esc closes).
function RainBackdrop({ stage, progress, onExplain, tokens, stageRef, progressRef }) {
  const canvasRef = React.useRef(null);
  const containerRef = React.useRef(null);
  const colsRef = React.useRef([]);
  const [expanded, setExpanded] = React.useState(false);
  const [headerOffset, setHeaderOffset] = React.useState(64);

  // Measure the header (first grid row) so the expanded overlay leaves it
  // visible. Re-measure on window resize.
  React.useLayoutEffect(() => {
    const measure = () => {
      const header = document.querySelector('[data-app-header]');
      if (header) setHeaderOffset(header.getBoundingClientRect().height);
    };
    measure();
    window.addEventListener('resize', measure);
    return () => window.removeEventListener('resize', measure);
  }, []);

  // Esc closes when expanded. Scope-aware so we don't hijack the tour or
  // inputs. Matches the pattern used elsewhere in app.jsx.
  React.useEffect(() => {
    if (!expanded) return;
    const onKey = (e) => {
      const t = e.target;
      if (t && (t.tagName === 'INPUT' || t.tagName === 'TEXTAREA' || t.isContentEditable)) return;
      if (e.key === 'Escape') { e.preventDefault(); setExpanded(false); }
    };
    window.addEventListener('keydown', onKey);
    return () => window.removeEventListener('keydown', onKey);
  }, [expanded]);

  // Canvas rendering loop. ResizeObserver re-initializes on container size
  // change. Per-frame work scales down for larger canvases (fewer columns,
  // shorter trail, 30fps cap) so expanded mode doesn't feel sluggish.
  React.useEffect(() => {
    const canvas = canvasRef.current; if (!canvas) return;
    const ctx = canvas.getContext('2d');
    const dpr = Math.min(2, window.devicePixelRatio || 1);
    let w = 0, h = 0;
    let fs = 11;
    let trail = 10;
    let frameIntervalMs = 0;  // 0 = uncapped (60fps); non-zero = cap
    const COL_SPACING = 0.7;

    const init = () => {
      w = canvas.clientWidth; h = canvas.clientHeight;
      if (w === 0 || h === 0) return;
      // Heuristic: corner panel (≤ ~400px wide) runs dense + full-trail at
      // 60fps. Once expanded (viewport-width), we're drawing ~5x the pixels
      // so dial back: larger glyphs → fewer columns, shorter trail, 30fps cap.
      const isLarge = w > 500;
      fs = isLarge ? 14 : 11;
      trail = isLarge ? 6 : 10;
      frameIntervalMs = isLarge ? 33 : 0; // 30fps when expanded, 60fps when corner
      canvas.width = w * dpr; canvas.height = h * dpr;
      ctx.setTransform(1, 0, 0, 1, 0, 0);
      ctx.scale(dpr, dpr);
      const colCount = Math.floor(w / (fs * COL_SPACING));
      colsRef.current = Array.from({ length: colCount }, () => ({
        y: Math.random() * h, speed: 30 + Math.random() * 80, seed: Math.random() * 9999,
      }));
    };
    init();

    let raf, last = performance.now();
    const draw = (now) => {
      const elapsed = now - last;
      // Frame cap: skip draws until we've accumulated enough real time.
      if (frameIntervalMs > 0 && elapsed < frameIntervalMs) {
        raf = requestAnimationFrame(draw);
        return;
      }
      const dt = elapsed / 1000;
      last = now;
      ctx.fillStyle = 'rgba(10,13,10,0.22)';
      ctx.fillRect(0, 0, w, h);
      ctx.font = `${fs}px JetBrains Mono, monospace`;
      const cols = colsRef.current;
      // Pre-compute the per-trail alpha strings so we don't re-concatenate
      // inside the hot loop.
      for (let i = 0; i < cols.length; i++) {
        const c = cols[i];
        c.y += c.speed * dt;
        if (c.y > h + 100) { c.y = -Math.random() * h; }
        const x = i * (fs * COL_SPACING) + 2;
        const ch = (Math.random() > 0.5 ? '+' : '-') + (Math.random() * 9.9).toFixed(1);
        ctx.fillStyle = 'rgba(170, 255, 190, 0.9)';
        ctx.fillText(ch, x, c.y);
        for (let t = 1; t < trail; t++) {
          ctx.fillStyle = `rgba(120, 220, 150, ${0.14 - t * 0.012})`;
          const tch = (Math.random() > 0.5 ? '+' : '-') + (Math.random() * 9.9).toFixed(1);
          ctx.fillText(tch, x, c.y - t * fs);
        }
      }
      raf = requestAnimationFrame(draw);
    };
    raf = requestAnimationFrame(draw);

    const ro = new ResizeObserver(() => init());
    ro.observe(canvas);
    return () => { cancelAnimationFrame(raf); ro.disconnect(); };
  }, []);

  const stageLabel = ['TOKENIZE','EMBED','ATTEND','MLP','SAMPLE'][stage];
  const pct = (progress * 100).toFixed(0).padStart(2, '0');

  // Reusable UI (same markup for both states — only container geometry changes)
  const chrome = (
    <>
      <canvas ref={canvasRef} style={{ position: 'absolute', inset: 0, width: '100%', height: '100%' }} />
      <div style={{
        position: 'absolute', inset: 0,
        background: 'linear-gradient(180deg, rgba(10,13,10,0.1), rgba(10,13,10,0.85))',
        pointerEvents: 'none',
      }} />
      <div style={{ position: 'absolute', top: 8, right: 8, display: 'flex', gap: 6, zIndex: 5 }}>
        {onExplain && <ExplainChip onClick={(e) => { e.stopPropagation(); onExplain(); }} />}
      </div>
      <div style={{
        position: 'absolute',
        left: expanded ? 24 : 12,
        bottom: expanded ? 20 : 10,
        right: expanded ? 24 : 12,
        display: 'flex', justifyContent: 'space-between', alignItems: 'baseline',
        fontFamily: 'var(--mono)',
        fontSize: expanded ? 14 : 10,
        color: 'var(--text-dim)',
        letterSpacing: '0.1em',
        pointerEvents: 'none',
      }}>
        <span>RESIDUAL·STREAM</span>
        <span style={{
          color: 'var(--ink)',
          textShadow: expanded ? '0 0 10px var(--ink)' : '0 0 6px var(--ink)',
        }}>
          {stageLabel} · {pct}%
        </span>
      </div>
    </>
  );

  // Two layers: an always-160px placeholder keeps the grid row stable,
  // and a position-aware inner container (same DOM element in both states
  // so the canvas stays mounted and ResizeObserver handles the dimension
  // change without re-initializing).
  const [hover, setHover] = React.useState(false);
  const hoverCue = hover && !expanded;
  const innerStyle = expanded
    ? {
        position: 'fixed',
        top: headerOffset, left: 0, right: 0, bottom: 0,
        background: 'rgba(10,13,10,0.98)',
        borderTop: '1px solid var(--line-strong)',
        overflow: 'hidden',
        zIndex: 80, // below help overlay (typically 999)
        cursor: 'zoom-out',
        animation: 'rain-expand-in 180ms ease-out',
      }
    : {
        position: 'absolute', inset: 0,
        border: `1px solid ${hoverCue ? 'color-mix(in oklch, var(--ink) 50%, transparent)' : 'var(--line)'}`,
        borderRadius: 8, overflow: 'hidden',
        background: 'rgba(0,0,0,0.5)',
        cursor: 'zoom-in',
        transition: 'border-color 180ms',
        boxShadow: hoverCue ? '0 0 0 1px color-mix(in oklch, var(--ink) 25%, transparent), 0 0 18px color-mix(in oklch, var(--ink) 15%, transparent)' : 'none',
      };

  return (
    <div
      data-help-target="stage-residual"
      style={{ position: 'relative', height: 160 }}
    >
      {expanded && (
        <div style={{
          position: 'absolute', inset: 0,
          border: '1px dashed var(--line-strong)', borderRadius: 8,
          display: 'flex', alignItems: 'center', justifyContent: 'center',
          fontFamily: 'var(--mono)', fontSize: 10, color: 'var(--text-faint)',
          letterSpacing: '0.1em', pointerEvents: 'none',
        }}>
          EXPANDED · ESC OR CLICK TO RETURN
        </div>
      )}
      <div
        ref={containerRef}
        style={innerStyle}
        onClick={() => setExpanded(v => !v)}
        onMouseEnter={() => setHover(true)}
        onMouseLeave={() => setHover(false)}
      >
        {chrome}
        {/* 3D scene disabled for now; 2D rain canvas fills the expanded view.
            Flip this flag to bring it back. Code lives in src/scene-3d.jsx. */}
        {false && expanded && typeof window !== 'undefined' && window.ResidualScene
          ? <window.ResidualScene
              active={expanded}
              tokens={tokens}
              stageRef={stageRef}
              progressRef={progressRef}
            />
          : null}
        {hoverCue && (
          <div style={{
            position: 'absolute', top: 10, left: 12,
            fontFamily: 'var(--mono)', fontSize: 9, letterSpacing: '0.14em',
            color: 'var(--ink)', textShadow: '0 0 6px var(--ink)',
            pointerEvents: 'none',
            opacity: 0.85,
            animation: 'rain-hint-in 160ms ease-out',
          }}>
            ↗ CLICK TO EXPAND
          </div>
        )}
      </div>
      <style>{`
        @keyframes rain-expand-in {
          from { opacity: 0; transform: translateY(4px); }
          to   { opacity: 1; transform: translateY(0); }
        }
        @keyframes rain-hint-in {
          from { opacity: 0; }
          to   { opacity: 0.85; }
        }
      `}</style>
    </div>
  );
}

// Small "?" chip matching the EXPLAIN buttons on each stage panel header.
function ExplainChip({ onClick }) {
  const [hover, setHover] = React.useState(false);
  return (
    <button
      onClick={onClick}
      onMouseEnter={() => setHover(true)}
      onMouseLeave={() => setHover(false)}
      title="Explain the residual stream"
      aria-label="Explain the residual stream"
      style={{
        width: 22, height: 22, padding: 0,
        display: 'inline-flex', alignItems: 'center', justifyContent: 'center',
        fontFamily: 'var(--mono)', fontSize: 11, lineHeight: 1,
        color: hover ? 'var(--bg)' : 'var(--ink)',
        background: hover ? 'var(--ink)' : 'transparent',
        border: `1px solid ${hover ? 'var(--ink)' : 'color-mix(in oklch, var(--ink) 45%, transparent)'}`,
        borderRadius: '50%',
        boxShadow: hover ? '0 0 10px color-mix(in oklch, var(--ink) 55%, transparent)' : 'none',
        cursor: 'pointer',
        transition: 'all 140ms',
      }}
    >?</button>
  );
}


const root = ReactDOM.createRoot(document.getElementById('root'));
root.render(<App />);