frontend/src/arena/components/TranscriptPanel.tsx

import { useState } from "react";
import type { EpisodeData, TranscriptTurn } from "../types";
import { highlightTrace } from "../utils/traceHighlight";

export interface DragHandleProps {
  draggable: true;
  onDragStart: (e: React.DragEvent) => void;
  onDragEnd: (e: React.DragEvent) => void;
}

interface TranscriptPanelProps {
  datasetName: string;
  repoName?: string;
  data: EpisodeData | undefined;
  dragHandleProps?: DragHandleProps;
}

const OUTCOME_STYLES = {
  win: { bg: "bg-green-900", text: "text-green-300", label: "WIN" },
  loss: { bg: "bg-red-900", text: "text-red-300", label: "LOSS" },
  error: { bg: "bg-yellow-900", text: "text-yellow-300", label: "ERROR" },
  unknown: { bg: "bg-gray-700", text: "text-gray-300", label: "?" },
};

const PLAYER_COLORS: Record<number, { bubble: string; label: string; name: string }> = {
  0: { bubble: "bg-purple-900/60 border-purple-700", label: "text-purple-400", name: "Player 0" },
  1: { bubble: "bg-orange-900/60 border-orange-700", label: "text-orange-400", name: "Player 1" },
  2: { bubble: "bg-purple-900/60 border-purple-700", label: "text-purple-400", name: "Player 2" },
  3: { bubble: "bg-teal-900/60 border-teal-700", label: "text-teal-400", name: "Player 3" },
};

function getPlayerColor(playerId: number) {
  return PLAYER_COLORS[playerId] || PLAYER_COLORS[0];
}

export default function TranscriptPanel({ datasetName, repoName, data, dragHandleProps }: TranscriptPanelProps) {
  if (!data) {
    return (
      <div className="h-full border border-gray-700 rounded-lg flex items-center justify-center">
        <div className="text-gray-500 text-sm">No data</div>
      </div>
    );
  }

  const outcomeStyle = OUTCOME_STYLES[data.outcome];
  const borderColor = data.outcome === "win" ? "border-green-600"
    : data.outcome === "loss" ? "border-red-600"
    : data.outcome === "error" ? "border-yellow-600"
    : "border-gray-700";

  return (
    <div className={`h-full border-2 ${borderColor} rounded-lg flex flex-col bg-gray-900/50`}>
      {/* Header */}
      <div className="px-3 py-2 border-b border-gray-700 shrink-0">
        <div className="flex items-center justify-between mb-1">
          <div className="flex items-center gap-2 min-w-0">
            <span className="text-sm font-semibold text-gray-200 truncate" title={repoName || datasetName}>
              {datasetName}
            </span>
            <span className={`px-1.5 py-0.5 text-[10px] rounded font-medium ${outcomeStyle.bg} ${outcomeStyle.text}`}>
              {outcomeStyle.label}
            </span>
          </div>
          <div className="flex items-center gap-1.5 shrink-0 ml-2">
            {dragHandleProps && (
              <span
                {...dragHandleProps}
                title="Drag to reorder"
                className="drag-handle text-gray-600 hover:text-gray-400 transition-colors"
              >
                <svg width="14" height="14" viewBox="0 0 16 16" fill="currentColor">
                  <circle cx="5" cy="3" r="1.5" />
                  <circle cx="11" cy="3" r="1.5" />
                  <circle cx="5" cy="8" r="1.5" />
                  <circle cx="11" cy="8" r="1.5" />
                  <circle cx="5" cy="13" r="1.5" />
                  <circle cx="11" cy="13" r="1.5" />
                </svg>
              </span>
            )}
          </div>
        </div>
        <div className="flex items-center gap-3 text-[10px] text-gray-500">
          <span>{data.model}</span>
          <span>{data.num_turns} turns</span>
          {data.reward !== null && <span>reward: {data.reward}</span>}
          {data.opponent_model && <span>vs {data.opponent_model}</span>}
        </div>
      </div>

      {/* Error banner */}
      {data.error && (
        <div className="px-3 py-1.5 bg-red-900/30 border-b border-red-800/50 text-xs text-red-300">
          Error: {data.error}
        </div>
      )}

      {/* Transcript chat */}
      <div className="flex-1 overflow-y-auto transcript-scroll px-3 py-2 space-y-3">
        {/* System prompt */}
        {data.system_prompt && (
          <SystemPromptBubble text={data.system_prompt} />
        )}
        {data.transcript.map((turn, i) => (
          <TurnBubble
            key={i}
            turn={turn}
            turnIndex={i}
            allTurns={data.transcript}
            systemPrompt={data.system_prompt}
            hasMultiplePlayers={data.opponent_model !== null}
          />
        ))}
      </div>
    </div>
  );
}


