/* global React, ReactDOM, parseModelResponse, formatThink, highlightSQL, DAGView, ResultTable, Icon, Step, useTweaks, TweaksPanel, TweakSection, TweakRadio, TweakToggle, lucide */
const { useState, useEffect, useMemo, useRef, useCallback } = React;

// ─── Backend API ──────────────────────────────────────────────
const API_BASE = ''; // same origin (FastAPI serves UI)

async function apiFetchCases() {
  const r = await fetch(`${API_BASE}/api/cases`);
  if (!r.ok) throw new Error(`/api/cases failed: ${r.status}`);
  return r.json();
}
async function apiRunLive(caseId, model = 'both') {
  const r = await fetch(`${API_BASE}/api/run`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ case_id: caseId, model }),
  });
  if (!r.ok) {
    const t = await r.text();
    throw new Error(`run failed (${r.status}): ${t}`);
  }
  return r.json();
}

// ─── Stage states for run orchestration ───────────────────────
function useRunOrchestrator() {
  const [stage, setStage] = useState('idle');
  const [helixStage, setHelixStage] = useState('idle');
  const [sonnetStage, setSonnetStage] = useState('idle');
  const timersRef = useRef([]);

  const stop = useCallback(() => {
    timersRef.current.forEach(clearTimeout);
    timersRef.current = [];
  }, []);

  const reset = useCallback(() => {
    stop();
    setStage('idle');
    setHelixStage('idle');
    setSonnetStage('idle');
  }, [stop]);

  const start = useCallback((helixLatency, sonnetLatency, speedMul = 1) => {
    stop();
    // Faster pacing — was 4.2-8.5s, now ~2.0-4.0s
    const playH = Math.min(4.0, Math.max(2.0, helixLatency / 8)) * speedMul;
    const playS = Math.min(3.4, Math.max(1.8, sonnetLatency / 8)) * speedMul;
    const seq = (durSec, setter) => {
      const ms = durSec * 1000;
      const t = (frac, val) => timersRef.current.push(setTimeout(() => setter(val), frac * ms));
      setter('reasoning');
      t(0.42, 'planning');
      t(0.60, 'coding');
      t(0.85, 'executing');
      t(1.0, 'done');
    };
    setStage('running');
    seq(playH, setHelixStage);
    seq(playS, setSonnetStage);
    const max = Math.max(playH, playS) * 1000 + 100;
    timersRef.current.push(setTimeout(() => setStage('done'), max));
  }, [stop]);

  const finishImmediately = useCallback(() => {
    stop();
    setHelixStage('done');
    setSonnetStage('done');
    setStage('done');
  }, [stop]);

  return { stage, helixStage, sonnetStage, start, reset, stop, finishImmediately };
}

function stepStateFromStage(runStage, stepIdx) {
  const map = { idle: -1, reasoning: 0, planning: 1, coding: 2, executing: 3, done: 4 };
  const cur = map[runStage] ?? -1;
  if (cur === -1) return 'idle';
  if (stepIdx < cur) return 'done';
  if (stepIdx === cur) return 'active';
  return 'pending';
}

