boocode/apps/web/src/components/control/PlaygroundTab.tsx

import { useState, useRef, useEffect, useCallback } from 'react';
import { cn } from '@/lib/utils';
import { Send, Loader2, Swords, Sparkles } from 'lucide-react';

interface PlaygroundTabProps {
  providerIds: string[];
}

interface ModelEntry {
  id: string;
  providerId: string;
}

interface ChatMessage {
  role: 'user' | 'assistant' | 'system';
  content: string;
}

export function PlaygroundTab({ providerIds }: PlaygroundTabProps) {
  const [models, setModels] = useState<ModelEntry[]>([]);
  const [selectedModel, setSelectedModel] = useState<string>('');
  const [selectedProvider, setSelectedProvider] = useState<string>('');
  const [temperature, setTemperature] = useState(0.7);
  const [topP, setTopP] = useState(0.9);
  const [maxTokens, setMaxTokens] = useState(1024);
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [input, setInput] = useState('');
  const [streaming, setStreaming] = useState(false);
  const [abMode, setAbMode] = useState(false);
  const [modelB, setModelB] = useState('');
  const [providerB, setProviderB] = useState('');
  const [responseA, setResponseA] = useState('');
  const [responseB, setResponseB] = useState('');
  const [streamingAb, setStreamingAb] = useState(false);
  const messagesEndRef = useRef<HTMLDivElement>(null);

  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages]);

  useEffect(() => {
    fetchModels();
  }, []);

  const fetchModels = useCallback(async () => {
    try {
      const res = await fetch('/api/control/playground/models');
      if (!res.ok) return;
      const data = await res.json() as { models: Array<{ providerId: string; models: string[] }> };
      const flattened: ModelEntry[] = [];
      for (const group of data.models) {
        for (const m of group.models) {
          flattened.push({ id: m, providerId: group.providerId });
        }
      }
      setModels(flattened);
      if (flattened.length > 0 && !selectedModel) {
        const first = flattened[0];
        if (first) {
          setSelectedModel(first.id);
          setSelectedProvider(first.providerId);
        }
      }
    } catch {
      // silent
    }
  }, [selectedModel]);

  const groupedModels = models.reduce((acc, m) => {
    if (!acc[m.providerId]) {
      acc[m.providerId] = [];
    }
    const group = acc[m.providerId];
    if (group) {
      group.push(m);
    }
    return acc;
  }, {} as Record<string, ModelEntry[]>);

  const handleSend = async () => {
    if (!input.trim() || !selectedModel || streaming) return;

    const userMsg: ChatMessage = { role: 'user', content: input.trim() };
    const newMessages = [...messages, userMsg, { role: 'assistant' as const, content: '' }];
    setMessages(newMessages);
    setInput('');
    setStreaming(true);

    try {
      const chatMessages = newMessages.slice(0, -1).map((m) => ({
        role: m.role,
        content: m.content,
      }));
      const res = await fetch('/api/control/playground/chat', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          providerId: selectedProvider,
          model: selectedModel,
          messages: chatMessages,
          temperature,
          topP,
          maxTokens,
        }),
      });

      if (!res.ok) {
        const err = await res.json().catch(() => ({}));
        setMessages((prev) => [...prev.slice(0, -1), { role: 'assistant', content: `Error: ${err.error || 'Request failed'}` }]);
        setStreaming(false);
        return;
      }

      const reader = res.body?.getReader();
      if (!reader) {
        setStreaming(false);
        return;
      }

      const decoder = new TextDecoder();
      let buffer = '';
      let assistantContent = '';

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop() ?? '';

        for (const line of lines) {
          const trimmed = line.trim();
          if (!trimmed) continue;
          if (trimmed === 'data: [DONE]') continue;

          const jsonStr = trimmed.startsWith('data: ') ? trimmed.slice(6) : trimmed;
          try {
            const parsed = JSON.parse(jsonStr);
            const delta = parsed.choices?.[0]?.delta?.content;
            if (delta) {
              assistantContent += delta;
              setMessages((prev) => {
                const updated = [...prev];
                updated[updated.length - 1] = { role: 'assistant', content: assistantContent };
                return updated;
              });
            }
          } catch {
            // skip
          }
        }
      }

      setStreaming(false);
    } catch (err) {
      const msg = (err as Error).message ?? String(err);
      setMessages((prev) => [...prev.slice(0, -1), { role: 'assistant', content: `Error: ${msg}` }]);
      setStreaming(false);
    }
  };

  const handleABCompare = async () => {
    if (!input.trim() || !selectedModel || !modelB || streamingAb) return;

    const userMsg: ChatMessage = { role: 'user', content: input.trim() };
    setMessages([...messages, userMsg]);
    setInput('');
    setResponseA('');
    setResponseB('');
    setStreamingAb(true);

    try {
      const res = await fetch('/api/control/playground/chat-ab', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          providerIdA: selectedProvider,
          modelA: selectedModel,
          providerIdB: providerB,
          modelB,
          messages: [...messages, userMsg],
          temperature,
          topP,
          maxTokens,
        }),
      });

      if (!res.ok) {
        setStreamingAb(false);
        return;
      }

      const reader = res.body?.getReader();
      if (!reader) {
        setStreamingAb(false);
        return;
      }

      const decoder = new TextDecoder();
      let buffer = '';

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop() ?? '';

        for (const line of lines) {
          const trimmed = line.trim();
          if (!trimmed) continue;

          const jsonStr = trimmed.startsWith('data: ') ? trimmed.slice(6) : trimmed;
          try {
            const parsed = JSON.parse(jsonStr);
            if (parsed.done) {
              if (parsed.lane === 'A') setStreamingAb((p) => {
                // Check if B is also done elsewhere
                return p;
              });
              continue;
            }
            if (parsed.raw) {
              const innerStr = parsed.raw.startsWith('data: ') ? parsed.raw.slice(6) : parsed.raw;
              const inner = JSON.parse(innerStr);
              const delta = inner.choices?.[0]?.delta?.content;
              if (delta) {
                if (parsed.lane === 'A') {
                  setResponseA((p) => p + delta);
                } else {
                  setResponseB((p) => p + delta);
                }
              }
            }
          } catch {
            // skip
          }
        }
      }

      setStreamingAb(false);
    } catch {
      setStreamingAb(false);
    }
  };

  const getArenaBattleUrl = () => {
    const prompt = encodeURIComponent(input || messages[messages.length - 1]?.content || '');
    const modelA = encodeURIComponent(selectedModel);
    const modelBParam = encodeURIComponent(modelB || '');
    return `/arena?prompt=${prompt}&models=${modelA},${modelBParam}`;
  };

  return (
    <div className="flex flex-col flex-1 min-h-0">
      {/* Model and param controls */}
      <div className="flex flex-wrap items-center gap-3 px-4 py-3 border-b border-border/40 shrink-0">
        <div className="flex items-center gap-2">
          <label className="text-xs text-muted-foreground">Host</label>
          <select
            value={selectedProvider}
            onChange={(e) => {
              setSelectedProvider(e.target.value);
              const firstModel = groupedModels[e.target.value]?.[0]?.id;
              if (firstModel) setSelectedModel(firstModel);
            }}
            className="bg-muted/50 border border-border/50 rounded px-2 py-1 text-sm"
          >
            {Object.keys(groupedModels).map((pid) => (
              <option key={pid} value={pid}>{pid}</option>
            ))}
          </select>
        </div>

        <div className="flex items-center gap-2">
          <label className="text-xs text-muted-foreground">Model</label>
          <select
            value={selectedModel}
            onChange={(e) => setSelectedModel(e.target.value)}
            className="bg-muted/50 border border-border/50 rounded px-2 py-1 text-sm min-w-[200px]"
          >
            {(groupedModels[selectedProvider] ?? []).map((m) => (
              <option key={m.id} value={m.id}>{m.id}</option>
            ))}
          </select>
        </div>

        <div className="flex items-center gap-2">
          <label className="text-xs text-muted-foreground">Temp</label>
          <input
            type="number"
            min={0}
            max={2}
            step={0.1}
            value={temperature}
            onChange={(e) => setTemperature(parseFloat(e.target.value) || 0.7)}
            className="w-16 bg-muted/50 border border-border/50 rounded px-2 py-1 text-sm"
          />
        </div>

        <div className="flex items-center gap-2">
          <label className="text-xs text-muted-foreground">Top P</label>
          <input
            type="number"
            min={0}
            max={1}
            step={0.05}
            value={topP}
            onChange={(e) => setTopP(parseFloat(e.target.value) || 0.9)}
            className="w-16 bg-muted/50 border border-border/50 rounded px-2 py-1 text-sm"
          />
        </div>

        <div className="flex items-center gap-2">
          <label className="text-xs text-muted-foreground">Max</label>
          <input
            type="number"
            min={1}
            max={8192}
            step={128}
            value={maxTokens}
            onChange={(e) => setMaxTokens(parseInt(e.target.value) || 1024)}
            className="w-20 bg-muted/50 border border-border/50 rounded px-2 py-1 text-sm"
          />
        </div>

        <button
          type="button"
          onClick={() => setAbMode(!abMode)}
          className={cn(
            'flex items-center gap-1 px-2 py-1 text-xs rounded transition-colors',
            abMode
              ? 'bg-accent/20 text-accent border border-accent/30'
              : 'text-muted-foreground hover:text-foreground border border-transparent'
          )}
        >
          <Swords className="size-3" />
          A/B
        </button>
      </div>

      {/* A/B model B selector */}
      {abMode && (
        <div className="flex items-center gap-3 px-4 py-2 border-b border-border/40 bg-muted/20">
          <label className="text-xs text-muted-foreground">Model B</label>
          <select
            value={providerB}
            onChange={(e) => {
              setProviderB(e.target.value);
              const firstModel = groupedModels[e.target.value]?.[0]?.id;
              if (firstModel) setModelB(firstModel);
            }}
            className="bg-muted/50 border border-border/50 rounded px-2 py-1 text-sm"
          >
            {Object.keys(groupedModels).map((pid) => (
              <option key={pid} value={pid}>{pid}</option>
            ))}
          </select>
          <select
            value={modelB}
            onChange={(e) => setModelB(e.target.value)}
            className="bg-muted/50 border border-border/50 rounded px-2 py-1 text-sm min-w-[200px]"
          >
            {(groupedModels[providerB] ?? []).map((m) => (
              <option key={m.id} value={m.id}>{m.id}</option>
            ))}
          </select>
        </div>
      )}

      {/* Chat area */}
      <div className="flex-1 flex flex-col min-h-0 overflow-hidden">
        {!abMode ? (
          <>
            {/* Messages */}
            <div className="flex-1 overflow-y-auto px-4 py-3 space-y-3">
              {messages.map((msg, i) => (
                <div
                  key={i}
                  className={cn(
                    'max-w-[80%] rounded-lg px-3 py-2 text-sm',
                    msg.role === 'user'
                      ? 'ml-auto bg-accent/20 text-accent-foreground'
                      : 'bg-muted/50 text-foreground'
                  )}
                >
                  {msg.content || (msg.role === 'assistant' && streaming ? <Loader2 className="size-4 animate-spin" /> : null)}
                </div>
              ))}
              <div ref={messagesEndRef} />
            </div>

            {/* Input */}
            <div className="flex items-center gap-2 px-4 py-3 border-t border-border/40 shrink-0">
              <textarea
                value={input}
                onChange={(e) => setInput(e.target.value)}
                onKeyDown={(e) => {
                  if (e.key === 'Enter' && !e.shiftKey) {
                    e.preventDefault();
                    handleSend();
                  }
                }}
                placeholder="Type a message..."
                className="flex-1 bg-muted/50 border border-border/50 rounded-lg px-3 py-2 text-sm resize-none min-h-[40px] max-h-[120px]"
                rows={1}
              />
              <button
                type="button"
                onClick={handleSend}
                disabled={streaming || !input.trim()}
                className={cn(
                  'p-2 rounded-lg transition-colors',
                  streaming || !input.trim()
                    ? 'text-muted-foreground/50'
                    : 'bg-accent/20 text-accent hover:bg-accent/30'
                )}
              >
                {streaming ? <Loader2 className="size-4 animate-spin" /> : <Send className="size-4" />}
              </button>
            </div>
          </>
        ) : (
          <>
            {/* A/B comparison */}
            <div className="flex-1 flex gap-2 px-4 py-3 min-h-0 overflow-hidden">
              <div className="flex-1 flex flex-col min-h-0 bg-muted/20 rounded-lg border border-border/30 overflow-hidden">
                <div className="px-3 py-1.5 text-xs font-medium text-muted-foreground border-b border-border/30 shrink-0">
                  Model A: {selectedModel}
                </div>
                <div className="flex-1 overflow-y-auto px-3 py-2 text-sm whitespace-pre-wrap">
                  {responseA || (streamingAb ? <Loader2 className="size-4 animate-spin" /> : 'Waiting...')}
                </div>
              </div>
              <div className="flex-1 flex flex-col min-h-0 bg-muted/20 rounded-lg border border-border/30 overflow-hidden">
                <div className="px-3 py-1.5 text-xs font-medium text-muted-foreground border-b border-border/30 shrink-0">
                  Model B: {modelB}
                </div>
                <div className="flex-1 overflow-y-auto px-3 py-2 text-sm whitespace-pre-wrap">
                  {responseB || (streamingAb ? <Loader2 className="size-4 animate-spin" /> : 'Waiting...')}
                </div>
              </div>
            </div>

            {/* A/B input */}
            <div className="flex items-center gap-2 px-4 py-3 border-t border-border/40 shrink-0">
              <textarea
                value={input}
                onChange={(e) => setInput(e.target.value)}
                onKeyDown={(e) => {
                  if (e.key === 'Enter' && !e.shiftKey) {
                    e.preventDefault();
                    handleABCompare();
                  }
                }}
                placeholder="Type a prompt for A/B comparison..."
                className="flex-1 bg-muted/50 border border-border/50 rounded-lg px-3 py-2 text-sm resize-none min-h-[40px]"
                rows={1}
              />
              <button
                type="button"
                onClick={handleABCompare}
                disabled={streamingAb || !input.trim() || !modelB}
                className={cn(
                  'p-2 rounded-lg transition-colors',
                  streamingAb || !input.trim() || !modelB
                    ? 'text-muted-foreground/50'
                    : 'bg-accent/20 text-accent hover:bg-accent/30'
                )}
              >
                {streamingAb ? <Loader2 className="size-4 animate-spin" /> : <Swords className="size-4" />}
              </button>
            </div>
          </>
        )}
      </div>

      {/* Battle in Arena link */}
      <div className="px-4 py-2 border-t border-border/40 shrink-0">
        <a
          href={getArenaBattleUrl()}
          target="_blank"
          rel="noopener noreferrer"
          className="flex items-center gap-1.5 text-xs text-muted-foreground hover:text-foreground transition-colors"
        >
          <Sparkles className="size-3" />
          Battle in Arena
        </a>
      </div>
    </div>
  );
}