frontend/src/pages/RunWithLlm.jsx

import { useState, useRef } from 'react'
import { Play, Terminal, Activity } from 'lucide-react'
import { llmRunStream } from '../lib/api.js'
import { RewardBreakdown } from '../components/RewardBreakdown.jsx'
import { TurnCard } from '../components/TurnCard.jsx'
import { KpiCard } from '../components/KpiCard.jsx'

const PRESET_ENDPOINTS = [
  { label: 'Ollama (localhost:11434)', url: 'http://localhost:11434/v1' },
  { label: 'Hugging Face Router',       url: 'https://router.huggingface.co/v1' },
  { label: 'OpenAI',                    url: 'https://api.openai.com/v1' },
  { label: 'Custom',                    url: '' },
]
const MODELS = [
  'qwen2.5:3b', 'qwen2.5:7b', 'qwen2.5:1.5b',
  'Qwen/Qwen2.5-7B-Instruct', 'Qwen/Qwen2.5-3B-Instruct',
  'meta-llama/Llama-3.2-3B-Instruct', 'gpt-4o-mini',
]

export function RunWithLlm() {
  const [endpoint, setEndpoint] = useState(PRESET_ENDPOINTS[0].label)
  const [customUrl, setCustomUrl] = useState('')
  const [apiKey, setApiKey] = useState('')
  const [model, setModel] = useState(MODELS[0])
  const [tier, setTier] = useState('T0')
  const [seed, setSeed] = useState(42)
  const [temperature, setTemperature] = useState(0.7)
  const [maxTurns, setMaxTurns] = useState(10)

  const [running, setRunning] = useState(false)
  const [turns, setTurns] = useState([])
  const [header, setHeader] = useState(null)
  const [done, setDone] = useState(null)
  const stopRef = useRef(null)

  function run() {
    setTurns([]); setHeader(null); setDone(null); setRunning(true)
    const preset = PRESET_ENDPOINTS.find(p => p.label === endpoint)
    const base_url = (customUrl || preset.url).replace(/\/$/, '')

    if (stopRef.current) stopRef.current()
    stopRef.current = llmRunStream(
      {
        base_url, api_key: apiKey, model,
        tier, seed, temperature, max_turns: maxTurns,
      },
      (ev) => {
        if (ev.kind === 'header') {
          setHeader(ev)
        } else if (ev.kind === 'turn') {
          // SSE event uses `kind` for event type and `kind_of` for action kind;
          // TurnCard expects `kind` = action kind, so rename here.
          setTurns(prev => [...prev, { ...ev, kind: ev.kind_of }])
        } else if (ev.kind === 'done') {
          setDone(ev)
          setRunning(false)
        } else if (ev.kind === 'error') {
          setHeader(h => ({ ...(h || {}), error: ev.message }))
          setRunning(false)
        }
      },
    )
  }

  return (
    <div className="grid grid-cols-[1fr_320px] gap-5">
      {/* ─── MAIN PANE ─── */}
      <div className="space-y-5">
        <div className="card">
          <div className="flex items-center gap-2 mb-3">
            <Terminal size={18} className="text-subtle" />
            <h3>Transcript</h3>
          </div>
          {!header && !turns.length && (
            <p className="text-muted italic text-sm">
              Configure the LLM on the right and hit <strong>Run episode</strong>.
              Each turn streams here as the model plays.
            </p>
          )}
          {header && (
            <div className="border-b border-border/50 pb-3 mb-3">
              <div className="text-sm text-muted">
                Model <span className="chip border-border text-ink">{header.model}</span>
                &nbsp;via&nbsp;
                <span className="chip border-border text-ink">{header.base_url}</span>
              </div>
              <div className="text-sm mt-2">
                <strong>Landscape:</strong> {header.landscape}
              </div>
              <div className="text-sm text-muted mt-1">
                Dim: <strong>{header.dim}</strong> · Initial budget:{' '}
                <strong>{header.budget}</strong>
              </div>
              {header.error && (
                <div className="mt-2 p-3 rounded border border-bad/40 bg-bad/10 text-bad">
                  {header.error}
                </div>
              )}
            </div>
          )}
          <div className="space-y-3">
            {turns.map((t, i) => <TurnCard key={i} {...t} />)}
          </div>
        </div>

        {done && <EpisodeDone done={done} />}
      </div>

      {/* ─── SIDEBAR ─── */}
      <aside className="card space-y-4 h-fit sticky top-4">
        <div>
          <h3 className="mb-1">Connect an LLM</h3>
          <p className="text-xs text-muted">
            Point at any OpenAI-compatible <code>/v1/chat/completions</code> endpoint.
          </p>
        </div>

        <div>
          <label className="label">Endpoint</label>
          <select
            className="input"
            value={endpoint}
            onChange={e => setEndpoint(e.target.value)}
          >
            {PRESET_ENDPOINTS.map(p => (
              <option key={p.label} value={p.label}>{p.label}</option>
            ))}
          </select>
        </div>

        <div>
          <label className="label">Model</label>
          <input
            className="input font-mono text-xs"
            list="model-list"
            value={model}
            onChange={e => setModel(e.target.value)}
          />
          <datalist id="model-list">
            {MODELS.map(m => <option key={m} value={m} />)}
          </datalist>
        </div>

        <div>
          <label className="label">Custom base URL (optional)</label>
          <input
            className="input font-mono text-xs"
            placeholder="http://localhost:8080/v1"
            value={customUrl}
            onChange={e => setCustomUrl(e.target.value)}
          />
        </div>

        <div>
          <label className="label">API key (optional)</label>
          <input
            type="password"
            className="input font-mono text-xs"
            placeholder="Bearer <key>"
            value={apiKey}
            onChange={e => setApiKey(e.target.value)}
          />
        </div>

        <hr className="border-border/50" />

        <h3>Episode config</h3>

        <div>
          <label className="label">Tier</label>
          <select className="input" value={tier} onChange={e => setTier(e.target.value)}>
            <option>T0</option><option>T1</option><option>T2</option>
          </select>
        </div>

        <RangeRow label={`Seed: ${seed}`} min={0} max={100} step={1}
                   value={seed} onChange={setSeed} />
        <RangeRow label={`Temperature: ${temperature.toFixed(2)}`}
                   min={0} max={1.5} step={0.05}
                   value={temperature} onChange={setTemperature} />
        <RangeRow label={`Max turns: ${maxTurns}`} min={3} max={15} step={1}
                   value={maxTurns} onChange={setMaxTurns} />

        <button
          className="btn-primary w-full py-3"
          disabled={running}
          onClick={run}
        >
          {running
            ? <><Activity size={16} className="animate-pulse" /> Running…</>
            : <><Play size={16} /> Run episode</>}
        </button>
      </aside>
    </div>
  )
}