// ─── Pane: per-model output ───────────────────────────────────
function ModelPane({ label, badgeKind, model, parsed, runStage, isWinner, openSteps, toggleStep, latencyShown }) {
  const stateFor = (i) => stepStateFromStage(runStage, i);
  const finished = runStage === 'done';
  const correct = model.correct;
  const verdictText = !finished ? null : (correct ? 'Match' : 'Wrong');
  return (
    <div className="pane" data-winner={isWinner ? 'true' : 'false'}>
      <div className="pane-header">
        <div className="pane-title">
          <h3>{label}</h3>
          <span className={`badge ${badgeKind}`}>{badgeKind === 'fine-tune' ? 'Fine-tune' : 'Baseline'}</span>
        </div>
        <div className="pane-status">
          {runStage === 'idle' && <span>Idle</span>}
          {(runStage !== 'idle' && runStage !== 'done') && (
            <>
              <span className="blinker"/>
              <span>{runStage === 'reasoning' ? 'Reasoning' : runStage === 'planning' ? 'Planning' : runStage === 'coding' ? 'Generating SQL' : 'Executing'}</span>
            </>
          )}
          {runStage === 'done' && (
            <>
              {correct
                ? <span className="check"><Icon name="check" size={9} stroke={3}/></span>
                : <span className="cross"><Icon name="x" size={9} stroke={3}/></span>}
              <span>{verdictText} · {latencyShown}s</span>
            </>
          )}
        </div>
      </div>

      <div className="pane-body">
        <Step num={1} title="Reasoning"
              meta={parsed.think ? `${Math.round(parsed.think.length/4)} tokens` : 'skipped'}
              state={stateFor(0)} open={openSteps[0]} onToggle={() => toggleStep(0)}>
          {parsed.think
            ? <div className="think-text">{formatThink(parsed.think)}</div>
            : <div style={{fontFamily:'var(--font-mono)', fontSize:12, color:'var(--fg-3)', padding:'8px 12px', background:'var(--bg-1)', border:'var(--hair)', borderRadius:4}}>
                — model skipped this block —
              </div>}
        </Step>

        <Step num={2} title="Pipeline DAG"
              meta={parsed.dag ? `${parsed.dag.nodes?.length ?? 0} nodes · ${parsed.dag.edges?.length ?? 0} edges` : 'invalid'}
              state={stateFor(1)} open={openSteps[1]} onToggle={() => toggleStep(1)}>
          {parsed.dag
            ? <DAGView dag={parsed.dag}/>
            : <div className="think-text">DAG block missing or unparseable.</div>}
        </Step>

        <Step num={3} title="CTE SQL"
              meta={parsed.code ? `${parsed.code.split('\n').length} lines` : 'missing'}
              state={stateFor(2)} open={openSteps[2]} onToggle={() => toggleStep(2)}>
          <div className="code-block">{highlightSQL(parsed.code)}</div>
        </Step>

        <Step num={4} title="Execution"
              meta={finished ? `${model.rows?.length ?? 0} rows returned` : 'pending'}
              state={stateFor(3)} open={openSteps[3]} onToggle={() => toggleStep(3)}>
          {!finished
            ? <div style={{fontFamily:'var(--font-mono)', fontSize:12, color:'var(--fg-3)'}}>Awaiting completion…</div>
            : (
              <>
                <div className={`result-summary ${correct ? 'match' : 'miss'}`}>
                  <span className={`verdict ${correct ? 'match' : 'miss'}`}>
                    <Icon name={correct ? 'check' : 'x'} size={14} stroke={2.5}/>
                    {correct ? 'Matches gold result' : 'Differs from gold'}
                  </span>
                  <span style={{color:'var(--fg-3)'}}>set-equal · {model.rows?.length ?? 0} rows</span>
                </div>
                {model.exec_error && (
                  <div className="think-text" style={{borderColor:'rgba(255,106,91,0.4)', color:'var(--danger)'}}>
                    SQL error: {model.exec_error}
                  </div>
                )}
                <ResultTable columns={model.columns || []} rows={model.rows || []} max={6}/>
              </>
            )}
        </Step>
      </div>
    </div>
  );
}

