import { useEffect, useState, useMemo } from "react";

// ── Types ─────────────────────────────────────────────────────────────────

type Template = "gpt-oss" | "qwen";

interface Message {
  role: "user" | "assistant";
  content: string;
}

interface Row {
  query_id: string;
  excerpt: string;
  messages_gpt:  Message[] | null;
  messages_qwen: Message[] | null;
}

// ── Color palette ─────────────────────────────────────────────────────────

const COLORS = {
  user:      { bg: "bg-amber-950/40",  border: "border-amber-600",  label: "text-amber-400",  tag: "USER MESSAGE"  },
  reasoning: { bg: "bg-purple-950/40", border: "border-purple-600", label: "text-purple-400", tag: "REASONING"     },
  tool_call: { bg: "bg-blue-950/40",   border: "border-blue-600",   label: "text-blue-400",   tag: "TOOL CALL"     },
  tool_resp: { bg: "bg-gray-800/60",   border: "border-gray-600",   label: "text-gray-400",   tag: "TOOL RESPONSE" },
};

// ── Block component ───────────────────────────────────────────────────────

function Block({ kind, label, children }: {
  kind: keyof typeof COLORS;
  label?: string;
  children: React.ReactNode;
}) {
  const c = COLORS[kind];
  return (
    <div className={`border-l-2 ${c.border} ${c.bg} pl-3 py-2 rounded-r mb-2`}>
      <div className={`text-[9px] font-bold uppercase tracking-widest mb-1.5 ${c.label}`}>
        {label ?? c.tag}
      </div>
      {children}
    </div>
  );
}

// ── Left panel: parse excerpt (newline-separated JSON objects) ────────────

interface ExcerptItem {
  type: "reasoning" | "function_call" | "function_call_output" | string;
  [key: string]: unknown;
}

function parseExcerptItems(excerpt: string): ExcerptItem[] {
  return excerpt
    .split(/\n\n/)
    .map(s => s.trim())
    .filter(Boolean)
    .flatMap(chunk => {
      try {
        const obj = JSON.parse(chunk);
        return typeof obj === "object" && obj !== null && "type" in obj ? [obj as ExcerptItem] : [];
      } catch {
        return [];
      }
    });
}

function ExcerptPanel({ excerpt, userContent }: { excerpt: string; userContent: string }) {
  const items = useMemo(() => parseExcerptItems(excerpt), [excerpt]);

  return (
    <div className="space-y-0">
      <Block kind="user">
        <pre className="text-xs text-gray-200 whitespace-pre-wrap font-mono leading-relaxed">
          {userContent}
        </pre>
      </Block>

      {items.map((item, i) => {
        const raw = JSON.stringify(item, null, 2);
        if (item.type === "reasoning") {
          return (
            <Block key={i} kind="reasoning">
              <pre className="text-xs text-gray-200 whitespace-pre-wrap font-mono leading-relaxed">{raw}</pre>
            </Block>
          );
        }
        if (item.type === "function_call") {
          return (
            <Block key={i} kind="tool_call" label={`TOOL CALL${item.name ? `: ${item.name}` : ""}`}>
              <pre className="text-xs text-gray-200 whitespace-pre-wrap font-mono leading-relaxed">{raw}</pre>
            </Block>
          );
        }
        if (item.type === "function_call_output") {
          return (
            <Block key={i} kind="tool_resp">
              <pre className="text-xs text-gray-200 whitespace-pre-wrap font-mono leading-relaxed">{raw}</pre>
            </Block>
          );
        }
        return (
          <div key={i} className="border-l-2 border-gray-700 pl-3 py-2 mb-2">
            <pre className="text-xs text-gray-400 whitespace-pre-wrap font-mono leading-relaxed">{raw}</pre>
          </div>
        );
      })}
    </div>
  );
}

// ── Right panel: split assistant content into typed segments ──────────────

type SegKind = "reasoning" | "tool_call";

/**
 * For gpt-oss: split at <tool_call>...</tool_call>.
 * Segments outside tags are reasoning; tags themselves are tool_call.
 */
function splitGptContent(content: string): { kind: SegKind; text: string }[] {
  const segs: { kind: SegKind; text: string }[] = [];
  const re = /(<tool_call>[\s\S]*?<\/tool_call>)/g;
  let last = 0;
  let m: RegExpExecArray | null;
  while ((m = re.exec(content)) !== null) {
    const pre = content.slice(last, m.index).trim();
    if (pre) segs.push({ kind: "reasoning", text: pre });
    segs.push({ kind: "tool_call", text: m[1] });
    last = m.index + m[1].length;
  }
  const tail = content.slice(last).trim();
  if (tail) segs.push({ kind: "reasoning", text: tail });
  return segs;
}

/**
 * For qwen: reasoning is inside <think>...</think>, tool calls in <tool_call>...</tool_call>.
 * Both may appear in one assistant message; extract all in order.
 */
