webapp/components/tree-visualizer/CreateTreeDialog.tsx

"use client";

import { useEffect, useMemo, useState } from "react";
import { useMutation } from "@tanstack/react-query";
import { AnimatePresence, motion } from "framer-motion";
import {
  AlertTriangle,
  CheckCircle2,
  Columns3,
  FileJson,
  FileUp,
  Loader2,
  X
} from "lucide-react";
import { Button } from "@/components/ui/Button";
import { Input } from "@/components/ui/Input";
import { Select } from "@/components/ui/Select";
import { Textarea } from "@/components/ui/Textarea";
import { Badge } from "@/components/ui/Badge";
import { TreeValidationPanel } from "@/components/tree-visualizer/TreeValidationPanel";
import {
  createTreeFromRaw,
  createTreesFromJsonl,
  parseTrajectory,
  previewBulkTreeImport
} from "@/lib/api";
import type { BulkImportPreview, ParseResult, TrajectoryTree } from "@/lib/tree/types";
import { cn } from "@/lib/utils";

interface CreateTreeDialogProps {
  open: boolean;
  onClose: () => void;
  onCreated: (tree: TrajectoryTree) => void;
}

type ImportMode = "single" | "bulk";

const exampleText = `<|trajectory|>{
  "v": 1,
  "id": "n1",
  "parent": null,
  "kind": "understand",
  "goal": "Restate the problem as a counting probability task",
  "skill": ["probability", "combinatorics"],
  "inputs": { "problem_refs": ["p0"], "node_refs": [] },
  "state_delta": { "claims": [], "equations": [], "constraints": [], "definitions": [] },
  "retrieval": { "needed": false, "source": "none", "queries": [], "evidence_refs": [] },
  "tool": null,
  "verification": { "kind": "none", "status": "pending", "notes": "" },
  "branch": { "mode": "main", "label": "direct" }
}<|end_of_trajectory|>
We need to reinterpret the question as a finite counting problem before choosing a probability model.`;

const exampleJsonl = JSON.stringify({ name: "Counting example", trajectory: exampleText }) + "\n";

function columnOptions(preview: BulkImportPreview | null) {
  return preview?.columns ?? [];
}

