// scenes-6.jsx — Caching internals I: the KV cache (what is actually stored)

function KVCell({ state, cell }) {
  const col = state === 'cached' ? COLORS.green
            : state === 'computing' ? COLORS.yellow
            : state === 'fresh' ? COLORS.blue
            : COLORS.inkFaint;
  const op = state === 'idle' ? 0.22 : 1;
  return (
    <div style={{
      width: cell, height: cell, borderRadius: 6,
      background: col + '22', border: `1.5px solid ${col}`, opacity: op,
      boxShadow: state === 'computing' ? `0 0 12px ${col}` : 'none',
      display: 'flex', flexDirection: 'column', justifyContent: 'center', gap: 3, padding: '0 5px',
      transition: 'none',
    }}>
      <div style={{ height: 3, borderRadius: 2, background: col, opacity: 0.95 }} />
      <div style={{ height: 3, borderRadius: 2, background: col, opacity: 0.5 }} />
    </div>
  );
}

function KVGrid({ x, y, cols, rows, cell, gap, colState }) {
  return (
    <div style={{ position: 'absolute', left: x, top: y, display: 'flex', gap }}>
      {Array.from({ length: cols }).map((_, c) => (
        <div key={c} style={{ display: 'flex', flexDirection: 'column', gap }}>
          {Array.from({ length: rows }).map((_, r) => (
            <KVCell key={r} state={colState(c)} cell={cell} />
          ))}
        </div>
      ))}
    </div>
  );
}

