mRemotify/frontend/src/components/Sidebar/ConnectionTree.tsx

import React, { useEffect, useState, useCallback } from 'react';
import {
  Tree,
  Button,
  Dropdown,
  Typography,
  Tooltip,
  Input,
  message,
} from 'antd';
import type { DataNode, DirectoryTreeProps } from 'antd/es/tree';
import type { MenuProps } from 'antd';
import {
  FolderOutlined,
  FolderOpenOutlined,
  PlusOutlined,
  ReloadOutlined,
  WindowsOutlined,
  CodeOutlined,
  MoreOutlined,
  EditOutlined,
  DeleteOutlined,
  FolderAddOutlined,
} from '@ant-design/icons';
import { useStore } from '../../store';
import {
  apiFolderList,
  apiFolderCreate,
  apiFolderDelete,
  apiConnectionList,
  apiConnectionCreate,
  apiConnectionDelete,
  apiConnectionUpdate,
} from '../../api/client';
import { Connection, ConnectionFormValues, Folder } from '../../types';
import { ConnectionModal } from '../Modals/ConnectionModal';

interface TreeItemData {
  type: 'folder' | 'connection';
  folder?: Folder;
  connection?: Connection;
}

interface ExtendedDataNode extends DataNode {
  itemData: TreeItemData;
  children?: ExtendedDataNode[];
}

const OsIcon: React.FC<{ osType?: string | null }> = ({ osType }) => {
  if (osType === 'windows') return <WindowsOutlined />;
  return <CodeOutlined />;
};

function buildTree(
  folders: Folder[],
  connections: Connection[],
  parentId: string | null = null
): ExtendedDataNode[] {
  const subFolders: ExtendedDataNode[] = folders
    .filter((f) => f.parentId === parentId)
    .map((folder) => ({
      key: `folder-${folder.id}`,
      title: folder.name,
      isLeaf: false,
      itemData: { type: 'folder' as const, folder },
      children: buildTree(folders, connections, folder.id),
    }));

  const leafConnections: ExtendedDataNode[] = connections
    .filter((c) => c.folderId === parentId)
    .map((connection) => ({
      key: `connection-${connection.id}`,
      title: connection.name,
      isLeaf: true,
      icon: <OsIcon osType={connection.osType} />,
      itemData: { type: 'connection' as const, connection },
    }));

  return [...subFolders, ...leafConnections];
}