function SystemPromptBubble({ text }: { text: string }) {
  const [expanded, setExpanded] = useState(true);

  return (
    <div className="mb-2">
      <button
        onClick={() => setExpanded(!expanded)}
        className="flex items-center gap-1.5 text-[10px] text-purple-400 hover:text-purple-300 transition-colors mb-1 font-semibold uppercase tracking-wider"
      >
        <span>{expanded ? "\u25BC" : "\u25B6"}</span>
        <span>System Prompt</span>
      </button>
      {expanded && (
        <div className="bg-purple-950/40 border border-purple-800/50 rounded-lg px-3 py-2">
          <pre className="text-xs leading-relaxed whitespace-pre-wrap font-mono text-purple-200">
            {text}
          </pre>
        </div>
      )}
    </div>
  );
}


interface TurnBubbleProps {
  turn: TranscriptTurn;
  turnIndex: number;
  allTurns: TranscriptTurn[];
  systemPrompt: string | null;
  hasMultiplePlayers: boolean;
}

function buildRawPrompt(
  turnIndex: number,
  allTurns: TranscriptTurn[],
  systemPrompt: string | null,
): object[] {
  const messages: object[] = [];
  if (systemPrompt) {
    messages.push({ role: "system", content: systemPrompt });
  }
  for (let i = 0; i < turnIndex; i++) {
    messages.push({ role: "user", content: allTurns[i].observation });
    messages.push({ role: "assistant", content: allTurns[i].action });
  }
  messages.push({ role: "user", content: allTurns[turnIndex].observation });
  return messages;
}

function TurnBubble({ turn, turnIndex, allTurns, systemPrompt, hasMultiplePlayers }: TurnBubbleProps) {
  const [thinkExpanded, setThinkExpanded] = useState(false);
  const [rawPromptExpanded, setRawPromptExpanded] = useState(false);
  const playerColor = getPlayerColor(turn.player_id);
  const thinkSegments = highlightTrace(turn.think_text);

  return (
    <div>
      {/* Turn number marker */}
      <div className="text-[10px] text-gray-600 mb-1">Turn {turn.turn}</div>

      {/* Observation (environment message) — left aligned */}
      {turn.observation && (
        <div className="flex justify-start mb-1.5">
          <div className="max-w-[90%]">
            <div className="text-[10px] text-gray-500 mb-0.5 font-semibold uppercase tracking-wider">ENV</div>
            <div className="bg-gray-800 border border-gray-700 rounded-lg rounded-tl-none px-3 py-2">
              <pre className="text-xs leading-relaxed whitespace-pre-wrap font-mono text-gray-300">
                {turn.observation}
              </pre>
            </div>
          </div>
        </div>
      )}

      {/* Action (model response) — right aligned */}
      <div className="flex justify-end">
        <div className="max-w-[90%]">
          <div className={`text-[10px] mb-0.5 font-semibold uppercase tracking-wider text-right ${playerColor.label}`}>
            {hasMultiplePlayers ? `${playerColor.name} (${turn.player_id === 0 ? "model" : "opponent"})` : "Model"}
          </div>

          <div className={`border rounded-lg rounded-tr-none px-3 py-2 ${playerColor.bubble}`}>
            {/* Raw Prompt — collapsible, shows full messages sent to API */}
            <div className="mb-2">
              <button
                onClick={() => setRawPromptExpanded(!rawPromptExpanded)}
                className="flex items-center gap-1 text-[10px] text-gray-500 hover:text-gray-400 transition-colors mb-1"
              >
                <span>{rawPromptExpanded ? "\u25BC" : "\u25B6"}</span>
                <span>Raw Prompt ({turnIndex * 2 + 1 + (systemPrompt ? 1 : 0)} messages)</span>
              </button>
              {rawPromptExpanded && (
                <pre className="text-[10px] leading-relaxed whitespace-pre-wrap font-mono bg-gray-950/80 border border-gray-800 rounded px-2 py-1.5 text-gray-400 max-h-80 overflow-y-auto">
                  {JSON.stringify(buildRawPrompt(turnIndex, allTurns, systemPrompt), null, 2)}
                </pre>
              )}
            </div>

            {/* Think section — collapsible */}
            {turn.think_text && (
              <div className="mb-2">
                <button
                  onClick={() => setThinkExpanded(!thinkExpanded)}
                  className="flex items-center gap-1 text-[10px] text-gray-500 hover:text-gray-400 transition-colors mb-1"
                >
                  <span>{thinkExpanded ? "\u25BC" : "\u25B6"}</span>
                  <span>Thinking ({turn.think_len.toLocaleString()} chars{turn.backtracks > 0 ? `, ${turn.backtracks} backtracks` : ""})</span>
                </button>
                {thinkExpanded && (
                  <pre className="text-xs leading-relaxed whitespace-pre-wrap font-mono border-l-2 border-gray-600 pl-2 mt-1">
                    {thinkSegments.map((seg, i) => (
                      <span key={i} className={seg.className}>{seg.text}</span>
                    ))}
                  </pre>
                )}
              </div>
            )}

            {/* Action text (the actual game move) */}
            <pre className="text-xs leading-relaxed whitespace-pre-wrap font-mono text-gray-100 font-medium">
              {turn.action_text || turn.action || "(no action)"}
            </pre>
          </div>
        </div>
      </div>
    </div>
  );
}