import { useState } from "react";
import type { RlmIterDetail } from "../types";

interface IterationDetailProps {
  data: RlmIterDetail;
}

function parsePromptMessages(promptStr: string): { role: string; content: string }[] {
  try {
    const parsed = JSON.parse(promptStr);
    if (Array.isArray(parsed)) return parsed;
  } catch { /* not JSON */ }
  return [{ role: "raw", content: promptStr }];
}

const roleColors: Record<string, string> = {
  system: "border-violet-500 bg-violet-950",
  user: "border-emerald-500 bg-emerald-950",
  assistant: "border-sky-500 bg-sky-950",
  raw: "border-gray-500 bg-gray-900",
};

export default function IterationDetail({ data }: IterationDetailProps) {
  const [promptExpanded, setPromptExpanded] = useState(false);
  const messages = parsePromptMessages(data.prompt);

  return (
    <div className="space-y-4 border border-gray-700 rounded-lg p-4 bg-gray-900">
      {/* Stats */}
      <div className="flex gap-4 text-xs text-gray-400">
        <span>Model: <span className="text-gray-200">{data.model}</span></span>
        <span>In: <span className="text-emerald-300">{(data.input_tokens / 1000).toFixed(1)}k</span></span>
        <span>Out: <span className="text-emerald-300">{(data.output_tokens / 1000).toFixed(1)}k</span></span>
        <span>Time: <span className="text-gray-200">{data.execution_time.toFixed(1)}s</span></span>
      </div>

      {/* Prompt section (collapsible) */}
      <div>
        <button
          className="flex items-center gap-2 text-sm font-semibold text-gray-300 hover:text-gray-100 mb-2"
          onClick={() => setPromptExpanded(!promptExpanded)}
        >
          <span className={`transform transition-transform ${promptExpanded ? "rotate-90" : ""}`}>
            &#9654;
          </span>
          Prompt ({messages.length} messages)
        </button>
        {promptExpanded && (
          <div className="space-y-2 ml-4">
            {messages.map((msg, i) => (
              <div
                key={i}
                className={`border-l-2 rounded-r-lg px-3 py-2 ${roleColors[msg.role] || roleColors.raw}`}
              >
                <div className="text-xs font-semibold text-gray-400 mb-1 uppercase">{msg.role}</div>
                <div className="text-sm text-gray-200 whitespace-pre-wrap max-h-96 overflow-y-auto">
                  {msg.content.length > 8000 ? msg.content.slice(0, 8000) + "\n...(truncated)" : msg.content}
                </div>
              </div>
            ))}
          </div>
        )}
      </div>

      {/* Response */}
      <div>
        <div className="text-sm font-semibold text-gray-300 mb-2">Response</div>
        <div className="bg-gray-800 border border-gray-700 rounded-lg p-3">
          <div className="text-sm text-gray-200 whitespace-pre-wrap max-h-96 overflow-y-auto font-mono">
            {data.response}
          </div>
        </div>
      </div>

      {/* Code Blocks */}
      {data.code_blocks.length > 0 && (
        <div>
          <div className="text-sm font-semibold text-gray-300 mb-2">
            Code Blocks ({data.code_blocks.length})
          </div>
          <div className="space-y-3">
            {data.code_blocks.map((cb, i) => (
              <div key={i} className="border border-gray-700 rounded-lg overflow-hidden">
                <div className="bg-gray-800 px-3 py-1.5 text-xs text-gray-400 border-b border-gray-700 flex items-center gap-2">
                  <span className="text-emerald-400 font-mono">python</span>
                  <span>Block {i + 1}</span>
                </div>
                <pre className="bg-gray-900 p-3 text-sm text-gray-200 overflow-x-auto font-mono leading-relaxed">
                  {cb.code}
                </pre>
                {cb.stdout && (
                  <div className="border-t border-gray-700">
                    <div className="bg-gray-800 px-3 py-1 text-xs text-gray-400">stdout</div>
                    <pre className="bg-emerald-950 p-3 text-sm text-emerald-200 overflow-x-auto font-mono">
                      {cb.stdout}
                    </pre>
                  </div>
                )}
              </div>
            ))}
          </div>
        </div>
      )}

      {/* Final Answer */}
      {data.final_answer && (
        <div className="bg-emerald-950 border border-emerald-700 rounded-lg p-4">
          <div className="text-xs font-semibold text-emerald-400 mb-2">Final Answer</div>
          <div className="text-sm text-gray-200 whitespace-pre-wrap max-h-96 overflow-y-auto">
            {data.final_answer}
          </div>
        </div>
      )}
    </div>
  );
}