export const ConnectionTree: React.FC = () => {
  const folders = useStore((s) => s.folders);
  const connections = useStore((s) => s.connections);
  const setFolders = useStore((s) => s.setFolders);
  const setConnections = useStore((s) => s.setConnections);
  const openSession = useStore((s) => s.openSession);

  const [loading, setLoading] = useState(false);
  const [modalOpen, setModalOpen] = useState(false);
  const [editingConnection, setEditingConnection] = useState<Connection | null>(null);
  const [newFolderParentId, setNewFolderParentId] = useState<string | null>(null);
  const [addingFolder, setAddingFolder] = useState(false);
  const [newFolderName, setNewFolderName] = useState('');

  const refresh = useCallback(async () => {
    setLoading(true);
    try {
      const [fRes, cRes] = await Promise.all([apiFolderList(), apiConnectionList()]);
      setFolders(fRes.data);
      setConnections(cRes.data);
    } catch {
      message.error('Failed to load connections');
    } finally {
      setLoading(false);
    }
  }, [setFolders, setConnections]);

  useEffect(() => {
    refresh();
  }, [refresh]);

  const treeData = buildTree(folders, connections);

  // --- Drop: move connection into folder ---
  const onDrop: DirectoryTreeProps['onDrop'] = async (info) => {
    const dragNode = info.dragNode as unknown as ExtendedDataNode;
    const dropNode = info.node as unknown as ExtendedDataNode;
    if (dragNode.itemData.type !== 'connection') return;

    const connectionId = dragNode.itemData.connection!.id;
    let targetFolderId: string | null = null;

    if (dropNode.itemData.type === 'folder') {
      targetFolderId = dropNode.itemData.folder!.id;
    } else if (dropNode.itemData.type === 'connection') {
      targetFolderId = dropNode.itemData.connection!.folderId ?? null;
    }

    try {
      await apiConnectionUpdate(connectionId, { folderId: targetFolderId });
      await refresh();
    } catch {
      message.error('Failed to move connection');
    }
  };

  // --- Double-click: open session ---
  const onDoubleClick = (_: React.MouseEvent, node: DataNode) => {
    const ext = node as ExtendedDataNode;
    if (ext.itemData.type === 'connection') {
      openSession(ext.itemData.connection!);
    }
  };

  // --- Context menu items ---
  const getContextMenu = (node: ExtendedDataNode): MenuProps['items'] => {
    if (node.itemData.type === 'connection') {
      return [
        {
          key: 'connect',
          label: 'Connect',
          onClick: () => openSession(node.itemData.connection!),
        },
        {
          key: 'edit',
          icon: <EditOutlined />,
          label: 'Edit',
          onClick: () => {
            setEditingConnection(node.itemData.connection!);
            setModalOpen(true);
          },
        },
        { type: 'divider' },
        {
          key: 'delete',
          icon: <DeleteOutlined />,
          label: 'Delete',
          danger: true,
          onClick: async () => {
            try {
              await apiConnectionDelete(node.itemData.connection!.id);
              await refresh();
            } catch {
              message.error('Failed to delete connection');
            }
          },
        },
      ];
    }

    return [
      {
        key: 'addConnection',
        icon: <PlusOutlined />,
        label: 'Add connection here',
        onClick: () => {
          setEditingConnection(null);
          setModalOpen(true);
        },
      },
      {
        key: 'addSubfolder',
        icon: <FolderAddOutlined />,
        label: 'Add subfolder',
        onClick: () => {
          setNewFolderParentId(node.itemData.folder!.id);
          setAddingFolder(true);
        },
      },
      { type: 'divider' },
      {
        key: 'deleteFolder',
        icon: <DeleteOutlined />,
        label: 'Delete folder',
        danger: true,
        onClick: async () => {
          try {
            await apiFolderDelete(node.itemData.folder!.id);
            await refresh();
          } catch {
            message.error('Failed to delete folder');
          }
        },
      },
    ];
  };

  // --- Save connection (create or update) ---
  const handleSave = async (values: ConnectionFormValues, id?: string) => {
    try {
      if (id) {
        await apiConnectionUpdate(id, values);
      } else {
        await apiConnectionCreate(values);
      }
      setModalOpen(false);
      setEditingConnection(null);
      await refresh();
    } catch {
      message.error('Failed to save connection');
    }
  };

  // --- Add folder ---
  const commitNewFolder = async () => {
    if (!newFolderName.trim()) {
      setAddingFolder(false);
      setNewFolderName('');
      return;
    }
    try {
      await apiFolderCreate({ name: newFolderName.trim(), parentId: newFolderParentId });
      await refresh();
    } catch {
      message.error('Failed to create folder');
    } finally {
      setAddingFolder(false);
      setNewFolderName('');
      setNewFolderParentId(null);
    }
  };

  const titleRender = (node: DataNode) => {
    const ext = node as ExtendedDataNode;
    return (
      <Dropdown menu={{ items: getContextMenu(ext) }} trigger={['contextMenu']}>
        <span style={{ display: 'inline-flex', alignItems: 'center', gap: 4, width: '100%' }}>
          <span
            style={{
              flex: 1,
              overflow: 'hidden',
              textOverflow: 'ellipsis',
              whiteSpace: 'nowrap',
            }}
          >
            {String(node.title)}
          </span>
          <Dropdown menu={{ items: getContextMenu(ext) }} trigger={['click']}>
            <MoreOutlined
              onClick={(e) => e.stopPropagation()}
              style={{ opacity: 0.45, fontSize: 12 }}
            />
          </Dropdown>
        </span>
      </Dropdown>
    );
  };

  return (
    <div style={{ display: 'flex', flexDirection: 'column', height: '100%' }}>
      {/* Toolbar */}
      <div
        style={{
          padding: '8px 8px 4px',
          borderBottom: '1px solid #f0f0f0',
          display: 'flex',
          gap: 4,
        }}
      >
        <Tooltip title="New connection">
          <Button
            size="small"
            icon={<PlusOutlined />}
            onClick={() => {
              setEditingConnection(null);
              setModalOpen(true);
            }}
          />
        </Tooltip>
        <Tooltip title="New folder">
          <Button
            size="small"
            icon={<FolderAddOutlined />}
            onClick={() => {
              setNewFolderParentId(null);
              setAddingFolder(true);
            }}
          />
        </Tooltip>
        <Tooltip title="Refresh">
          <Button size="small" icon={<ReloadOutlined />} loading={loading} onClick={refresh} />
        </Tooltip>
      </div>

      {/* Inline folder name input */}
      {addingFolder && (
        <div style={{ padding: '4px 8px' }}>
          <Input
            size="small"
            autoFocus
            placeholder="Folder name"
            value={newFolderName}
            onChange={(e) => setNewFolderName(e.target.value)}
            onPressEnter={commitNewFolder}
            onBlur={commitNewFolder}
          />
        </div>
      )}

      {/* Tree */}
      <div style={{ flex: 1, overflow: 'auto', padding: '4px 0' }}>
        {treeData.length === 0 && !loading ? (
          <Typography.Text type="secondary" style={{ padding: '8px 16px', display: 'block' }}>
            No connections yet. Click + to add one.
          </Typography.Text>
        ) : (
          <Tree
            treeData={treeData}
            draggable={{ icon: false }}
            blockNode
            showIcon
            onDrop={onDrop}
            onDoubleClick={onDoubleClick}
            titleRender={titleRender}
            icon={(node) => {
              const ext = node as unknown as ExtendedDataNode;
              if (ext.itemData?.type === 'folder') {
                return (node as { expanded?: boolean }).expanded ? (
                  <FolderOpenOutlined />
                ) : (
                  <FolderOutlined />
                );
              }
              return null;
            }}
          />
        )}
      </div>

      <ConnectionModal
        open={modalOpen}
        connection={editingConnection}
        folders={folders}
        onClose={() => {
          setModalOpen(false);
          setEditingConnection(null);
        }}
        onSave={handleSave}
      />
    </div>
  );
};