export function CreateTreeDialog({ open, onClose, onCreated }: CreateTreeDialogProps) {
  const [mode, setMode] = useState<ImportMode>("single");
  const [title, setTitle] = useState("");
  const [rawText, setRawText] = useState("");
  const [parseResult, setParseResult] = useState<ParseResult | null>(null);
  const [jsonlText, setJsonlText] = useState("");
  const [bulkPreview, setBulkPreview] = useState<BulkImportPreview | null>(null);
  const [trajectoryColumn, setTrajectoryColumn] = useState("");
  const [titleColumn, setTitleColumn] = useState("");

  useEffect(() => {
    if (!open) return;
    setParseResult(null);
    setBulkPreview(null);
  }, [open]);

  const parseMutation = useMutation({
    mutationFn: () => parseTrajectory(rawText),
    onSuccess: ({ parse }) => {
      setParseResult(parse);
      if (!title.trim()) setTitle(parse.title);
    }
  });

  const createMutation = useMutation({
    mutationFn: () => createTreeFromRaw({ title: title.trim() || parseResult?.title, rawText }),
    onSuccess: ({ tree }) => {
      onCreated(tree);
      setRawText("");
      setTitle("");
      setParseResult(null);
      onClose();
    }
  });

  const bulkPreviewMutation = useMutation({
    mutationFn: () =>
      previewBulkTreeImport({
        rawJsonl: jsonlText,
        trajectoryColumn: trajectoryColumn || undefined,
        titleColumn: titleColumn || undefined
      }),
    onSuccess: ({ preview }) => {
      setBulkPreview(preview);
      setTrajectoryColumn(preview.trajectoryColumn);
      setTitleColumn(preview.titleColumn);
    }
  });

  const bulkCreateMutation = useMutation({
    mutationFn: () =>
      createTreesFromJsonl({
        rawJsonl: jsonlText,
        trajectoryColumn,
        titleColumn: titleColumn || undefined
      }),
    onSuccess: ({ tree }) => {
      if (tree) onCreated(tree);
      setJsonlText("");
      setBulkPreview(null);
      setTrajectoryColumn("");
      setTitleColumn("");
      onClose();
    }
  });

  const activeError = parseMutation.error || createMutation.error || bulkPreviewMutation.error || bulkCreateMutation.error;
  const bulkRows = useMemo(() => bulkPreview?.rows.slice(0, 80) ?? [], [bulkPreview?.rows]);

  const loadJsonlFile = async (file: File | undefined) => {
    if (!file) return;
    setJsonlText(await file.text());
    setBulkPreview(null);
    setTrajectoryColumn("");
    setTitleColumn("");
  };

  return (
    <AnimatePresence>
      {open && (
        <motion.div
          className="fixed inset-0 z-50 flex items-center justify-center bg-ink-900/20 px-4 py-6 backdrop-blur-sm"
          initial={{ opacity: 0 }}
          animate={{ opacity: 1 }}
          exit={{ opacity: 0 }}
        >
          <motion.div
            initial={{ opacity: 0, scale: 0.96, y: 16 }}
            animate={{ opacity: 1, scale: 1, y: 0 }}
            exit={{ opacity: 0, scale: 0.98, y: 8 }}
            transition={{ duration: 0.2 }}
            className="flex max-h-[88vh] w-full max-w-5xl flex-col overflow-hidden rounded-xl border border-paper-300 bg-white/92 shadow-soft backdrop-blur-xl"
          >
            <div className="flex items-center justify-between border-b border-paper-300 px-5 py-4">
              <div>
                <div className="flex items-center gap-2">
                  <FileJson className="h-4 w-4 text-claude-500" />
                  <h2 className="text-base font-semibold text-ink-900">Create Tree</h2>
                </div>
                <p className="mt-1 text-sm text-ink-500">
                  Paste one trace or bulk import JSONL rows with trajectory columns.
                </p>
              </div>
              <Button size="icon" variant="ghost" onClick={onClose} title="Close dialog">
                <X className="h-4 w-4" />
              </Button>
            </div>

            <div className="border-b border-paper-300 px-5 py-3">
              <div className="inline-flex rounded-lg border border-paper-300 bg-paper-50 p-1">
                {(["single", "bulk"] as const).map((item) => (
                  <button
                    key={item}
                    type="button"
                    className={cn(
                      "rounded-md px-3 py-1.5 text-sm font-semibold transition",
                      mode === item ? "bg-white text-claude-700 shadow-sm" : "text-ink-500 hover:text-ink-900"
                    )}
                    onClick={() => setMode(item)}
                  >
                    {item === "single" ? "Single trace" : "Bulk JSONL"}
                  </button>
                ))}
              </div>
            </div>

            {mode === "single" ? (
              <div className="grid min-h-0 flex-1 grid-cols-1 gap-4 overflow-auto p-5 scrollbar-thin lg:grid-cols-[minmax(0,1fr)_380px]">
                <div className="space-y-4">
                  <div>
                    <label className="mb-2 block text-xs font-semibold uppercase text-ink-500">Tree title</label>
                    <Input value={title} onChange={(event) => setTitle(event.target.value)} placeholder="Optional title" />
                  </div>
                  <div>
                    <div className="mb-2 flex items-center justify-between gap-2">
                      <label className="text-xs font-semibold uppercase text-ink-500">Raw trajectory text</label>
                      <Button size="sm" variant="ghost" onClick={() => setRawText(exampleText)}>
                        Use example
                      </Button>
                    </div>
                    <Textarea
                      value={rawText}
                      onChange={(event) => {
                        setRawText(event.target.value);
                        setParseResult(null);
                      }}
                      placeholder="<|trajectory|>{ JSON }<|end_of_trajectory|> reasoning text..."
                      className="min-h-[430px] font-mono text-xs leading-5"
                    />
                  </div>
                </div>

                <div className="space-y-4">
                  <div className="rounded-lg border border-paper-300 bg-paper-50/80 p-4">
                    <div className="mb-3 flex items-center justify-between">
                      <h3 className="text-sm font-semibold text-ink-900">Parse preview</h3>
                      {parseResult && (
                        <Badge status={parseResult.validation.ok ? "valid" : "error"}>
                          {parseResult.nodes.length} nodes
                        </Badge>
                      )}
                    </div>
                    {!parseResult ? (
                      <p className="text-sm leading-6 text-ink-500">
                        Run preview to inspect extracted trajectory blocks before saving.
                      </p>
                    ) : (
                      <div className="space-y-3">
                        <div className="flex items-center gap-2 text-sm">
                          {parseResult.validation.ok ? (
                            <CheckCircle2 className="h-4 w-4 text-green-600" />
                          ) : (
                            <AlertTriangle className="h-4 w-4 text-red-600" />
                          )}
                          <span className="font-medium text-ink-900">{parseResult.title}</span>
                        </div>
                        <div className="max-h-56 space-y-2 overflow-auto pr-1 scrollbar-thin">
                          {parseResult.blocks.map((block) => (
                            <div key={block.index} className="rounded-md border border-paper-200 bg-white px-3 py-2">
                              <div className="flex items-center justify-between gap-2">
                                <span className="font-mono text-[11px] text-ink-400">block {block.index + 1}</span>
                                {block.error ? <Badge status="error">error</Badge> : <Badge status="valid">{block.node?.id}</Badge>}
                              </div>
                              {block.node && <p className="mt-1 truncate text-xs text-ink-700">{block.node.goal}</p>}
                              {block.error && <p className="mt-1 text-xs text-red-600">{block.error}</p>}
                            </div>
                          ))}
                        </div>
                      </div>
                    )}
                  </div>
                  <TreeValidationPanel validation={parseResult?.validation ?? null} />
                </div>
              </div>
            ) : (
              <div className="grid min-h-0 flex-1 grid-cols-1 gap-4 overflow-auto p-5 scrollbar-thin lg:grid-cols-[minmax(0,1fr)_400px]">
                <div className="space-y-4">
                  <div className="rounded-lg border border-paper-300 bg-paper-50/70 p-3">
                    <div className="flex flex-wrap items-center justify-between gap-2">
                      <div>
                        <div className="flex items-center gap-2 text-sm font-semibold text-ink-900">
                          <FileUp className="h-4 w-4 text-claude-500" />
                          JSONL source
                        </div>
                        <p className="mt-1 text-xs text-ink-500">One JSON object per line. Choose the trace and name columns after preview.</p>
                      </div>
                      <div className="flex gap-2">
                        <Button size="sm" variant="ghost" onClick={() => setJsonlText(exampleJsonl)}>
                          Use example
                        </Button>
                        <label className="inline-flex h-8 cursor-pointer items-center gap-2 rounded-lg border border-paper-300 bg-white/84 px-3 text-xs font-medium text-ink-700 shadow-sm transition hover:border-claude-300 hover:bg-claude-50">
                          Upload JSONL
                          <input
                            type="file"
                            accept=".jsonl,.ndjson,application/jsonl,text/plain"
                            className="hidden"
                            onChange={(event) => void loadJsonlFile(event.target.files?.[0])}
                          />
                        </label>
                      </div>
                    </div>
                  </div>
                  <Textarea
                    value={jsonlText}
                    onChange={(event) => {
                      setJsonlText(event.target.value);
                      setBulkPreview(null);
                      setTrajectoryColumn("");
                      setTitleColumn("");
                    }}
                    placeholder='{"name":"Example","trajectory":"<|trajectory|>{...}<|end_of_trajectory|> reasoning..."}'
                    className="min-h-[430px] font-mono text-xs leading-5"
                  />
                </div>

                <div className="space-y-4">
                  <div className="rounded-lg border border-paper-300 bg-paper-50/80 p-4">
                    <div className="mb-3 flex items-center justify-between">
                      <h3 className="text-sm font-semibold text-ink-900">Bulk preview</h3>
                      {bulkPreview && (
                        <Badge status={bulkPreview.validRows === bulkPreview.rowCount ? "valid" : "warning"}>
                          {bulkPreview.validRows}/{bulkPreview.rowCount} valid
                        </Badge>
                      )}
                    </div>

                    {!bulkPreview ? (
                      <p className="text-sm leading-6 text-ink-500">
                        Preview JSONL to detect columns, then choose which column contains trajectories and which column names the saved trees.
                      </p>
                    ) : (
                      <div className="space-y-4">
                        <div className="grid grid-cols-1 gap-3 sm:grid-cols-2">
                          <div>
                            <label className="mb-2 block text-xs font-semibold uppercase text-ink-500">Trajectory column</label>
                            <Select
                              value={trajectoryColumn}
                              onChange={(event) => {
                                setTrajectoryColumn(event.target.value);
                                setBulkPreview(null);
                              }}
                            >
                              {columnOptions(bulkPreview).map((column) => (
                                <option key={column.name} value={column.name}>
                                  {column.name}
                                </option>
                              ))}
                            </Select>
                          </div>
                          <div>
                            <label className="mb-2 block text-xs font-semibold uppercase text-ink-500">Name column</label>
                            <Select
                              value={titleColumn}
                              onChange={(event) => {
                                setTitleColumn(event.target.value);
                                setBulkPreview(null);
                              }}
                            >
                              <option value="">Use parsed title</option>
                              {columnOptions(bulkPreview).map((column) => (
                                <option key={column.name} value={column.name}>
                                  {column.name}
                                </option>
                              ))}
                            </Select>
                          </div>
                        </div>

                        <div className="rounded-lg border border-paper-200 bg-white/70 p-3">
                          <div className="mb-2 flex items-center gap-2 text-xs font-semibold uppercase text-ink-400">
                            <Columns3 className="h-4 w-4 text-claude-500" />
                            Columns
                          </div>
                          <div className="max-h-28 space-y-1 overflow-auto scrollbar-thin">
                            {bulkPreview.columns.map((column) => (
                              <div key={column.name} className="grid grid-cols-[120px_1fr] gap-2 text-xs">
                                <span className="truncate font-semibold text-ink-700">{column.name}</span>
                                <span className="truncate text-ink-500">{column.sample || "empty"}</span>
                              </div>
                            ))}
                          </div>
                        </div>

                        <div className="max-h-56 space-y-2 overflow-auto pr-1 scrollbar-thin">
                          {bulkRows.map((row) => (
                            <div key={row.line} className="rounded-md border border-paper-200 bg-white px-3 py-2">
                              <div className="flex items-center justify-between gap-2">
                                <span className="font-mono text-[11px] text-ink-400">line {row.line}</span>
                                <Badge status={row.ok ? "valid" : "error"}>{row.ok ? `${row.nodeCount} nodes` : "error"}</Badge>
                              </div>
                              <p className="mt-1 truncate text-xs font-medium text-ink-700">{row.title}</p>
                              {row.error && <p className="mt-1 text-xs text-red-600">{row.error}</p>}
                            </div>
                          ))}
                        </div>
                      </div>
                    )}
                  </div>
                </div>
              </div>
            )}

            {activeError && (
              <div className="border-t border-red-100 bg-red-50 px-5 py-3 text-sm text-red-700">
                {activeError.message}
              </div>
            )}

            <div className="flex items-center justify-between border-t border-paper-300 px-5 py-4">
              <p className="text-xs text-ink-500">
                Raw text is preserved separately from the normalized tree file.
              </p>
              {mode === "single" ? (
                <div className="flex gap-2">
                  <Button onClick={() => parseMutation.mutate()} disabled={!rawText.trim() || parseMutation.isPending}>
                    {parseMutation.isPending && <Loader2 className="h-4 w-4 animate-spin" />}
                    Preview Parse
                  </Button>
                  <Button
                    variant="primary"
                    onClick={() => createMutation.mutate()}
                    disabled={!rawText.trim() || createMutation.isPending}
                  >
                    {createMutation.isPending && <Loader2 className="h-4 w-4 animate-spin" />}
                    Create Tree
                  </Button>
                </div>
              ) : (
                <div className="flex gap-2">
                  <Button
                    onClick={() => bulkPreviewMutation.mutate()}
                    disabled={!jsonlText.trim() || bulkPreviewMutation.isPending}
                  >
                    {bulkPreviewMutation.isPending && <Loader2 className="h-4 w-4 animate-spin" />}
                    Preview JSONL
                  </Button>
                  <Button
                    variant="primary"
                    onClick={() => bulkCreateMutation.mutate()}
                    disabled={
                      !jsonlText.trim() ||
                      !trajectoryColumn ||
                      !bulkPreview ||
                      bulkPreview.validRows !== bulkPreview.rowCount ||
                      bulkCreateMutation.isPending
                    }
                  >
                    {bulkCreateMutation.isPending && <Loader2 className="h-4 w-4 animate-spin" />}
                    Import {bulkPreview?.validRows ?? 0} Trees
                  </Button>
                </div>
              )}
            </div>
          </motion.div>
        </motion.div>
      )}
    </AnimatePresence>
  );
}