"use client";
import { useState, useMemo, useEffect, useRef } from "react";
import {
  DndContext,
  DragOverlay,
  KeyboardSensor,
  PointerSensor,
  useSensor,
  useSensors,
  closestCenter,
  type DragEndEvent,
  type DragMoveEvent,
  type DragStartEvent,
  type DragOverEvent,
  type UniqueIdentifier,
} from "@dnd-kit/core";
import {
  SortableContext,
  useSortable,
  arrayMove,
  verticalListSortingStrategy,
  sortableKeyboardCoordinates,
} from "@dnd-kit/sortable";
import { CSS } from "@dnd-kit/utilities";
import {
  GripVertical,
  Plus,
  Trash2,
  Edit3,
  ChevronDown,
  ChevronRight,
  Users,
  Building2,
  Shield,
  Star,
  CornerDownRight,
  Loader2,
  Check,
  UserMinus,
} from "lucide-react";
import { cn } from "@/lib/utils";
import Avatar from "./ui/Avatar";
import Badge from "./ui/Badge";
import Modal from "./ui/Modal";

export interface OrgNode {
  id: number;
  departmentId: number | null;
  parentId: number | null;
  name: string;
  type: "department" | "district" | "squad" | "unit" | "position";
  color: string | null;
  sortOrder: number | null;
}

interface Employee {
  id: number;
  firstName: string;
  lastName: string;
  rank: string | null;
  role: string;
  orgNodeId: number | null;
  avatarColor: string | null;
  badgeNumber: string | null;
  isActive?: boolean | null;
}

interface OrgChartProps {
  nodes: OrgNode[];
  employees: Employee[];
  onNodesChange: () => void;
  onEmployeesChange: () => void;
  departmentId: number;
}

const NODE_TYPE_ICONS: Record<string, React.ElementType> = {
  department: Building2,
  district: Shield,
  squad: Star,
  unit: Users,
  position: Users,
};

const NODE_TYPES = ["district", "squad", "unit", "position"] as const;
const INDENT = 28;
const MAX_DEPTH = 4;

// ---------------------------------------------------------------- tree helpers
interface FlatNode {
  id: number;
  parentId: number | null;
  depth: number;
  node: OrgNode;
}

function flattenTree(nodes: OrgNode[]): FlatNode[] {
  const byParent = new Map<number | null, OrgNode[]>();
  for (const n of nodes) {
    const key = n.parentId ?? null;
    byParent.set(key, [...(byParent.get(key) ?? []), n]);
  }
  for (const [, list] of byParent) {
    list.sort((a, b) => (a.sortOrder ?? 0) - (b.sortOrder ?? 0) || a.id - b.id);
  }
  const out: FlatNode[] = [];
  const walk = (parentId: number | null, depth: number) => {
    for (const n of byParent.get(parentId) ?? []) {
      out.push({ id: n.id, parentId, depth, node: n });
      walk(n.id, depth + 1);
    }
  };
  walk(null, 0);
  return out;
}

function getDescendantIds(items: FlatNode[], id: number): number[] {
  const startIndex = items.findIndex((i) => i.id === id);
  if (startIndex === -1) return [];
  const baseDepth = items[startIndex].depth;
  const out: number[] = [];
  for (let i = startIndex + 1; i < items.length; i++) {
    if (items[i].depth <= baseDepth) break;
    out.push(items[i].id);
  }
  return out;
}

function removeDescendants(items: FlatNode[], ids: number[]): FlatNode[] {
  const excluded = new Set<number>();
  for (const id of ids) getDescendantIds(items, id).forEach((d) => excluded.add(d));
  return items.filter((i) => !excluded.has(i.id));
}