function splitQwenContent(content: string): { kind: SegKind; text: string }[] {
  const segs: { kind: SegKind; text: string }[] = [];
  const re = /(<think>[\s\S]*?<\/think>|<tool_call>[\s\S]*?<\/tool_call>)/g;
  let last = 0;
  let m: RegExpExecArray | null;
  while ((m = re.exec(content)) !== null) {
    const gap = content.slice(last, m.index).trim();
    if (gap) segs.push({ kind: "reasoning", text: gap });
    if (m[1].startsWith("<think>")) {
      segs.push({ kind: "reasoning", text: m[1] });
    } else {
      segs.push({ kind: "tool_call", text: m[1] });
    }
    last = m.index + m[1].length;
  }
  const tail = content.slice(last).trim();
  if (tail) segs.push({ kind: "reasoning", text: tail });
  return segs;
}

function AssistantBlock({ content, template }: { content: string; template: Template }) {
  const segs = template === "qwen" ? splitQwenContent(content) : splitGptContent(content);
  const hasToolCall = segs.some(s => s.kind === "tool_call");
  const outerColor = COLORS.reasoning;

  return (
    <div className={`border-l-2 ${outerColor.border} ${outerColor.bg} pl-3 py-2 rounded-r mb-2`}>
      <div className={`text-[9px] font-bold uppercase tracking-widest mb-1.5 ${outerColor.label}`}>
        {hasToolCall ? "REASONING + TOOL CALL" : "REASONING"}
      </div>
      {segs.map((seg, j) => {
        if (seg.kind === "tool_call") {
          const tc = COLORS.tool_call;
          return (
            <div key={j} className={`border-l-2 ${tc.border} ${tc.bg} pl-2 py-1.5 rounded-r mb-1.5`}>
              <div className={`text-[9px] font-bold uppercase tracking-widest mb-1 ${tc.label}`}>TOOL CALL</div>
              <pre className="text-xs text-gray-200 whitespace-pre-wrap font-mono leading-relaxed">{seg.text}</pre>
            </div>
          );
        }
        return (
          <pre key={j} className="text-xs text-gray-200 whitespace-pre-wrap font-mono leading-relaxed mb-1.5">
            {seg.text}
          </pre>
        );
      })}
    </div>
  );
}

function MessagesPanel({ messages, template }: { messages: Message[]; template: Template }) {
  return (
    <div className="space-y-0">
      {messages.map((msg, i) => {
        if (msg.role === "user" && i === 0) {
          return (
            <Block key={i} kind="user">
              <pre className="text-xs text-gray-200 whitespace-pre-wrap font-mono leading-relaxed">{msg.content}</pre>
            </Block>
          );
        }
        if (msg.role === "user") {
          return (
            <Block key={i} kind="tool_resp">
              <pre className="text-xs text-gray-200 whitespace-pre-wrap font-mono leading-relaxed">{msg.content}</pre>
            </Block>
          );
        }
        return <AssistantBlock key={i} content={msg.content} template={template} />;
      })}
    </div>
  );
}

// ── Legend ────────────────────────────────────────────────────────────────

function Legend() {
  return (
    <div className="flex items-center gap-4 px-4 py-1.5 bg-gray-900/80 border-b border-gray-800 shrink-0 flex-wrap">
      <span className="text-[9px] font-bold uppercase tracking-widest text-gray-500 mr-1">Legend</span>
      {(Object.entries(COLORS) as [keyof typeof COLORS, typeof COLORS[keyof typeof COLORS]][]).map(([k, c]) => (
        <span key={k} className={`flex items-center gap-1.5 text-[10px] ${c.label}`}>
          <span className={`inline-block w-2.5 h-2.5 rounded-sm border ${c.border} ${c.bg}`} />
          {c.tag}
        </span>
      ))}
    </div>
  );
}

// ── Template dropdown ─────────────────────────────────────────────────────

function TemplateDropdown({ value, onChange }: { value: Template; onChange: (t: Template) => void }) {
  return (
    <select
      value={value}
      onChange={e => onChange(e.target.value as Template)}
      className="text-xs bg-gray-800 border border-gray-600 text-gray-200 rounded px-2 py-1 ml-2 cursor-pointer"
    >
      <option value="gpt-oss">gpt-oss</option>
      <option value="qwen">qwen</option>
    </select>
  );
}

// ── Main component ────────────────────────────────────────────────────────

