web_ui/frontend/src/components/graph/GraphToolbar.tsx

import { useState } from "react";
import {
  CheckCircle2,
  Layout,
  Save,
  FolderOpen,
  FilePlus,
  Download,
  Upload,
  PlayCircle,
  Flag,
  Trash2,
  BookTemplate,
} from "lucide-react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Separator } from "@/components/ui/separator";
import {
  Dialog,
  DialogContent,
  DialogHeader,
  DialogTitle,
  DialogDescription,
} from "@/components/ui/dialog";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Badge } from "@/components/ui/badge";
import { useGraphStore, type AgentNodeData, type GraphTemplate } from "@/stores/graphStore";
import type { GraphListItem } from "@/types/graph";

// ── Predefined graph templates ──────────────────────────────────────────
const GRAPH_TEMPLATES: GraphTemplate[] = [
  {
    name: "Linear Pipeline",
    description: "Sequential: Researcher -> Writer -> Reviewer",
    agents: [
      { agent_id: "researcher", display_name: "Researcher", persona: "Expert researcher", description: "Gathers information", tools: ["web_search"] },
      { agent_id: "writer", display_name: "Writer", persona: "Technical writer", description: "Writes content", tools: [] },
      { agent_id: "reviewer", display_name: "Reviewer", persona: "Quality reviewer", description: "Reviews output", tools: [] },
    ],
    edges: [
      { source: "researcher", target: "writer", weight: 1.0, condition: null },
      { source: "writer", target: "reviewer", weight: 1.0, condition: null },
    ],
    start_node: "researcher",
    end_node: "reviewer",
  },
  {
    name: "Fan-Out / Fan-In",
    description: "Planner fans out to parallel workers, Aggregator collects results",
    agents: [
      { agent_id: "planner", display_name: "Planner", persona: "Task planner", description: "Breaks task into subtasks", tools: [] },
      { agent_id: "worker_a", display_name: "Worker A", persona: "Specialist A", description: "Handles subtask A", tools: [] },
      { agent_id: "worker_b", display_name: "Worker B", persona: "Specialist B", description: "Handles subtask B", tools: [] },
      { agent_id: "aggregator", display_name: "Aggregator", persona: "Synthesizer", description: "Combines results", tools: [] },
    ],
    edges: [
      { source: "planner", target: "worker_a", weight: 1.0, condition: null },
      { source: "planner", target: "worker_b", weight: 1.0, condition: null },
      { source: "worker_a", target: "aggregator", weight: 1.0, condition: null },
      { source: "worker_b", target: "aggregator", weight: 1.0, condition: null },
    ],
    start_node: "planner",
    end_node: "aggregator",
  },
  {
    name: "Conditional Branching",
    description: "Writer outputs, conditional edges route to Editor or Fixer based on response",
    agents: [
      { agent_id: "writer", display_name: "Writer", persona: "Content writer", description: "Generates draft", tools: [] },
      { agent_id: "editor", display_name: "Editor", persona: "Editor", description: "Polishes good drafts", tools: [] },
      { agent_id: "fixer", display_name: "Fixer", persona: "Error corrector", description: "Fixes problematic drafts", tools: [] },
      { agent_id: "publisher", display_name: "Publisher", persona: "Publisher", description: "Publishes final version", tools: [] },
    ],
    edges: [
      { source: "writer", target: "editor", weight: 0.9, condition: "contains:APPROVED" },
      { source: "writer", target: "fixer", weight: 0.7, condition: "source_failed" },
      { source: "editor", target: "publisher", weight: 1.0, condition: null },
      { source: "fixer", target: "publisher", weight: 1.0, condition: null },
    ],
    start_node: "writer",
    end_node: "publisher",
  },
  {
    name: "Review Loop",
    description: "Coder writes code, Reviewer checks quality; conditional edges route back or forward",
    agents: [
      { agent_id: "planner", display_name: "Planner", persona: "Architect", description: "Plans the approach", tools: [] },
      { agent_id: "coder", display_name: "Coder", persona: "Software developer", description: "Writes code", tools: ["code_interpreter"] },
      { agent_id: "reviewer", display_name: "Reviewer", persona: "Code reviewer", description: "Reviews code quality", tools: [] },
      { agent_id: "finalizer", display_name: "Finalizer", persona: "Integrator", description: "Delivers final result", tools: [] },
    ],
    edges: [
      { source: "planner", target: "coder", weight: 1.0, condition: null },
      { source: "coder", target: "reviewer", weight: 1.0, condition: null },
      { source: "reviewer", target: "finalizer", weight: 0.9, condition: "contains:LGTM" },
      { source: "reviewer", target: "coder", weight: 0.5, condition: "contains:REVISE" },
    ],
    start_node: "planner",
    end_node: "finalizer",
  },
  {
    name: "Diamond with Weights",
    description: "Dispatcher -> two parallel paths with different weights -> Merger",
    agents: [
      { agent_id: "start", display_name: "Dispatcher", persona: "Task dispatcher", description: "Dispatches task", tools: [] },
      { agent_id: "fast_path", display_name: "Fast Path", persona: "Quick processor", description: "Fast but less thorough", tools: [] },
      { agent_id: "deep_path", display_name: "Deep Path", persona: "Deep analyzer", description: "Slow but thorough", tools: ["web_search"] },
      { agent_id: "merger", display_name: "Merger", persona: "Result merger", description: "Merges results", tools: [] },
    ],
    edges: [
      { source: "start", target: "fast_path", weight: 0.6, condition: null },
      { source: "start", target: "deep_path", weight: 1.0, condition: null },
      { source: "fast_path", target: "merger", weight: 0.6, condition: null },
      { source: "deep_path", target: "merger", weight: 1.0, condition: null },
    ],
    start_node: "start",
    end_node: "merger",
  },
];