/** Works out the parent + depth an in-flight drag would land on. */
function getProjection(
  items: FlatNode[],
  activeId: number,
  overId: number,
  dragOffsetX: number
): { depth: number; parentId: number | null } {
  const overIndex = items.findIndex((i) => i.id === overId);
  const activeIndex = items.findIndex((i) => i.id === activeId);
  if (overIndex === -1 || activeIndex === -1) return { depth: 0, parentId: null };

  const activeItem = items[activeIndex];
  const newItems = arrayMove(items, activeIndex, overIndex);
  const previousItem = newItems[overIndex - 1];
  const nextItem = newItems[overIndex + 1];

  const dragDepth = Math.round(dragOffsetX / INDENT);
  const projectedDepth = activeItem.depth + dragDepth;
  const maxDepth = Math.min(previousItem ? previousItem.depth + 1 : 0, MAX_DEPTH);
  const minDepth = nextItem ? nextItem.depth : 0;

  let depth = projectedDepth;
  if (projectedDepth >= maxDepth) depth = maxDepth;
  else if (projectedDepth < minDepth) depth = minDepth;

  const parentId = (() => {
    if (depth === 0 || !previousItem) return null;
    if (depth === previousItem.depth) return previousItem.parentId;
    if (depth > previousItem.depth) return previousItem.id;
    const candidate = newItems
      .slice(0, overIndex)
      .reverse()
      .find((i) => i.depth === depth);
    return candidate?.parentId ?? null;
  })();

  return { depth, parentId };
}

// ---------------------------------------------------------------- row
function OrgRow({
  item,
  depth,
  employees,
  collapsed,
  hasChildren,
  onToggle,
  onEdit,
  onDelete,
  onAddChild,
  onMoveEmployee,
  isGhost,
}: {
  item: FlatNode;
  depth: number;
  employees: Employee[];
  collapsed: boolean;
  hasChildren: boolean;
  onToggle: (id: number) => void;
  onEdit: (n: OrgNode) => void;
  onDelete: (n: OrgNode) => void;
  onAddChild: (id: number) => void;
  onMoveEmployee: (employeeId: number, nodeId: number | null) => void;
  isGhost?: boolean;
}) {
  const [dropActive, setDropActive] = useState(false);
  const { attributes, listeners, setNodeRef, setActivatorNodeRef, transform, transition, isDragging } =
    useSortable({ id: item.id });

  const node = item.node;
  const Icon = NODE_TYPE_ICONS[node.type] ?? Users;
  const nodeEmployees = employees.filter((e) => e.orgNodeId === node.id && e.isActive !== false);
  const color = node.color ?? "#3B82F6";

  return (
    <div
      ref={setNodeRef}
      style={{
        transform: CSS.Translate.toString(transform),
        transition,
        paddingLeft: depth * INDENT,
      }}
      className={cn("relative", isDragging && "opacity-40")}
    >
      {depth > 0 && (
        <CornerDownRight
          size={14}
          className="absolute text-slate-300 top-4"
          style={{ left: depth * INDENT - 18 }}
        />
      )}
      <div
        onDragOver={(e) => {
          e.preventDefault();
          setDropActive(true);
        }}
        onDragLeave={() => setDropActive(false)}
        onDrop={(e) => {
          e.preventDefault();
          setDropActive(false);
          const empId = parseInt(e.dataTransfer.getData("employeeId"));
          if (!Number.isNaN(empId)) onMoveEmployee(empId, node.id);
        }}
        className={cn(
          "mb-1.5 rounded-xl border transition-all",
          dropActive ? "border-dashed border-2 border-blue-500 bg-blue-50" : "border-slate-200 bg-white",
          isGhost && "ring-2 ring-blue-400 shadow-lg"
        )}
        style={{ borderLeft: `4px solid ${color}` }}
      >
        <div className="group flex items-center gap-2 p-2.5">
          <button
            ref={setActivatorNodeRef}
            {...attributes}
            {...listeners}
            className="cursor-grab active:cursor-grabbing text-slate-300 hover:text-slate-600 touch-none p-1 -m-1"
            aria-label={`Drag ${node.name}`}
          >
            <GripVertical size={16} />
          </button>

          <button
            onClick={() => onToggle(node.id)}
            className={cn("text-slate-400 hover:text-slate-700", !hasChildren && "invisible")}
            aria-label="Toggle children"
          >
            {collapsed ? <ChevronRight size={16} /> : <ChevronDown size={16} />}
          </button>

          <div
            className="w-8 h-8 rounded-lg flex items-center justify-center flex-shrink-0"
            style={{ backgroundColor: `${color}22`, color }}
          >
            <Icon size={16} />
          </div>

          <div className="flex-1 min-w-0">
            <div className="flex items-center gap-2 flex-wrap">
              <span className="font-semibold text-slate-800 text-sm">{node.name}</span>
              <Badge color={color}>{node.type}</Badge>
            </div>
            <span className="text-xs text-slate-400">
              {nodeEmployees.length} officer{nodeEmployees.length !== 1 ? "s" : ""}
            </span>
          </div>

          <div className="flex items-center gap-1 lg:opacity-0 group-hover:opacity-100 transition-opacity">
            <button
              onClick={() => onAddChild(node.id)}
              className="p-1.5 rounded-lg bg-blue-50 text-blue-500 hover:bg-blue-100"
              title="Add sub-unit"
            >
              <Plus size={14} />
            </button>
            <button
              onClick={() => onEdit(node)}
              className="p-1.5 rounded-lg bg-slate-50 text-slate-500 hover:bg-slate-100"
              title="Edit unit"
            >
              <Edit3 size={14} />
            </button>
            <button
              onClick={() => onDelete(node)}
              className="p-1.5 rounded-lg bg-red-50 text-red-400 hover:bg-red-100"
              title="Delete unit"
            >
              <Trash2 size={14} />
            </button>
          </div>
        </div>

        {nodeEmployees.length > 0 && (
          <div className="flex flex-wrap gap-1.5 px-3 pb-2.5">
            {nodeEmployees.map((emp) => (
              <div
                key={emp.id}
                draggable
                onDragStart={(e) => {
                  e.dataTransfer.effectAllowed = "move";
                  e.dataTransfer.setData("employeeId", String(emp.id));
                }}
                className="flex items-center gap-2 bg-slate-50 border border-slate-200 rounded-lg pl-1.5 pr-2 py-1 cursor-grab active:cursor-grabbing hover:border-blue-300 hover:bg-blue-50 transition-colors"
                title="Drag to another unit to reassign"
              >
                <Avatar
                  firstName={emp.firstName}
                  lastName={emp.lastName}
                  color={emp.avatarColor ?? "#3B82F6"}
                  size="xs"
                />
                <div className="leading-tight">
                  <p className="text-xs font-medium text-slate-700">
                    {emp.firstName} {emp.lastName}
                  </p>
                  <p className="text-[10px] text-slate-400">{emp.rank ?? emp.role}</p>
                </div>
              </div>
            ))}
          </div>
        )}
      </div>
    </div>
  );
}