export default function SftDiffApp() {
  const [data, setData]               = useState<Row[]>([]);
  const [loading, setLoading]         = useState(true);
  const [error, setError]             = useState<string | null>(null);
  const [selectedIdx, setSelectedIdx] = useState(0);
  const [search, setSearch]           = useState("");
  const [template, setTemplate]       = useState<Template>("gpt-oss");

  useEffect(() => {
    setLoading(true);
    fetch("/api/sft-diff/")
      .then(r => { if (!r.ok) throw new Error(r.statusText); return r.json(); })
      .then((d: { rows: Row[] }) => { setData(d.rows); setLoading(false); })
      .catch(e => { setError(e.message); setLoading(false); });
  }, []);

  const filtered = useMemo(() => {
    if (!search.trim()) return data;
    const q = search.toLowerCase();
    return data.filter(r =>
      r.query_id.toLowerCase().includes(q) ||
      (r.messages_gpt?.[0]?.content ?? r.messages_qwen?.[0]?.content ?? "").toLowerCase().includes(q)
    );
  }, [data, search]);

  const current = filtered[selectedIdx] ?? null;

  const activeMessages: Message[] | null = current
    ? (template === "qwen" ? current.messages_qwen : current.messages_gpt)
    : null;

  // User content for the left panel — prefer gpt-oss, fall back to qwen (same content)
  const userContent = current?.messages_gpt?.[0]?.content ?? current?.messages_qwen?.[0]?.content ?? "";

  if (loading) return <div className="h-full flex items-center justify-center text-gray-400">Loading…</div>;
  if (error)   return <div className="h-full flex items-center justify-center text-red-400">Error: {error}</div>;

  return (
    <div className="h-full flex overflow-hidden bg-gray-950 text-gray-100">

      {/* ── Sidebar ──────────────────────────────────────────────── */}
      <div className="w-60 shrink-0 flex flex-col border-r border-gray-800 bg-gray-900">
        <div className="px-2 py-1.5 border-b border-gray-800">
          <input
            type="text"
            placeholder="Search query ID or question…"
            value={search}
            onChange={e => { setSearch(e.target.value); setSelectedIdx(0); }}
            className="w-full bg-gray-800 border border-gray-700 text-gray-200 text-xs rounded px-2 py-1.5 placeholder-gray-600"
          />
          <div className="text-[10px] text-gray-600 mt-1">{filtered.length} / {data.length}</div>
        </div>

        <div className="flex-1 overflow-y-auto">
          {filtered.map((row, i) => {
            const content = row.messages_gpt?.[0]?.content ?? row.messages_qwen?.[0]?.content ?? "";
            const qMatch = content.match(/Question:\s*([\s\S]{0,120})/);
            const preview = qMatch ? qMatch[1].trim().replace(/\n/g, " ") : content.slice(0, 80);
            return (
              <button key={row.query_id + i} onClick={() => setSelectedIdx(i)}
                className={`w-full text-left px-3 py-2 border-b border-gray-800/50 text-xs transition-colors ${
                  selectedIdx === i
                    ? "bg-blue-900/40 text-blue-200 border-l-2 border-l-blue-500"
                    : "text-gray-400 hover:bg-gray-800"
                }`}
              >
                <div className="font-medium text-gray-200 mb-0.5">#{row.query_id}</div>
                <div className="text-[10px] text-gray-500 leading-snug line-clamp-2">{preview}</div>
              </button>
            );
          })}
        </div>
      </div>

      {/* ── Main area ────────────────────────────────────────────── */}
      {current ? (
        <div className="flex-1 flex flex-col min-w-0 overflow-hidden">
          <Legend />

          <div className="flex-1 flex overflow-hidden min-w-0">

            {/* Left: original excerpt */}
            <div className="flex-1 flex flex-col min-w-0 border-r border-gray-800 overflow-hidden">
              <div className="px-3 py-1.5 bg-gray-900/60 border-b border-gray-800 shrink-0">
                <span className="text-[11px] font-semibold text-amber-400 uppercase tracking-widest">Original</span>
                <span className="text-[10px] text-gray-500 ml-2">excerpt field — raw JSON items</span>
              </div>
              <div className="flex-1 overflow-y-auto p-3">
                <ExcerptPanel excerpt={current.excerpt} userContent={userContent} />
              </div>
            </div>

            {/* Right: converted messages */}
            <div className="flex-1 flex flex-col min-w-0 overflow-hidden">
              <div className="px-3 py-1.5 bg-gray-900/60 border-b border-gray-800 shrink-0 flex items-center">
                <span className="text-[11px] font-semibold text-sky-400 uppercase tracking-widest">Converted</span>
                <TemplateDropdown value={template} onChange={setTemplate} />
                {activeMessages && (
                  <span className="text-[10px] text-gray-500 ml-2">{activeMessages.length} messages</span>
                )}
              </div>
              <div className="flex-1 overflow-y-auto p-3">
                {activeMessages ? (
                  <MessagesPanel messages={activeMessages} template={template} />
                ) : (
                  <div className="text-gray-500 text-xs italic mt-4">
                    Not available for {template} template.
                  </div>
                )}
              </div>
            </div>

          </div>
        </div>
      ) : (
        <div className="flex-1 flex items-center justify-center text-gray-500">No record selected.</div>
      )}
    </div>
  );
}