Prompt Playground

"use client";

import * as React from "react";
import { Slider } from "radix-ui";
import { ArrowsRightLeftIcon, BeakerIcon, BoltIcon, PlayIcon, SparklesIcon, StopIcon } from "@heroicons-animated/react";

import { cn } from "@/lib/utils";
import { Button } from "@/components/ui/button";
import { MessageMarkdown } from "@/components/pandacoderz-ui/message";
import { ModelSelector, ModelSelectorContent, ModelSelectorRadioGroup, ModelSelectorRadioItem, ModelSelectorTrigger } from "@/components/pandacoderz-ui/model-selector";
import { TextShimmer } from "@/components/pandacoderz-ui/text-shimmer";
import { UsageMeter } from "@/components/pandacoderz-ui/usage-meter";
import { sleep, streamText } from "@/lib/mock-stream";

type ModelDef = { value: string; title: string; description: string; inputPer1M: number; outputPer1M: number; speed: number; icon: React.ComponentType<{ className?: string }> };

const Sparkles = ({ className }: { className?: string }) => <SparklesIcon size={16} className={className} />;
const Bolt = ({ className }: { className?: string }) => <BoltIcon size={16} className={className} />;

const models: ModelDef[] = [
  { value: "claude-fable-5-1", title: "Claude Fable 5.1", description: "Most capable", inputPer1M: 15, outputPer1M: 75, speed: 22, icon: Sparkles },
  { value: "claude-opus-5", title: "Claude Opus 5", description: "Deep reasoning", inputPer1M: 15, outputPer1M: 75, speed: 26, icon: Sparkles },
  { value: "claude-sonnet-5", title: "Claude Sonnet 5", description: "Balanced", inputPer1M: 3, outputPer1M: 15, speed: 14, icon: Bolt },
  { value: "claude-haiku-4-5", title: "Claude Haiku 4.5", description: "Fastest", inputPer1M: 0.8, outputPer1M: 4, speed: 8, icon: Bolt },
];

const outputs: Record<string, string> = {
  "claude-fable-5-1": `**Subject:** Your workspace is ready\n\nHi Maya,\n\nYour team workspace is live. Three things worth doing first:\n\n1. Invite a teammate so you can see shared threads.\n2. Connect a data source; most teams start with their docs.\n3. Pin the assistant to your sidebar for one-click access.\n\nReply to this email if anything feels off. A person reads every message.\n\n— The team`,
  "claude-opus-5": `**Subject:** Welcome aboard, Maya\n\nYour workspace is set up and ready. To get value quickly:\n\n- **Invite your team.** Shared context makes the assistant far more useful.\n- **Connect docs.** Answers get grounded in your own material.\n- **Try a block.** Drop the chat into any page in under a minute.\n\nWe're here if you need us.`,
  "claude-sonnet-5": `**Subject:** You're in\n\nHi Maya, your workspace is ready. Start by inviting a teammate, connecting a docs source, and pinning the assistant to your sidebar. Reply any time; a real person reads every message.`,
  "claude-haiku-4-5": `**Subject:** Workspace ready\n\nHi Maya, you're all set. Invite your team, connect your docs, and pin the assistant. Reply with any questions.`,
};

type Result = { text: string; status: "idle" | "running" | "done"; latencyMs: number; inputTokens: number; outputTokens: number };
const emptyResult: Result = { text: "", status: "idle", latencyMs: 0, inputTokens: 0, outputTokens: 0 };

function Param({ label, value, min, max, step, onChange, format = (v) => String(v) }: { label: string; value: number; min: number; max: number; step: number; onChange: (v: number) => void; format?: (v: number) => string }) {
  return (
    <div className="flex flex-col gap-2">
      <div className="flex items-center justify-between text-xs"><span className="text-muted-foreground">{label}</span><span className="font-mono tabular-nums">{format(value)}</span></div>
      <Slider.Root value={[value]} min={min} max={max} step={step} onValueChange={([v]) => onChange(v)} className="relative flex h-4 w-full touch-none select-none items-center">
        <Slider.Track className="relative h-1 w-full grow overflow-hidden rounded-full bg-muted"><Slider.Range className="absolute h-full bg-brand" /></Slider.Track>
        <Slider.Thumb aria-label={label} className="block size-3.5 rounded-full border border-brand bg-background shadow-xs outline-none focus-visible:ring-[3px] focus-visible:ring-ring/50" />
      </Slider.Root>
    </div>
  );
}