// ─── Hero ───────────────────────────────────────────────────
function Hero() {
  return (
    <header className="hero">
      <div className="hero-eyebrow">
        <span className="cursor"/>
        Helix CodeDAG · Live model showcase
      </div>
      <h1 className="hero-headline">
        Plain English in. <em>Working data pipelines</em> out.
      </h1>
      <p className="hero-sub">
        Ask a database question in plain English. Helix returns
        <em style={{fontFamily:'var(--font-serif)', fontStyle:'italic', color:'var(--signal)'}}> three things</em>:
        the reasoning, a typed pipeline DAG, and CTE SQL ready to drop into DBT, Dagster, or Foundry.
      </p>
      <p className="hero-sub" style={{marginTop:-32, fontSize:18, color:'var(--fg-3)'}}>
        Below, run our 32B specialist side by side with Sonnet 4.6 — real questions,
        real SQLite, set-based row comparison. No mocks.
      </p>
      <div className="headline-numbers">
        <div className="hn-cell">
          <div className="hn-label">Helix · SFT (deployed)</div>
          <div className="hn-value">50.1<span className="unit">% EX</span><span className="delta">+9.0</span></div>
          <div className="hn-caption">What you're running in this demo. A 32B specialist on a couple of H100s, on-prem.</div>
        </div>
        <div className="hn-cell">
          <div className="hn-label">Sonnet 4.6 · CTE+DAG</div>
          <div className="hn-value">41.1<span className="unit">% EX</span></div>
          <div className="hn-caption">Frontier baseline, zero-shot, same prompt.</div>
        </div>
        <div className="hn-cell">
          <div className="hn-label">Helix · +RL stage</div>
          <div className="hn-value">53.3<span className="unit">% EX</span><span className="delta">+12.2</span></div>
          <div className="hn-caption">Same model after reinforcement learning with execution rewards.</div>
        </div>
      </div>
      <p style={{marginTop:24, fontSize:13, color:'var(--fg-3)', fontFamily:'var(--font-mono)', letterSpacing:'0.04em'}}>
        Benchmark · 1,432 BIRD-dev samples · 11 SQLite databases · set-based row comparison.
      </p>
    </header>
  );
}

function SchemaBlock({ ddl }) {
  const [open, setOpen] = useState(false);
  const tableCount = (ddl.match(/CREATE TABLE/gi) || []).length;
  return (
    <div className="schema-block">
      <div className="schema-header" onClick={() => setOpen(o => !o)}>
        <div className="schema-header-left">
          <Icon name="database" size={14}/>
          <span>Schema</span>
        </div>
        <div className="schema-tables">
          <b>{tableCount}</b> tables · click to {open ? 'hide' : 'inspect'}
        </div>
      </div>
      {open && <pre className="schema-body">{ddl}</pre>}
    </div>
  );
}

function VerdictCard({ caseData }) {
  const helixOK = caseData.helix.correct;
  const sonnetOK = caseData.sonnet.correct;
  let summary;
  if (helixOK && !sonnetOK)       summary = <>Helix <em>got it right</em> on a question Sonnet missed.</>;
  else if (!helixOK && sonnetOK)  summary = <>Sonnet won this round. We keep the failing cases honest.</>;
  else if (helixOK && sonnetOK)   summary = <>Both <em>correct</em>. Both pipelines match the gold rows.</>;
  else                            summary = <>Both missed this one — uncommon edge case.</>;
  return (
    <div className="verdict-card fade-in">
      <div className="vc-summary">{summary}</div>
      <div>
        <div className="vc-stat-label">Helix · Result</div>
        <div className="vc-stat-value helix">{helixOK ? '✓' : '✗'}</div>
        <div className="vc-stat-meta">{caseData.helix.latency_s}s · {caseData.helix.rows?.length ?? 0} rows</div>
      </div>
      <div>
        <div className="vc-stat-label">Sonnet · Result</div>
        <div className="vc-stat-value sonnet">{sonnetOK ? '✓' : '✗'}</div>
        <div className="vc-stat-meta">{caseData.sonnet.latency_s}s · {caseData.sonnet.rows?.length ?? 0} rows</div>
      </div>
    </div>
  );
}