interface GraphToolbarProps {
  onValidate: () => void;
  onRun: () => void;
  isRunning: boolean;
}

export function GraphToolbar({ onValidate, onRun, isRunning }: GraphToolbarProps) {
  const {
    graphName,
    setGraphName,
    autoLayout,
    saveGraph,
    loadGraph,
    newGraph,
    loadTemplate,
    fetchSavedGraphs,
    savedGraphs,
    nodes,
    validationErrors,
    validationWarnings,
    setStartNode,
    setEndNode,
    startNode,
    endNode,
    toGraphRequest,
  } = useGraphStore();

  const [saving, setSaving] = useState(false);
  const [loadOpen, setLoadOpen] = useState(false);
  const [startEndOpen, setStartEndOpen] = useState(false);
  const [templatesOpen, setTemplatesOpen] = useState(false);

  const handleSave = async () => {
    setSaving(true);
    try {
      await saveGraph();
    } catch {
      // handle error
    } finally {
      setSaving(false);
    }
  };

  const handleLoadOpen = async () => {
    await fetchSavedGraphs();
    setLoadOpen(true);
  };

  const handleLoad = async (graph: GraphListItem) => {
    await loadGraph(graph.graph_id);
    setLoadOpen(false);
  };

  const handleExport = () => {
    const data = toGraphRequest();
    const blob = new Blob([JSON.stringify(data, null, 2)], { type: "application/json" });
    const url = URL.createObjectURL(blob);
    const a = document.createElement("a");
    a.href = url;
    a.download = `${graphName.replace(/\s+/g, "-").toLowerCase()}.json`;
    a.click();
    URL.revokeObjectURL(url);
  };

  const handleImport = () => {
    const input = document.createElement("input");
    input.type = "file";
    input.accept = ".json";
    input.onchange = async (e) => {
      const file = (e.target as HTMLInputElement).files?.[0];
      if (!file) return;
      const text = await file.text();
      try {
        const data = JSON.parse(text);
        if (data.name) setGraphName(data.name);
      } catch {
        // invalid JSON
      }
    };
    input.click();
  };

  const handleLoadTemplate = (template: GraphTemplate) => {
    loadTemplate(template);
    setTemplatesOpen(false);
  };

  return (
    <>
      <div className="flex items-center gap-2 border-b bg-card px-3 py-2">
        <Input
          value={graphName}
          onChange={(e) => setGraphName(e.target.value)}
          className="h-8 w-48 text-sm font-medium"
        />

        <Separator orientation="vertical" className="h-6" />

        <Button variant="outline" size="sm" onClick={onValidate} className="gap-1">
          <CheckCircle2 className="h-3.5 w-3.5" />
          Validate
        </Button>

        {validationErrors.length > 0 && (
          <Badge variant="destructive" className="text-xs">
            {validationErrors.length} error{validationErrors.length > 1 ? "s" : ""}
          </Badge>
        )}
        {validationWarnings.length > 0 && (
          <Badge variant="secondary" className="text-xs">
            {validationWarnings.length} warning{validationWarnings.length > 1 ? "s" : ""}
          </Badge>
        )}

        <Button variant="outline" size="sm" onClick={autoLayout} className="gap-1">
          <Layout className="h-3.5 w-3.5" />
          Auto Layout
        </Button>

        <Button variant="outline" size="sm" onClick={() => setStartEndOpen(true)} className="gap-1">
          <Flag className="h-3.5 w-3.5" />
          Start/End
        </Button>

        <Separator orientation="vertical" className="h-6" />

        <Button variant="outline" size="sm" onClick={handleSave} disabled={saving} className="gap-1">
          <Save className="h-3.5 w-3.5" />
          {saving ? "Saving..." : "Save"}
        </Button>

        <Button variant="outline" size="sm" onClick={handleLoadOpen} className="gap-1">
          <FolderOpen className="h-3.5 w-3.5" />
          Load
        </Button>

        <Button variant="outline" size="sm" onClick={() => setTemplatesOpen(true)} className="gap-1">
          <BookTemplate className="h-3.5 w-3.5" />
          Templates
        </Button>

        <Button variant="outline" size="sm" onClick={newGraph} className="gap-1">
          <FilePlus className="h-3.5 w-3.5" />
          New
        </Button>

        <Button
          variant="outline"
          size="sm"
          onClick={() => { if (nodes.length === 0 || window.confirm("Clear all nodes and edges?")) newGraph(); }}
          disabled={nodes.length === 0}
          className="gap-1 text-destructive hover:text-destructive"
        >
          <Trash2 className="h-3.5 w-3.5" />
          Clear
        </Button>

        <Button variant="ghost" size="icon" className="h-8 w-8" onClick={handleExport}>
          <Download className="h-3.5 w-3.5" />
        </Button>

        <Button variant="ghost" size="icon" className="h-8 w-8" onClick={handleImport}>
          <Upload className="h-3.5 w-3.5" />
        </Button>

        <div className="flex-1" />

        <Button onClick={onRun} disabled={isRunning || nodes.length === 0} className="gap-1">
          <PlayCircle className="h-4 w-4" />
          {isRunning ? "Running..." : "Execute"}
        </Button>
      </div>

      {/* Load dialog */}
      <Dialog open={loadOpen} onOpenChange={setLoadOpen}>
        <DialogContent>
          <DialogHeader>
            <DialogTitle>Load Workflow</DialogTitle>
            <DialogDescription>Select a saved workflow to load.</DialogDescription>
          </DialogHeader>
          <ScrollArea className="max-h-[400px]">
            {savedGraphs.length === 0 ? (
              <p className="p-4 text-center text-sm text-muted-foreground">No saved workflows.</p>
            ) : (
              <div className="space-y-2 p-1">
                {savedGraphs.map((g) => (
                  <div
                    key={g.graph_id}
                    className="flex items-center justify-between rounded-md border p-3 cursor-pointer hover:bg-accent"
                    onClick={() => handleLoad(g)}
                  >
                    <div>
                      <div className="font-medium text-sm">{g.name}</div>
                      <div className="text-xs text-muted-foreground">
                        {g.agent_count} agents, {g.edge_count} edges
                      </div>
                    </div>
                    <div className="text-xs text-muted-foreground">
                      {g.updated_at ? new Date(g.updated_at).toLocaleDateString() : ""}
                    </div>
                  </div>
                ))}
              </div>
            )}
          </ScrollArea>
        </DialogContent>
      </Dialog>

      {/* Templates dialog */}
      <Dialog open={templatesOpen} onOpenChange={setTemplatesOpen}>
        <DialogContent className="sm:max-w-[520px]">
          <DialogHeader>
            <DialogTitle>Graph Templates</DialogTitle>
            <DialogDescription>Start from a predefined topology with agents and conditional edges.</DialogDescription>
          </DialogHeader>
          <ScrollArea className="max-h-[400px]">
            <div className="space-y-2 p-1">
              {GRAPH_TEMPLATES.map((t) => (
                <div
                  key={t.name}
                  className="rounded-md border p-3 cursor-pointer hover:bg-accent transition-colors"
                  onClick={() => handleLoadTemplate(t)}
                >
                  <div className="font-medium text-sm">{t.name}</div>
                  <div className="text-xs text-muted-foreground mt-0.5">{t.description}</div>
                  <div className="flex gap-2 mt-1.5">
                    <Badge variant="secondary" className="text-[10px]">
                      {t.agents.length} agents
                    </Badge>
                    <Badge variant="secondary" className="text-[10px]">
                      {t.edges.length} edges
                    </Badge>
                    {t.edges.some((e) => e.condition) && (
                      <Badge variant="outline" className="text-[10px] text-indigo-600">
                        conditional
                      </Badge>
                    )}
                    {t.edges.some((e) => e.weight < 1.0) && (
                      <Badge variant="outline" className="text-[10px] text-amber-600">
                        weighted
                      </Badge>
                    )}
                  </div>
                </div>
              ))}
            </div>
          </ScrollArea>
        </DialogContent>
      </Dialog>

      {/* Start/End node dialog */}
      <Dialog open={startEndOpen} onOpenChange={setStartEndOpen}>
        <DialogContent>
          <DialogHeader>
            <DialogTitle>Set Start/End Nodes</DialogTitle>
            <DialogDescription>
              Select which agent starts and ends the workflow execution.
            </DialogDescription>
          </DialogHeader>
          <div className="space-y-4">
            <div className="space-y-2">
              <label className="text-sm font-medium">Start Node</label>
              <div className="space-y-1">
                {nodes.map((n) => (
                  <div
                    key={n.id}
                    className={`flex items-center gap-2 rounded-md border p-2 cursor-pointer transition-colors ${
                      startNode === n.id ? "border-green-500 bg-green-50 dark:bg-green-950" : "hover:bg-accent"
                    }`}
                    onClick={() => setStartNode(startNode === n.id ? null : n.id)}
                  >
                    <span className="text-sm">{(n.data as unknown as AgentNodeData).displayName}</span>
                    {startNode === n.id && <Badge className="ml-auto text-xs">START</Badge>}
                  </div>
                ))}
              </div>
            </div>
            <div className="space-y-2">
              <label className="text-sm font-medium">End Node</label>
              <div className="space-y-1">
                {nodes.map((n) => (
                  <div
                    key={n.id}
                    className={`flex items-center gap-2 rounded-md border p-2 cursor-pointer transition-colors ${
                      endNode === n.id ? "border-orange-500 bg-orange-50 dark:bg-orange-950" : "hover:bg-accent"
                    }`}
                    onClick={() => setEndNode(endNode === n.id ? null : n.id)}
                  >
                    <span className="text-sm">{(n.data as unknown as AgentNodeData).displayName}</span>
                    {endNode === n.id && <Badge className="ml-auto text-xs">END</Badge>}
                  </div>
                ))}
              </div>
            </div>
          </div>
        </DialogContent>
      </Dialog>
    </>
  );
}