function ModelPicker({ value, onChange }: { value: string; onChange: (v: string) => void }) {
  return (
    <ModelSelector value={value} onValueChange={onChange} items={models}>
      <ModelSelectorTrigger variant="ghost" className="h-8 text-xs" />
      <ModelSelectorContent className="w-72" align="start">
        <ModelSelectorRadioGroup value={value} onValueChange={onChange}>
          {models.map((m) => <ModelSelectorRadioItem key={m.value} value={m.value} title={m.title} description={m.description} icon={m.icon} />)}
        </ModelSelectorRadioGroup>
      </ModelSelectorContent>
    </ModelSelector>
  );
}

export type PromptPlaygroundProps = { className?: string };

export default function PromptPlayground({ className }: PromptPlaygroundProps) {
  const [system, setSystem] = React.useState("You are a concise product writer. Keep emails under 120 words and end with a clear next step.");
  const [prompt, setPrompt] = React.useState("Write a welcome email for a new user named Maya who just created a team workspace.");
  const [temperature, setTemperature] = React.useState(0.7);
  const [maxTokens, setMaxTokens] = React.useState(512);
  const [topP, setTopP] = React.useState(0.95);
  const [modelA, setModelA] = React.useState(models[0].value);
  const [modelB, setModelB] = React.useState(models[2].value);
  const [results, setResults] = React.useState<[Result, Result]>([emptyResult, emptyResult]);
  const abortRef = React.useRef<AbortController | null>(null);

  const running = results.some((r) => r.status === "running");
  const inputTokens = Math.round((system.length + prompt.length) / 4);

  const runOne = async (idx: 0 | 1, model: string, signal: AbortSignal) => {
    const def = models.find((m) => m.value === model)!;
    const set = (fn: (r: Result) => Result) => setResults((prev) => { const next = [...prev] as [Result, Result]; next[idx] = fn(next[idx]); return next; });
    const started = performance.now();
    set(() => ({ ...emptyResult, status: "running", inputTokens }));
    await sleep(300 + def.speed * 15, signal);
    let out = "";
    for await (const piece of streamText(outputs[model] ?? outputs["claude-sonnet-5"], signal, def.speed)) {
      out += piece;
      set((r) => ({ ...r, text: out, outputTokens: Math.round(out.length / 4), latencyMs: performance.now() - started }));
    }
    set((r) => ({ ...r, status: "done", latencyMs: performance.now() - started }));
  };

  const run = async () => {
    abortRef.current?.abort();
    const controller = new AbortController();
    abortRef.current = controller;
    try {
      await Promise.all([runOne(0, modelA, controller.signal), runOne(1, modelB, controller.signal)]);
    } catch {
      setResults((prev) => prev.map((r) => (r.status === "running" ? { ...r, status: "done" } : r)) as [Result, Result]);
    }
  };
  const stop = () => abortRef.current?.abort();
  React.useEffect(() => () => abortRef.current?.abort(), []);

  const cost = (r: Result, model: string) => {
    const def = models.find((m) => m.value === model)!;
    return (r.inputTokens / 1e6) * def.inputPer1M + (r.outputTokens / 1e6) * def.outputPer1M;
  };

  return (
    <div data-slot="prompt-playground" className={cn("flex h-full min-h-0 w-full flex-col overflow-hidden rounded-3xl border bg-background shadow-xs", className)}>
      <header className="flex h-12 shrink-0 items-center justify-between gap-3 border-b px-4">
        <div className="flex items-center gap-2 text-sm font-medium">
          <span className="flex size-6 items-center justify-center rounded-md bg-brand text-primary-foreground"><BeakerIcon size={14} className="flex" /></span>
          Playground
        </div>
        <div className="flex items-center gap-1">
          <Button size="sm" variant="ghost" className="rounded-full" onClick={() => { setModelA(modelB); setModelB(modelA); }}><ArrowsRightLeftIcon size={14} className="flex" /> Swap</Button>
          {running ? (
            <Button size="sm" variant="secondary" className="rounded-full" onClick={stop}><StopIcon size={14} className="flex" /> Stop</Button>
          ) : (
            <Button size="sm" className="rounded-full" onClick={run}><PlayIcon size={14} className="flex" /> Run both</Button>
          )}
        </div>
      </header>

      <div className="grid min-h-0 flex-1 lg:grid-cols-[minmax(0,18rem)_minmax(0,1fr)]">
        <aside className="flex min-h-0 flex-col gap-5 overflow-y-auto border-b bg-surface p-4 lg:border-r lg:border-b-0">
          <label className="flex flex-col gap-1.5">
            <span className="text-xs font-medium text-muted-foreground">System prompt</span>
            <textarea value={system} onChange={(e) => setSystem(e.target.value)} rows={5} className="resize-none rounded-xl border bg-background p-3 text-sm leading-6 outline-none focus-visible:ring-[3px] focus-visible:ring-ring/50" />
          </label>
          <div className="flex flex-col gap-4">
            <Param label="Temperature" value={temperature} min={0} max={1} step={0.05} onChange={setTemperature} format={(v) => v.toFixed(2)} />
            <Param label="Max tokens" value={maxTokens} min={64} max={4096} step={64} onChange={setMaxTokens} />
            <Param label="Top P" value={topP} min={0} max={1} step={0.01} onChange={setTopP} format={(v) => v.toFixed(2)} />
          </div>
          <UsageMeter label="Context used" value={inputTokens + Math.max(...results.map((r) => r.outputTokens))} max={200_000} showPercent size="sm" />
        </aside>

        <div className="flex min-h-0 flex-col">
          <div className="shrink-0 border-b p-3">
            <label className="flex flex-col gap-1.5">
              <span className="px-1 text-xs font-medium text-muted-foreground">User message</span>
              <textarea value={prompt} onChange={(e) => setPrompt(e.target.value)} rows={2} className="resize-none rounded-xl border bg-background p-3 text-sm leading-6 outline-none focus-visible:ring-[3px] focus-visible:ring-ring/50" />
            </label>
          </div>
          <div className="grid min-h-0 flex-1 md:grid-cols-2">
            {([[modelA, setModelA, results[0]], [modelB, setModelB, results[1]]] as const).map(([model, setModel, result], i) => {
              const def = models.find((m) => m.value === model)!;
              return (
                <div key={i} className={cn("flex min-h-0 flex-col", i === 0 && "border-b md:border-r md:border-b-0")}>
                  <div className="flex h-10 shrink-0 items-center justify-between gap-2 border-b px-2">
                    <ModelPicker value={model} onChange={setModel} />
                    <span className="px-2 text-[11px] text-muted-foreground">${def.inputPer1M}/{def.outputPer1M} per 1M</span>
                  </div>
                  <div className="min-h-0 flex-1 overflow-y-auto p-4">
                    {result.status === "idle" ? (
                      <p className="text-sm text-muted-foreground">Run to see output from {def.title}.</p>
                    ) : result.text ? (
                      <MessageMarkdown isAnimating={result.status === "running"}>{result.text}</MessageMarkdown>
                    ) : (
                      <TextShimmer className="text-sm text-muted-foreground" invertLight>Waiting for first token…</TextShimmer>
                    )}
                  </div>
                  <footer className="flex h-9 shrink-0 items-center gap-4 border-t px-4 font-mono text-[11px] tabular-nums text-muted-foreground">
                    <span>{result.latencyMs ? `${(result.latencyMs / 1000).toFixed(1)}s` : "—"}</span>
                    <span>{result.inputTokens} in</span>
                    <span>{result.outputTokens} out</span>
                    <span className="ml-auto text-foreground">${cost(result, model).toFixed(4)}</span>
                  </footer>
                </div>
              );
            })}
          </div>
        </div>
      </div>
    </div>
  );
}

/**
 * Helpers for scripted streaming in demos. Blocks use these to fake a model
 * response token by token without touching the network. Swap them for a real
 * fetch in production; the UI code does not change.
 */

export function sleep(ms: number, signal?: AbortSignal) {
  return new Promise<void>((resolve, reject) => {
    if (signal?.aborted) return reject(signal.reason);
    const id = setTimeout(resolve, ms);
    signal?.addEventListener(
      "abort",
      () => {
        clearTimeout(id);
        reject(signal.reason);
      },
      { once: true },
    );
  });
}

/** Split text into word-ish chunks so streaming looks natural. */
export function chunk(text: string): string[] {
  return text.match(/\S+\s*|\s+/g) ?? [text];
}

export function jitter(base: number) {
  return base + Math.random() * base;
}

/** Yield a string piece by piece with a small random delay between pieces. */
export async function* streamText(
  text: string,
  signal?: AbortSignal,
  delayMs = 14,
): AsyncGenerator<string> {
  for (const piece of chunk(text)) {
    if (signal?.aborted) return;
    yield piece;
    await sleep(jitter(delayMs), signal);
  }
}

let counter = 0;
export const uid = (prefix = "id") => `${prefix}-${Date.now().toString(36)}-${++counter}`;