function RangeRow({ label, min, max, step, value, onChange }) {
  return (
    <div>
      <label className="label">{label}</label>
      <input
        type="range"
        min={min} max={max} step={step} value={value}
        onChange={e => onChange(Number(e.target.value))}
        className="w-full accent-accent"
      />
    </div>
  )
}

function EpisodeDone({ done }) {
  const reward = done.reward
  const speedup = done.speedup_vs_adam
  const myProg = done.my_progress
  const adamProg = done.adam_progress

  const rewardTone =
    reward >= 0.5 ? 'good' : reward >= 0 ? 'warn' : 'bad'

  // --- Speedup display ---
  // "Speedup" only makes sense when both optimizers descended. Handle the
  // degenerate cases cleanly rather than showing "-0.44×" (mathematically
  // correct, semantically nonsense).
  let speedupDisplay, speedupTone, speedupSub
  if (myProg < 0) {
    // Our optimizer went uphill
    speedupDisplay = 'diverged'
    speedupTone = 'bad'
    speedupSub = `f moved +${Math.abs(myProg).toFixed(2)} (wrong direction)`
  } else if (adamProg <= 0) {
    // Adam itself couldn't descend — unfair denominator
    speedupDisplay = myProg > 0 ? '∞' : '—'
    speedupTone = myProg > 0 ? 'good' : 'warn'
    speedupSub = 'Adam made no progress on this landscape'
  } else {
    const f = speedup < 100 ? speedup.toFixed(2) : Math.round(speedup).toString()
    speedupDisplay = `${f}×`
    speedupTone = speedup >= 1.0 ? 'good' : 'warn'
    speedupSub = `descent ${myProg.toFixed(2)} vs Adam ${adamProg.toFixed(2)}`
  }

  // --- Verdict ---
  let verdict, verdictTone, verdictSub
  if (myProg < 0) {
    verdict = 'Diverged'
    verdictTone = 'bad'
    verdictSub = 'optimizer moved away from the minimum'
  } else if (adamProg <= 0) {
    verdict = myProg > 0 ? 'Succeeds where Adam fails' : 'Tied · both stuck'
    verdictTone = myProg > 0 ? 'good' : 'warn'
    verdictSub = `you: ${myProg.toFixed(2)}, Adam: ${adamProg.toFixed(2)}`
  } else if (speedup >= 1.5) {
    verdict = 'Beats Adam'
    verdictTone = 'good'
    verdictSub = `${((speedup - 1) * 100).toFixed(0)}% further than Adam`
  } else if (speedup >= 1.1) {
    verdict = 'Edges Adam'
    verdictTone = 'good'
    verdictSub = `${((speedup - 1) * 100).toFixed(0)}% further than Adam`
  } else if (speedup >= 0.9) {
    verdict = 'Matches Adam'
    verdictTone = 'warn'
    verdictSub = 'within ±10% of Adam'
  } else {
    verdict = 'Behind Adam'
    verdictTone = 'bad'
    verdictSub = `covered ${(speedup * 100).toFixed(0)}% of Adam's descent`
  }

  return (
    <div className="card"
         style={{ background: 'linear-gradient(180deg, rgba(226,135,99,0.07) 0%, rgba(42,40,36,0) 60%)' }}>
      <div className="flex items-baseline gap-3 mb-4">
        <span className="chip border-accent text-accent uppercase tracking-wider text-[0.7rem]">
          Episode complete
        </span>
        <span className="text-subtle text-sm">
          ended by <code className="text-muted">{done.reason}</code>
        </span>
      </div>

      <div className="grid grid-cols-3 gap-3 mb-5">
        <KpiCard
          label="Reward"
          value={reward.toFixed(3)}
          sub="GRPO scalar · range [−1.5, +1.8]"
          tone={rewardTone}
          sign={reward >= 0 ? '+' : ''} />
        <KpiCard
          label="Speedup vs Adam"
          value={speedupDisplay}
          sub={speedupSub}
          tone={speedupTone} />
        <KpiCard
          label="Verdict"
          value={verdict}
          sub={verdictSub}
          tone={verdictTone} />
      </div>

      <RewardBreakdown breakdown={done.breakdown} total={reward} />
    </div>
  )
}