// ─── Main App ───────────────────────────────────────────────
function App() {
  const [tweaks, setTweak] = useTweaks(/*EDITMODE-BEGIN*/{
    "speed": "normal",
    "showSchema": true,
    "openAllSteps": false,
    "liveApi": false,
    "theme": "light"
  }/*EDITMODE-END*/);

  // sync html[data-theme] whenever the tweak changes
  useEffect(() => {
    document.documentElement.setAttribute('data-theme', tweaks.theme === 'dark' ? 'dark' : 'light');
  }, [tweaks.theme]);

  const [cases, setCases] = useState(null);
  const [loadError, setLoadError] = useState(null);
  const [caseIdx, setCaseIdx] = useState(0); // default to the first ⭐ Helix-wins case
  const [runError, setRunError] = useState(null);
  const [runningLive, setRunningLive] = useState(false);

  const orch = useRunOrchestrator();
  const { helixStage, sonnetStage, stage, start, reset, finishImmediately } = orch;

  const initialOpen = [true, true, true, true];
  const [helixOpen, setHelixOpen] = useState(initialOpen);
  const [sonnetOpen, setSonnetOpen] = useState(initialOpen);
  const toggleH = (i) => setHelixOpen(s => s.map((v, j) => j === i ? !v : v));
  const toggleS = (i) => setSonnetOpen(s => s.map((v, j) => j === i ? !v : v));

  // initial load
  useEffect(() => {
    apiFetchCases()
      .then(setCases)
      .catch(e => setLoadError(e.message));
  }, []);

  // reflow lucide icons whenever DOM updates (must run on every render — no early-return above)
  useEffect(() => { if (window.lucide) window.lucide.createIcons(); });

  if (loadError) {
    return <div style={{padding:40, fontFamily:'var(--font-mono)', color:'var(--danger)'}}>Failed to load cases: {loadError}</div>;
  }
  if (!cases) {
    return <div style={{padding:40, fontFamily:'var(--font-mono)', color:'var(--fg-3)'}}>Loading…</div>;
  }

  const caseData = cases[caseIdx];
  const helixParsed = { think: caseData.helix.think, dag: caseData.helix.dag, code: caseData.helix.code };
  const sonnetParsed = { think: caseData.sonnet.think, dag: caseData.sonnet.dag, code: caseData.sonnet.code };

  const switchCase = (i) => {
    setCaseIdx(i);
    setRunError(null);
    reset();
    setHelixOpen(initialOpen);
    setSonnetOpen(initialOpen);
  };

  const speedMul = tweaks.speed === 'fast' ? 0.5 : tweaks.speed === 'slow' ? 1.6 : 1;

  const handleRun = async () => {
    if (stage !== 'idle') {
      reset();
      setRunError(null);
      return;
    }
    setRunError(null);
    if (tweaks.liveApi) {
      setRunningLive(true);
      start(caseData.helix.latency_s || 12, caseData.sonnet.latency_s || 9, speedMul);
      try {
        const fresh = await apiRunLive(caseData.id, 'both');
        const next = [...cases];
        next[caseIdx] = fresh;
        setCases(next);
        finishImmediately();
      } catch (e) {
        setRunError(e.message);
        reset();
      } finally {
        setRunningLive(false);
      }
    } else {
      start(caseData.helix.latency_s, caseData.sonnet.latency_s, speedMul);
    }
  };

  const winner = stage === 'done' ?
    (caseData.helix.correct && !caseData.sonnet.correct ? 'helix'
     : caseData.sonnet.correct && !caseData.helix.correct ? 'sonnet' : null)
    : null;

  return (
    <>
      <nav className="nav">
        <div className="nav-brand">
          <span className="dot"/>
          Helix · CodeDAG
        </div>
        <div className="nav-meta">
          {tweaks.liveApi && <span style={{color:'var(--signal)'}}><b>LIVE</b></span>}
        </div>
      </nav>

      <Hero/>

      <section className="section">
        <div className="section-label">01 · Pick a question</div>
        <h2 className="section-title">Four real questions, three real databases.</h2>
        <p className="section-sub">
          Each case is a natural-language question against a real SQLite database.
          The two starred ⭐ cases are where Helix gets it right and Sonnet doesn't.
          On the other two, both models land the gold rows.
        </p>
        <p className="section-sub" style={{marginTop:-24, fontSize:14, color:'var(--fg-3)', fontFamily:'var(--font-mono)', letterSpacing:'0.03em'}}>
          Live model: <b style={{color:'var(--signal)'}}>helix-codedag-sft-32b-v0</b> · SFT release at <b style={{color:'var(--fg-1)'}}>50.1% EX</b>.
        </p>
        <div className="case-picker">
          {cases.map((c, i) => {
            const star = c.highlight && c.highlight.startsWith('⭐');
            return (
              <button key={c.id} className="case-card"
                      data-active={i === caseIdx ? 'true' : 'false'}
                      onClick={() => switchCase(i)}>
                <div className="case-card-id">
                  <span>{c.id}</span>
                  {star && <span className="case-card-star"><Icon name="sparkles" size={11}/> Helix wins</span>}
                </div>
                <div className="case-card-name">{c.name}</div>
              </button>
            );
          })}
        </div>

        <div className="case-detail">
          <div>
            <div className="case-detail-label">Selected · {caseData.id}</div>
            <p className="case-detail-query">{caseData.query}</p>
            <div className="case-detail-meta">
              <span className="pill"><Icon name="database" size={11}/> &nbsp;{caseData.db_id}.sqlite</span>
              <span className="pill">expects {caseData.gold_row_count} {caseData.gold_row_count === 1 ? 'row' : 'rows'}</span>
              {caseData.evidence && <span className="pill evidence">hint provided</span>}
            </div>
            {caseData.evidence && (
              <p style={{marginTop:18, fontSize:14, color:'var(--fg-2)', fontStyle:'italic', maxWidth:'62ch', lineHeight:1.5}}>
                Hint: {caseData.evidence}
              </p>
            )}
            {runError && (
              <p style={{marginTop:18, fontSize:13, color:'var(--danger)', fontFamily:'var(--font-mono)'}}>
                {runError}
              </p>
            )}
          </div>
          <div className="case-detail-actions">
            <button className="btn" onClick={handleRun} disabled={runningLive && stage !== 'done'}>
              {stage === 'idle' && <><Icon name="play" size={14} stroke={2.5}/> {tweaks.liveApi ? 'Run live on Foundry' : 'Replay cached run'}</>}
              {stage === 'running' && <><Icon name="loader-circle" size={14} stroke={2.5}/> Running…</>}
              {stage === 'done' && <><Icon name="rotate-ccw" size={14} stroke={2.5}/> Reset & re-run</>}
            </button>
          </div>
        </div>

        {tweaks.showSchema && <SchemaBlock ddl={caseData.schema_ddl}/>}
      </section>

      <section className="section-stage">
        <div className="section-label">02 · Live comparison</div>
        <h2 className="section-title" style={{marginBottom: 36}}>
          Same prompt. Same database. Two models.
        </h2>

        {stage === 'idle' ? (
          <div className="run-prompt fade-in">
            <Icon name="mouse-pointer-click" size={20}/>
            <div>
              <div className="run-prompt-title">Pick a question above, then hit <b>Run</b>.</div>
              <div className="run-prompt-sub">Both models will generate a pipeline DAG, write CTE SQL, and execute against the real database side by side.</div>
            </div>
          </div>
        ) : (
          <>
            <div className="stage">
              <ModelPane label="Helix SFT 32B" badgeKind="fine-tune"
                         model={caseData.helix} parsed={helixParsed}
                         runStage={helixStage} isWinner={winner === 'helix'}
                         openSteps={helixOpen} toggleStep={toggleH}
                         latencyShown={caseData.helix.latency_s}/>
              <ModelPane label="Sonnet 4.6" badgeKind="baseline"
                         model={caseData.sonnet} parsed={sonnetParsed}
                         runStage={sonnetStage} isWinner={winner === 'sonnet'}
                         openSteps={sonnetOpen} toggleStep={toggleS}
                         latencyShown={caseData.sonnet.latency_s}/>
            </div>

            {stage === 'done' && <VerdictCard caseData={caseData}/>}

            {stage === 'done' && (
              <div className="gold fade-in">
                <div className="gold-title"><Icon name="bookmark-check" size={14}/> Gold reference SQL</div>
                <pre className="gold-sql">{caseData.gold_sql}</pre>
              </div>
            )}
          </>
        )}
      </section>

      <section className="section">
        <div className="section-label">03 · Why this works</div>
        <h2 className="section-title">A specialist, verified by execution.</h2>
        <p className="section-sub">
          Helix isn't a generalist that happens to write SQL. It's tuned for one job — turning questions
          into deployable data pipelines — and every training sample was <em style={{fontFamily:'var(--font-serif)', color:'var(--signal)', fontStyle:'italic'}}>proven correct</em> against
          a real database before training began.
        </p>
        <div className="story">
          <div className="story-step">
            <div className="step-num-large">i.</div>
            <h4>Built for pipelines, not chat.</h4>
            <p>One job, tuned end-to-end: <em style={{fontFamily:'var(--font-mono)', fontStyle:'normal', fontSize:'0.92em'}}>think → DAG → CTE SQL</em>. Every DAG node maps 1:1 to a CTE — drop the output straight into DBT, Dagster, or Foundry with lineage built in.</p>
          </div>
          <div className="story-step">
            <div className="step-num-large">ii.</div>
            <h4>Every sample, proven correct.</h4>
            <p>11,086 training samples from BIRD and Spider. Each one's SQL was run against a real database and only kept if the rows matched the gold result. No "looks reasonable" data ever reached training.</p>
          </div>
          <div className="story-step">
            <div className="step-num-large">iii.</div>
            <h4>Two training stages.</h4>
            <p>The deployed model is the <b>SFT</b> stage — 50.1% EX, already +9 over Sonnet. Layering reinforcement learning on top, with an execution-based reward, adds another +3.2 points (<b>53.3%</b>). Same prompt. Same on-prem footprint. Better correctness.</p>
          </div>
        </div>
      </section>

      <footer className="footer">
        <div>
          <a href="https://huggingface.co/AI71ai/helix-codedag-sft-32b-v0" target="_blank" rel="noopener noreferrer">
            helix-codedag-sft-32b-v0 · model card ↗
          </a>
        </div>
        <div>
          Analysis or suggestions? <a href="mailto:yingbiao.huang@ai71.ai">yingbiao.huang@ai71.ai</a>
        </div>
      </footer>

      <TweaksPanel title="Tweaks">
        <TweakSection title="Run">
          <TweakToggle label="Live API (call Foundry on Run)"
                       value={tweaks.liveApi}
                       onChange={v => setTweak('liveApi', v)}/>
          <TweakRadio label="Playback speed" value={tweaks.speed}
                      onChange={v => setTweak('speed', v)}
                      options={[
                        { value: 'fast', label: 'Fast' },
                        { value: 'normal', label: 'Normal' },
                        { value: 'slow', label: 'Slow' },
                      ]}/>
        </TweakSection>
        <TweakSection title="Layout">
          <TweakRadio label="Theme" value={tweaks.theme}
                      onChange={v => setTweak('theme', v)}
                      options={[
                        { value: 'light', label: 'Light' },
                        { value: 'dark', label: 'Dark' },
                      ]}/>
          <TweakToggle label="Show schema preview"
                       value={tweaks.showSchema}
                       onChange={v => setTweak('showSchema', v)}/>
          <TweakToggle label="Auto-expand all output blocks"
                       value={tweaks.openAllSteps}
                       onChange={v => {
                         setTweak('openAllSteps', v);
                         setHelixOpen([v, v, v, v]);
                         setSonnetOpen([v, v, v, v]);
                       }}/>
        </TweakSection>
      </TweaksPanel>
    </>
  );
}

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