// ════════════════════════════════════════════════════════════════════════════
// SCENE — INSIDE THE CACHE: KV STATES  (length 24s)
// ════════════════════════════════════════════════════════════════════════════
function SceneKVCache() {
  const { localTime: lt } = useSprite();
  const setup = ramp(lt, 0.2, 1.0);

  const PFX = 8, rows = 5, cell = 36, gap = 8;
  const cols = PFX + 1;                       // +1 tail token (appears in 2nd call)
  const gx = 470, gy = 360;
  const gridW = cols * cell + (cols - 1) * gap;
  const headY = gy - 54;

  // ── phases ─────────────────────────────────────────────────────────────────
  // prefill: columns 0..7 light up sequentially 2.6→6.2
  const prefillAt = (c) => 2.6 + c * 0.42;
  const stored   = lt > 8.2;
  const cacheFill = clamp(ramp(lt, 8.0, 9.8), 0, 1);
  const cacheGlow = pulse(lt, 7.8, 24, 0.6);
  const turn2    = lt > 11.2;                 // second call
  const tailComputeAt = 14.4;

  const colState = (c) => {
    if (c < PFX) {                            // prefix columns
      if (turn2) return lt > 11.6 + c * 0.16 ? 'cached' : 'fresh';
      if (lt > prefillAt(c) + 0.3) return 'fresh';
      if (lt > prefillAt(c)) return 'computing';
      return 'idle';
    }
    // tail column (the one new token)
    if (!turn2) return 'idle';
    if (lt > tailComputeAt + 0.3) return 'fresh';
    if (lt > tailComputeAt) return 'computing';
    return 'idle';
  };

  const cacheXY = [1640, gy + (rows * cell + (rows - 1) * gap) / 2];
  const flowStore = lt > 8.0 && lt < 9.8;
  const flowRead  = turn2 && lt > 11.4 && lt < 13.6;

  return (
    <>
      <Bg accent={COLORS.green} />
      <Eyebrow lt={lt} a={0.4} b={24} n="10" label="Inside the cache · KV states" color={COLORS.green} />

      {/* call badge */}
      <div style={{ position:'absolute', left: gx + gridW/2, top: 150, transform:'translateX(-50%)', opacity:setup, textAlign:'center', whiteSpace:'nowrap' }}>
        <div style={{ fontFamily:FONTS.mono, fontSize:19, letterSpacing:'0.14em', textTransform:'uppercase',
          color: turn2 ? COLORS.green : COLORS.blue }}>
          {turn2 ? 'second call · prefill skipped' : 'first call · full prefill'}
        </div>
      </div>

      {/* token header row */}
      {setup > 0.3 && Array.from({ length: cols }).map((_, c) => {
        const st = colState(c);
        const col = st === 'cached' ? COLORS.green : st === 'computing' ? COLORS.yellow : st === 'fresh' ? COLORS.blue : COLORS.inkFaint;
        const isTail = c === PFX;
        const o = isTail ? ramp(lt, 11.2, 11.9) : setup;
        if (o <= 0.01) return null;
        return (
          <div key={c} style={{
            position:'absolute', left: gx + c*(cell+gap), top: headY, width: cell, height: 30, opacity:o,
            borderRadius:5, background: col+'33', border:`1.4px solid ${col}`,
          }} />
        );
      })}
      <div style={{ position:'absolute', left: gx, top: headY - 30, fontFamily:FONTS.sans, fontSize:17, color:COLORS.inkDim, opacity:setup }}>
        tokens →
      </div>

      {/* the grid */}
      {setup > 0.2 && <KVGrid x={gx} y={gy} cols={cols} rows={rows} cell={cell} gap={gap} colState={colState} />}

      {/* layers bracket on the left */}
      {setup > 0.4 && (
        <div style={{ position:'absolute', left: gx - 30, top: gy, height: rows*cell + (rows-1)*gap, width: 6,
          borderLeft:`3px solid ${COLORS.inkFaint}`, borderTop:`3px solid ${COLORS.inkFaint}`, borderBottom:`3px solid ${COLORS.inkFaint}`, borderRadius:'6px 0 0 6px', opacity:setup }}>
          <div style={{ position:'absolute', left:-150, top:'50%', width:150, transform:'translateY(-50%) rotate(-90deg)', transformOrigin:'center',
            textAlign:'center', fontFamily:FONTS.sans, fontSize:17, fontWeight:600, color:COLORS.inkDim, whiteSpace:'nowrap' }}>
            attention layers
          </div>
        </div>
      )}

      {/* legend: each cell = K/V */}
      <div style={{ position:'absolute', left: gx, top: gy + rows*cell + (rows-1)*gap + 22, opacity: pulse(lt,1.2,24,0.5),
        fontFamily:FONTS.mono, fontSize:18, color:COLORS.inkDim, display:'flex', gap:14, alignItems:'center', whiteSpace:'nowrap' }}>
        <span style={{ display:'inline-flex', flexDirection:'column', gap:3, width:26, height:26, padding:4, borderRadius:5, border:`1.4px solid ${COLORS.inkDim}`, flexShrink:0 }}>
          <span style={{ height:3, background:COLORS.inkDim, borderRadius:2 }} /><span style={{ height:3, background:COLORS.inkDim, opacity:0.5, borderRadius:2 }} />
        </span>
        <span>one cell = the <b style={{color:COLORS.ink}}>Key</b> &amp; <b style={{color:COLORS.ink}}>Value</b> vectors for one token, at one layer</span>
      </div>

      {/* cache cylinder */}
      <CacheStore x={cacheXY[0]} y={cacheXY[1]} scale={ramp(lt,6.4,7.4,Easing.easeOutBack)} fill={cacheFill} glow={cacheGlow} label="KV cache" />
      <Wire from={[gx + gridW + 6, cacheXY[1]]} to={[cacheXY[0]-96, cacheXY[1]]}
        color={COLORS.green + ((flowStore||flowRead)?'':'44')} width={2.5} draw={ramp(lt,6.6,7.4)}
        flow={flowStore ? ((lt%0.8)/0.8) : (flowRead ? (1-((lt%0.8)/0.8)) : null)} />
      <div style={{ position:'absolute', left:(gx+gridW+cacheXY[0])/2, top: cacheXY[1]-30, transform:'translate(-50%,-50%)',
        fontFamily:FONTS.sans, fontSize:16, fontWeight:600, color:COLORS.green,
        opacity:(flowStore?pulse(lt,8.0,9.8,0.3):0)+(flowRead?pulse(lt,11.4,13.6,0.3):0) }}>
        {flowStore ? 'write KV ↘' : 'load KV ↖'}
      </div>

      <Caption lt={lt} a={1.2} b={2.6}>
        What does the cache actually <b>hold</b>? Not text — tensors.
      </Caption>
      <Caption lt={lt} a={2.8} b={6.4}>
        Processing a prompt runs a <b style={{color:COLORS.yellow}}>prefill</b>: every token, at every layer, computes a Key &amp; Value.
      </Caption>
      <Caption lt={lt} a={6.6} b={8.0} color={COLORS.yellow}>
        That matrix math is the expensive part of the forward pass.
      </Caption>
      <Caption lt={lt} a={8.2} b={11.0} color={COLORS.green}>
        Caching writes those <b>KV tensors</b> to memory, keyed by the exact token prefix.
      </Caption>
      <Caption lt={lt} a={11.6} b={14.2}>
        Next call: matching columns are <b style={{color:COLORS.green}}>loaded</b> — their prefill is skipped entirely.
      </Caption>
      <Caption lt={lt} a={14.6} b={18.0}>
        Only the <b style={{color:COLORS.blue}}>new token</b> runs a real forward pass.
      </Caption>
      <Caption lt={lt} a={18.2} b={24} color={COLORS.green}>
        You still pay a small read fee — memory I/O isn't free — hence ~10%, not zero.
      </Caption>
    </>
  );
}

Object.assign(window, { SceneKVCache, KVGrid, KVCell });