// ---------------------------------------------------------------- main
export default function OrgChart({
  nodes,
  employees,
  onNodesChange,
  onEmployeesChange,
  departmentId,
}: OrgChartProps) {
  const [localNodes, setLocalNodes] = useState<OrgNode[]>(nodes);
  const [collapsedIds, setCollapsedIds] = useState<Set<number>>(new Set());
  const [activeId, setActiveId] = useState<number | null>(null);
  const [overId, setOverId] = useState<UniqueIdentifier | null>(null);
  const [offsetX, setOffsetX] = useState(0);
  const [saveState, setSaveState] = useState<"idle" | "saving" | "saved" | "error">("idle");
  const [editingNode, setEditingNode] = useState<OrgNode | null>(null);
  const [showAddModal, setShowAddModal] = useState(false);
  const [addParentId, setAddParentId] = useState<number | null>(null);
  const [newNode, setNewNode] = useState({ name: "", type: "squad" as (typeof NODE_TYPES)[number], color: "#3B82F6" });
  const [unassignedDropActive, setUnassignedDropActive] = useState(false);
  const savedTimer = useRef<ReturnType<typeof setTimeout> | null>(null);

  // Keep local mirror in sync with server data (unless mid-drag)
  useEffect(() => {
    if (activeId === null) setLocalNodes(nodes);
  }, [nodes, activeId]);

  useEffect(() => () => { if (savedTimer.current) clearTimeout(savedTimer.current); }, []);

  const sensors = useSensors(
    useSensor(PointerSensor, { activationConstraint: { distance: 5 } }),
    useSensor(KeyboardSensor, { coordinateGetter: sortableKeyboardCoordinates })
  );

  const flattened = useMemo(() => {
    const full = flattenTree(localNodes);
    const collapsedParents = [...collapsedIds].filter((id) => full.some((f) => f.id === id));
    const withoutCollapsed = removeDescendants(full, collapsedParents);
    return activeId !== null ? removeDescendants(withoutCollapsed, [activeId]) : withoutCollapsed;
  }, [localNodes, collapsedIds, activeId]);

  const sortableIds = useMemo(() => flattened.map((f) => f.id), [flattened]);

  const projection = useMemo(() => {
    if (activeId === null || overId === null) return null;
    return getProjection(flattened, activeId, Number(overId), offsetX);
  }, [flattened, activeId, overId, offsetX]);

  const flashSaved = () => {
    setSaveState("saved");
    if (savedTimer.current) clearTimeout(savedTimer.current);
    savedTimer.current = setTimeout(() => setSaveState("idle"), 2200);
  };

  const persistTree = async (updatedNodes: OrgNode[]) => {
    setSaveState("saving");
    const ordered = flattenTree(updatedNodes);
    const counters = new Map<number | null, number>();
    const items = ordered.map((f) => {
      const key = f.parentId;
      const next = counters.get(key) ?? 0;
      counters.set(key, next + 1);
      return { id: f.id, parentId: f.parentId, sortOrder: next };
    });

    try {
      const res = await fetch("/api/org-nodes/reorder", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ items, departmentId }),
      });
      if (!res.ok) throw new Error("Reorder rejected");
      flashSaved();
      onNodesChange();
    } catch (e) {
      console.error(e);
      setSaveState("error");
      setLocalNodes(nodes); // roll back to server truth
    }
  };

  const handleDragStart = (event: DragStartEvent) => {
    setActiveId(Number(event.active.id));
    setOverId(event.active.id);
    setOffsetX(0);
  };

  const handleDragMove = (event: DragMoveEvent) => setOffsetX(event.delta.x);

  const handleDragOver = (event: DragOverEvent) => setOverId(event.over?.id ?? null);

  const handleDragEnd = (event: DragEndEvent) => {
    const currentProjection = projection;
    const activeNodeId = activeId;
    setActiveId(null);
    setOverId(null);
    setOffsetX(0);

    const { over } = event;
    if (!over || activeNodeId === null || !currentProjection) return;

    // Rebuild ordering from the flattened view, including hidden descendants
    const full = flattenTree(localNodes);
    const visible = removeDescendants(
      removeDescendants(full, [...collapsedIds]),
      [activeNodeId]
    );
    const overIndex = visible.findIndex((i) => i.id === Number(over.id));
    const activeFull = full.find((i) => i.id === activeNodeId);
    if (!activeFull || overIndex === -1) return;

    const noChange =
      activeFull.parentId === currentProjection.parentId && Number(over.id) === activeNodeId;
    if (noChange) return;

    // Order siblings: place active right after the item it was dropped over
    const workingOrder = visible.map((i) => i.id);
    workingOrder.splice(overIndex, 0, activeNodeId);

    const updated = localNodes.map((n) =>
      n.id === activeNodeId ? { ...n, parentId: currentProjection.parentId } : n
    );

    // Assign sort orders following the drop order within each parent bucket
    const rank = new Map<number, number>();
    workingOrder.forEach((id, i) => rank.set(id, i));
    const withOrder = updated.map((n) => ({
      ...n,
      sortOrder: rank.has(n.id) ? (rank.get(n.id) as number) : (n.sortOrder ?? 999),
    }));

    setLocalNodes(withOrder);
    void persistTree(withOrder);
  };

  const handleMoveEmployee = async (employeeId: number, nodeId: number | null) => {
    setSaveState("saving");
    try {
      const res = await fetch("/api/employees/assign", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ employeeId, orgNodeId: nodeId }),
      });
      if (!res.ok) throw new Error("assign failed");
      flashSaved();
      onEmployeesChange();
    } catch (e) {
      console.error(e);
      setSaveState("error");
    }
  };

  const toggleCollapse = (id: number) => {
    setCollapsedIds((prev) => {
      const next = new Set(prev);
      if (next.has(id)) next.delete(id);
      else next.add(id);
      return next;
    });
  };

  const handleAddNode = async () => {
    if (!newNode.name.trim()) return;
    setSaveState("saving");
    await fetch("/api/org-nodes", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        departmentId,
        parentId: addParentId,
        name: newNode.name.trim(),
        type: newNode.type,
        color: newNode.color,
        sortOrder: localNodes.length,
      }),
    });
    setShowAddModal(false);
    setNewNode({ name: "", type: "squad", color: "#3B82F6" });
    flashSaved();
    onNodesChange();
  };

  const handleEditNode = async () => {
    if (!editingNode) return;
    setSaveState("saving");
    await fetch("/api/org-nodes", {
      method: "PUT",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(editingNode),
    });
    setEditingNode(null);
    flashSaved();
    onNodesChange();
  };

  const handleDeleteNode = async (node: OrgNode) => {
    if (!confirm(`Delete "${node.name}" and all sub-units? Officers will be moved to Unassigned.`)) return;
    setSaveState("saving");
    await fetch(`/api/org-nodes?id=${node.id}`, { method: "DELETE" });
    flashSaved();
    onNodesChange();
    onEmployeesChange();
  };

  const activeNode = activeId !== null ? localNodes.find((n) => n.id === activeId) : null;
  const unassigned = employees.filter((e) => e.orgNodeId === null && e.isActive !== false);
  const fullFlat = flattenTree(localNodes);

  return (
    <div>
      {/* Toolbar */}
      <div className="flex flex-wrap justify-between items-center gap-3 mb-4">
        <div className="flex items-center gap-3">
          <p className="text-sm text-slate-500">
            Drag the <GripVertical size={13} className="inline -mt-0.5" /> handle to re-order or nest units
            (drag right to indent) • Drag officer chips between units
          </p>
          {saveState === "saving" && (
            <span className="flex items-center gap-1.5 text-xs text-blue-600 font-medium">
              <Loader2 size={13} className="animate-spin" /> Saving…
            </span>
          )}
          {saveState === "saved" && (
            <span className="flex items-center gap-1.5 text-xs text-green-600 font-medium">
              <Check size={13} /> Saved &amp; schedule updated
            </span>
          )}
          {saveState === "error" && (
            <span className="text-xs text-red-600 font-medium">Save failed — reverted</span>
          )}
        </div>
        <button
          onClick={() => {
            setAddParentId(null);
            setShowAddModal(true);
          }}
          className="flex items-center gap-2 px-4 py-2 bg-blue-600 text-white rounded-xl text-sm font-medium hover:bg-blue-700 transition-colors"
        >
          <Plus size={16} />
          Add District / Unit
        </button>
      </div>

      <DndContext
        sensors={sensors}
        collisionDetection={closestCenter}
        onDragStart={handleDragStart}
        onDragMove={handleDragMove}
        onDragOver={handleDragOver}
        onDragEnd={handleDragEnd}
        onDragCancel={() => {
          setActiveId(null);
          setOverId(null);
          setOffsetX(0);
        }}
      >
        <SortableContext items={sortableIds} strategy={verticalListSortingStrategy}>
          <div className="min-h-[80px]">
            {flattened.length === 0 && (
              <div className="text-center py-10 text-slate-400 border-2 border-dashed border-slate-200 rounded-xl">
                <Building2 size={32} className="mx-auto mb-2 opacity-40" />
                <p className="text-sm font-medium">No units yet</p>
                <p className="text-xs">Add a district to start building your org chart</p>
              </div>
            )}
            {flattened.map((item) => {
              const isActiveItem = item.id === activeId;
              const depth = isActiveItem && projection ? projection.depth : item.depth;
              const hasChildren = fullFlat.some((f) => f.parentId === item.id);
              return (
                <OrgRow
                  key={item.id}
                  item={item}
                  depth={depth}
                  employees={employees}
                  collapsed={collapsedIds.has(item.id)}
                  hasChildren={hasChildren}
                  onToggle={toggleCollapse}
                  onEdit={setEditingNode}
                  onDelete={handleDeleteNode}
                  onAddChild={(id) => {
                    setAddParentId(id);
                    setShowAddModal(true);
                  }}
                  onMoveEmployee={handleMoveEmployee}
                />
              );
            })}
          </div>
        </SortableContext>

        <DragOverlay dropAnimation={null}>
          {activeNode && (
            <div
              className="flex items-center gap-2 p-3 bg-white rounded-xl shadow-2xl border-2 border-blue-400"
              style={{ borderLeft: `4px solid ${activeNode.color ?? "#3B82F6"}` }}
            >
              <GripVertical size={16} className="text-slate-400" />
              <span className="font-semibold text-slate-800 text-sm">{activeNode.name}</span>
              <Badge color={activeNode.color ?? "#3B82F6"}>{activeNode.type}</Badge>
            </div>
          )}
        </DragOverlay>
      </DndContext>

      {/* Unassigned pool doubles as a drop target to un-assign officers */}
      <div
        onDragOver={(e) => {
          e.preventDefault();
          setUnassignedDropActive(true);
        }}
        onDragLeave={() => setUnassignedDropActive(false)}
        onDrop={(e) => {
          e.preventDefault();
          setUnassignedDropActive(false);
          const empId = parseInt(e.dataTransfer.getData("employeeId"));
          if (!Number.isNaN(empId)) handleMoveEmployee(empId, null);
        }}
        className={cn(
          "mt-6 rounded-xl border-2 border-dashed p-4 transition-colors",
          unassignedDropActive ? "border-amber-400 bg-amber-50" : "border-slate-200 bg-slate-50"
        )}
      >
        <div className="flex items-center gap-2 mb-3">
          <UserMinus size={15} className="text-slate-400" />
          <h4 className="text-sm font-semibold text-slate-600">Unassigned Officers</h4>
          <Badge color="#94A3B8">{unassigned.length}</Badge>
          <span className="text-xs text-slate-400 ml-auto">Drop here to remove from a unit</span>
        </div>
        <div className="flex flex-wrap gap-1.5">
          {unassigned.length === 0 && (
            <p className="text-xs text-slate-400 italic">Everyone is assigned to a unit.</p>
          )}
          {unassigned.map((emp) => (
            <div
              key={emp.id}
              draggable
              onDragStart={(e) => {
                e.dataTransfer.effectAllowed = "move";
                e.dataTransfer.setData("employeeId", String(emp.id));
              }}
              className="flex items-center gap-2 bg-white border border-slate-200 rounded-lg pl-1.5 pr-2 py-1 cursor-grab active:cursor-grabbing hover:border-blue-300"
            >
              <Avatar
                firstName={emp.firstName}
                lastName={emp.lastName}
                color={emp.avatarColor ?? "#3B82F6"}
                size="xs"
              />
              <div className="leading-tight">
                <p className="text-xs font-medium text-slate-700">
                  {emp.firstName} {emp.lastName}
                </p>
                <p className="text-[10px] text-slate-400">{emp.rank ?? emp.role}</p>
              </div>
            </div>
          ))}
        </div>
      </div>

      {/* Add modal */}
      <Modal isOpen={showAddModal} onClose={() => setShowAddModal(false)} title="Add Organizational Unit" size="sm">
        <div className="space-y-4">
          {addParentId !== null && (
            <div className="text-xs bg-blue-50 text-blue-700 rounded-lg px-3 py-2">
              Adding beneath <strong>{localNodes.find((n) => n.id === addParentId)?.name}</strong>
            </div>
          )}
          <div>
            <label className="block text-sm font-medium text-slate-700 mb-1">Unit Name</label>
            <input
              type="text"
              value={newNode.name}
              onChange={(e) => setNewNode({ ...newNode, name: e.target.value })}
              placeholder="e.g., District Alpha, Squad 1…"
              className="w-full border border-slate-300 rounded-xl px-3 py-2 text-sm focus:outline-none focus:ring-2 focus:ring-blue-500"
            />
          </div>
          <div>
            <label className="block text-sm font-medium text-slate-700 mb-1">Type</label>
            <select
              value={newNode.type}
              onChange={(e) => setNewNode({ ...newNode, type: e.target.value as (typeof NODE_TYPES)[number] })}
              className="w-full border border-slate-300 rounded-xl px-3 py-2 text-sm focus:outline-none focus:ring-2 focus:ring-blue-500"
            >
              {NODE_TYPES.map((t) => (
                <option key={t} value={t}>
                  {t.charAt(0).toUpperCase() + t.slice(1)}
                </option>
              ))}
            </select>
          </div>
          <div>
            <label className="block text-sm font-medium text-slate-700 mb-1">Color</label>
            <input
              type="color"
              value={newNode.color}
              onChange={(e) => setNewNode({ ...newNode, color: e.target.value })}
              className="w-12 h-10 rounded-lg border border-slate-300 cursor-pointer"
            />
          </div>
          <div className="flex gap-3 pt-2">
            <button
              onClick={handleAddNode}
              disabled={!newNode.name.trim()}
              className="flex-1 py-2 bg-blue-600 text-white rounded-xl text-sm font-medium hover:bg-blue-700 disabled:opacity-50"
            >
              Add Unit
            </button>
            <button
              onClick={() => setShowAddModal(false)}
              className="px-4 py-2 border border-slate-200 rounded-xl text-sm text-slate-600 hover:bg-slate-50"
            >
              Cancel
            </button>
          </div>
        </div>
      </Modal>

      {/* Edit modal */}
      <Modal isOpen={!!editingNode} onClose={() => setEditingNode(null)} title="Edit Unit" size="sm">
        {editingNode && (
          <div className="space-y-4">
            <div>
              <label className="block text-sm font-medium text-slate-700 mb-1">Name</label>
              <input
                type="text"
                value={editingNode.name}
                onChange={(e) => setEditingNode({ ...editingNode, name: e.target.value })}
                className="w-full border border-slate-300 rounded-xl px-3 py-2 text-sm focus:outline-none focus:ring-2 focus:ring-blue-500"
              />
            </div>
            <div>
              <label className="block text-sm font-medium text-slate-700 mb-1">Type</label>
              <select
                value={editingNode.type}
                onChange={(e) => setEditingNode({ ...editingNode, type: e.target.value as OrgNode["type"] })}
                className="w-full border border-slate-300 rounded-xl px-3 py-2 text-sm focus:outline-none focus:ring-2 focus:ring-blue-500"
              >
                {NODE_TYPES.map((t) => (
                  <option key={t} value={t}>
                    {t.charAt(0).toUpperCase() + t.slice(1)}
                  </option>
                ))}
              </select>
            </div>
            <div>
              <label className="block text-sm font-medium text-slate-700 mb-1">Color</label>
              <input
                type="color"
                value={editingNode.color ?? "#3B82F6"}
                onChange={(e) => setEditingNode({ ...editingNode, color: e.target.value })}
                className="w-12 h-10 rounded-lg border border-slate-300 cursor-pointer"
              />
            </div>
            <div className="flex gap-3 pt-2">
              <button
                onClick={handleEditNode}
                className="flex-1 py-2 bg-blue-600 text-white rounded-xl text-sm font-medium hover:bg-blue-700"
              >
                Save Changes
              </button>
              <button
                onClick={() => setEditingNode(null)}
                className="px-4 py-2 border border-slate-200 rounded-xl text-sm text-slate-600 hover:bg-slate-50"
              >
                Cancel
              </button>
            </div>
          </div>
        )}
      </Modal>
    </div>
  );
}