import { useState, useCallback, useRef, useEffect, useMemo } from 'react';
import ReactFlow, {
  MiniMap,
  Controls,
  Background,
  useNodesState,
  useEdgesState,
  addEdge,
  Panel,
  useReactFlow,
  ReactFlowProvider,
  Handle,
  Position,
  BaseEdge,
  EdgeLabelRenderer,
  getSmoothStepPath,
  getBezierPath
} from 'reactflow';
import 'reactflow/dist/style.css';
import axios from 'axios';
import planPlusLogo from '../assets/img/planPlusLogo.png'; 

// Import icons from react-icons
import { FaDatabase, FaTable, FaFlask, FaArrowRight, FaPlus, FaTimes, FaChevronUp, FaChevronDown } from 'react-icons/fa';
import { BiSolidData } from 'react-icons/bi';
import { AiFillFolder } from 'react-icons/ai';
import { BsFileEarmarkSpreadsheet } from 'react-icons/bs';

// Import mock data from DataflowCanvas
import mockApiData from './mockData';

// Modal component for creating entities
const Modal = ({ isOpen, title, onClose, children }) => {
  if (!isOpen) return null;
  
  return (
    <div style={{
      position: 'fixed',
      top: 0,
      left: 0,
      right: 0,
      bottom: 0,
      backgroundColor: 'rgba(0, 0, 0, 0.5)',
      display: 'flex',
      alignItems: 'center',
      justifyContent: 'center',
      zIndex: 1000
    }}>
      <div style={{
        backgroundColor: 'white',
        borderRadius: '8px',
        width: '700px',
        maxWidth: '90%',
        maxHeight: '90vh',
        boxShadow: '0 4px 20px rgba(0, 0, 0, 0.2)',
        overflow: 'hidden',
        display: 'flex',
        flexDirection: 'column'
      }}>
        <div style={{
          display: 'flex',
          justifyContent: 'space-between',
          alignItems: 'center',
          padding: '16px 20px',
          borderBottom: '1px solid #eee',
          backgroundColor: '#f8f8f8'
        }}>
          <h3 style={{ margin: 0, color: '#333', fontSize: '18px' }}>{title}</h3>
          <button 
            onClick={onClose}
            style={{
              background: 'none',
              border: 'none',
              cursor: 'pointer',
              fontSize: '20px',
              color: '#666'
            }}
          >
            <FaTimes />
          </button>
        </div>
        <div style={{ 
          padding: '20px',
          overflowY: 'auto',
          flex: 1
        }}>
          {children}
        </div>
      </div>
    </div>
  );
};

// Database Creation Form with nested schema and table creation
const DatabaseForm = ({ onSave, onCancel }) => {
  const [formData, setFormData] = useState({
    name: '',
    description: '',
    url: '',
    key: '',
    type: 'PostgreSQL',
    schemas: [] // Array to hold schemas
  });
  
  const [showSchemaForm, setShowSchemaForm] = useState(false);
  const [currentSchemaIndex, setCurrentSchemaIndex] = useState(null);
  const [showTableForm, setShowTableForm] = useState(false);
  const [currentSchemaForTable, setCurrentSchemaForTable] = useState(null);
  
  const handleChange = (e) => {
    const { name, value } = e.target;
    setFormData(prev => ({ ...prev, [name]: value }));
  };
  
  const handleSubmit = (e) => {
    e.preventDefault();
    onSave(formData);
  };
  
  const fetchId = () => {
    // In a real app, this would make an API call
    alert('Fetching ID from URL: ' + formData.url);
  };
  
  // Add a new schema to the database
  const addSchema = (schemaData) => {
    if (currentSchemaIndex !== null) {
      // Edit existing schema
      const updatedSchemas = [...formData.schemas];
      updatedSchemas[currentSchemaIndex] = {
        ...updatedSchemas[currentSchemaIndex],
        ...schemaData,
      };
      setFormData(prev => ({ ...prev, schemas: updatedSchemas }));
    } else {
      // Add new schema
      const newSchema = {
        id: `schema-${Date.now()}`,
        name: schemaData.name,
        description: schemaData.description,
        tables: [] // Initialize with empty tables array
      };
      setFormData(prev => ({ 
        ...prev, 
        schemas: [...prev.schemas, newSchema] 
      }));
    }
    setShowSchemaForm(false);
    setCurrentSchemaIndex(null);
  };
  
  // Add a new table to a schema
  const addTable = (tableData) => {
    const schemaIndex = formData.schemas.findIndex(s => s.id === currentSchemaForTable);
    if (schemaIndex !== -1) {
      const newTable = {
        id: `table-${Date.now()}`,
        name: tableData.name,
        description: tableData.description,
        type: tableData.type
      };
      
      const updatedSchemas = [...formData.schemas];
      updatedSchemas[schemaIndex] = {
        ...updatedSchemas[schemaIndex],
        tables: [...updatedSchemas[schemaIndex].tables, newTable]
      };
      
      setFormData(prev => ({ ...prev, schemas: updatedSchemas }));
    }
    setShowTableForm(false);
    setCurrentSchemaForTable(null);
  };
  
  // Remove a schema
  const removeSchema = (index) => {
    const updatedSchemas = [...formData.schemas];
    updatedSchemas.splice(index, 1);
    setFormData(prev => ({ ...prev, schemas: updatedSchemas }));
  };
  
  // Remove a table
  const removeTable = (schemaIndex, tableIndex) => {
    const updatedSchemas = [...formData.schemas];
    updatedSchemas[schemaIndex].tables.splice(tableIndex, 1);
    setFormData(prev => ({ ...prev, schemas: updatedSchemas }));
  };
  
  return (
    <div>
      {showSchemaForm ? (
        <div>
          <div style={{ 
            display: 'flex', 
            justifyContent: 'space-between', 
            alignItems: 'center',
            marginBottom: '16px'
          }}>
            <h4 style={{ margin: 0 }}>
              {currentSchemaIndex !== null ? 'Edit Schema' : 'Add Schema'}
            </h4>
            <button
              onClick={() => {
                setShowSchemaForm(false);
                setCurrentSchemaIndex(null);
              }}
              style={{
                background: 'none',
                border: 'none',
                fontSize: '18px',
                cursor: 'pointer',
                color: '#666'
              }}
            >
              ×
            </button>
          </div>
          
          <SchemaForm 
            onSave={addSchema} 
            onCancel={() => {
              setShowSchemaForm(false);
              setCurrentSchemaIndex(null);
            }}
            initialData={currentSchemaIndex !== null ? formData.schemas[currentSchemaIndex] : null}
          />
        </div>
      ) : showTableForm ? (
        <div>
          <div style={{ 
            display: 'flex', 
            justifyContent: 'space-between', 
            alignItems: 'center',
            marginBottom: '16px'
          }}>
            <h4 style={{ margin: 0 }}>Add Table</h4>
            <button
              onClick={() => {
                setShowTableForm(false);
                setCurrentSchemaForTable(null);
              }}
              style={{
                background: 'none',
                border: 'none',
                fontSize: '18px',
                cursor: 'pointer',
                color: '#666'
              }}
            >
              ×
            </button>
          </div>
          
          <TableForm 
            onSave={addTable} 
            onCancel={() => {
              setShowTableForm(false);
              setCurrentSchemaForTable(null);
            }}
          />
        </div>
      ) : (
        <form onSubmit={handleSubmit}>
          <div style={{ marginBottom: '16px' }}>
            <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
              Name
            </label>
            <input
              type="text"
              name="name"
              value={formData.name}
              onChange={handleChange}
              placeholder="Name"
              style={{
                width: '100%',
                padding: '10px',
                borderRadius: '4px',
                border: '1px solid #ddd',
                fontSize: '14px'
              }}
              required
            />
          </div>
          
          <div style={{ marginBottom: '16px' }}>
            <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
              Description
            </label>
            <textarea
              name="description"
              value={formData.description}
              onChange={handleChange}
              placeholder="Description"
              style={{
                width: '100%',
                padding: '10px',
                borderRadius: '4px',
                border: '1px solid #ddd',
                fontSize: '14px',
                minHeight: '60px',
                resize: 'vertical'
              }}
            />
          </div>
          
          <div style={{ marginBottom: '16px' }}>
            <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
              URL
            </label>
            <div style={{ display: 'flex', gap: '10px' }}>
              <input
                type="url"
                name="url"
                value={formData.url}
                onChange={handleChange}
                placeholder="http://"
                style={{
                  flex: 1,
                  padding: '10px',
                  borderRadius: '4px',
                  border: '1px solid #ddd',
                  fontSize: '14px'
                }}
              />
              <button
                type="button"
                onClick={fetchId}
                style={{
                  padding: '10px 15px',
                  backgroundColor: '#1890ff',
                  color: 'white',
                  border: 'none',
                  borderRadius: '4px',
                  cursor: 'pointer',
                  fontSize: '14px'
                }}
              >
                Fetch ID
              </button>
            </div>
          </div>
          
          <div style={{ marginBottom: '16px' }}>
            <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
              Key
            </label>
            <input
              type="text"
              name="key"
              value={formData.key}
              onChange={handleChange}
              placeholder="http://"
              style={{
                width: '100%',
                padding: '10px',
                borderRadius: '4px',
                border: '1px solid #ddd',
                fontSize: '14px'
              }}
            />
          </div>
          
          <div style={{ marginBottom: '16px' }}>
            <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
              Database Type
            </label>
            <select
              name="type"
              value={formData.type}
              onChange={handleChange}
              style={{
                width: '100%',
                padding: '10px',
                borderRadius: '4px',
                border: '1px solid #ddd',
                fontSize: '14px',
                backgroundColor: 'white'
              }}
            >
              <option value="PostgreSQL">PostgreSQL</option>
              <option value="MySQL">MySQL</option>
              <option value="Oracle">Oracle</option>
              <option value="SQL Server">SQL Server</option>
              <option value="MongoDB">MongoDB</option>
            </select>
          </div>
          
          {/* Schemas Section removed as per requirements */}
          
          <div style={{ display: 'flex', justifyContent: 'flex-end', gap: '10px', marginTop: '24px' }}>
            <button
              type="button"
              onClick={onCancel}
              style={{
                padding: '10px 20px',
                backgroundColor: '#f5f5f5',
                color: '#333',
                border: '1px solid #ddd',
                borderRadius: '4px',
                cursor: 'pointer',
                fontSize: '14px'
              }}
            >
              Cancel
            </button>
            <button
              type="submit"
              style={{
                padding: '10px 20px',
                backgroundColor: '#00a99d',
                color: 'white',
                border: 'none',
                borderRadius: '4px',
                cursor: 'pointer',
                fontSize: '14px'
              }}
            >
              Save Database
            </button>
          </div>
        </form>
      )}
    </div>
  );
};

// Schema Creation Form
const SchemaForm = ({ onSave, onCancel, parentDatabase, initialData }) => {
  const [formData, setFormData] = useState({
    name: initialData ? initialData.name : '',
    description: initialData ? initialData.description : ''
  });
  
  const handleChange = (e) => {
    const { name, value } = e.target;
    setFormData(prev => ({ ...prev, [name]: value }));
  };
  
  const handleSubmit = (e) => {
    e.preventDefault();
    onSave(parentDatabase ? { ...formData, parentDatabase } : formData);
  };
  
  return (
    <form onSubmit={handleSubmit}>
      <div style={{ marginBottom: '16px' }}>
        <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
          Schema Name
        </label>
        <input
          type="text"
          name="name"
          value={formData.name}
          onChange={handleChange}
          placeholder="Schema Name"
          style={{
            width: '100%',
            padding: '10px',
            borderRadius: '4px',
            border: '1px solid #ddd',
            fontSize: '14px'
          }}
          required
        />
      </div>
      
      <div style={{ marginBottom: '24px' }}>
        <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
          Description
        </label>
        <textarea
          name="description"
          value={formData.description}
          onChange={handleChange}
          placeholder="Schema Description"
          style={{
            width: '100%',
            padding: '10px',
            borderRadius: '4px',
            border: '1px solid #ddd',
            fontSize: '14px',
            minHeight: '60px',
            resize: 'vertical'
          }}
        />
      </div>
      
      <div style={{ display: 'flex', justifyContent: 'flex-end', gap: '10px' }}>
        <button
          type="button"
          onClick={onCancel}
          style={{
            padding: '10px 20px',
            backgroundColor: '#f5f5f5',
            color: '#333',
            border: '1px solid #ddd',
            borderRadius: '4px',
            cursor: 'pointer',
            fontSize: '14px'
          }}
        >
          Cancel
        </button>
        <button
          type="submit"
          style={{
            padding: '10px 20px',
            backgroundColor: '#00a99d',
            color: 'white',
            border: 'none',
            borderRadius: '4px',
            cursor: 'pointer',
            fontSize: '14px'
          }}
        >
          {initialData ? 'Update Schema' : 'Add Schema'}
        </button>
      </div>
    </form>
  );
};

// Entity Selector Component
const EntitySelector = ({ 
  activeEntityType,
  setActiveEntityType,
  databases,
  schemas,
  selectedDatabaseForSchema,
  setSelectedDatabaseForSchema,
  selectedSchemaForTable,
  setSelectedSchemaForTable,
  onSaveDatabase,
  onSaveSchema,
  onSaveTable,
  onCancel
}) => {
  // Determine which databases are available for schema creation
  const availableDatabases = databases || [];
  
  // Determine which schemas are available for table creation
  const availableSchemas = schemas ? schemas.filter(schema => 
    selectedDatabaseForSchema ? schema.dbId === selectedDatabaseForSchema : true
  ) : [];
  
  return (
    <div>
      {/* Entity Type Selector - Only Database */}
      <div style={{ 
        display: 'flex', 
        justifyContent: 'center', 
        marginBottom: '20px',
        borderRadius: '8px',
        overflow: 'hidden',
        border: '1px solid #eee'
      }}>
        <button
          onClick={() => {
            setActiveEntityType('database');
            setSelectedDatabaseForSchema(null);
            setSelectedSchemaForTable(null);
          }}
          style={{
            flex: 1,
            padding: '12px',
            border: 'none',
            background: '#00a99d',
            color: 'white',
            cursor: 'pointer',
            display: 'flex',
            flexDirection: 'column',
            alignItems: 'center',
            gap: '5px',
            transition: 'all 0.2s'
          }}
        >
          <svg width="20" height="20" viewBox="0 0 41 39" fill="none" xmlns="http://www.w3.org/2000/svg">
            <rect width="40.5" height="38.7" rx="3.6" fill="url(#db_paint0_linear)"/>
            <path d="M19.8845 24.789C17.0714 24.789 14.2465 24.0672 12.912 22.5097C12.8725 22.683 12.8462 22.8607 12.8462 23.0493V25.8844C12.8462 28.3962 16.4937 29.5392 19.8845 29.5392C23.2753 29.5392 26.9228 28.3962 26.9228 25.8844V23.0493C26.9228 22.8607 26.8964 22.6837 26.8569 22.5104C25.5217 24.068 22.6976 24.7904 19.8837 24.7904L19.8845 24.789ZM19.8845 19.3083C17.0714 19.3083 14.2465 18.5865 12.9127 17.0289C12.8706 17.2053 12.8486 17.3858 12.8469 17.5671V20.4022C12.8469 22.9133 16.4937 24.0563 19.8845 24.0563C23.2753 24.0563 26.9228 22.9133 26.9228 20.4015V17.5657C26.9228 17.3778 26.8964 17.2001 26.8569 17.0268C25.5217 18.5843 22.6976 19.3068 19.8837 19.3068L19.8845 19.3083ZM19.8845 8.42944C16.4937 8.42944 12.8462 9.57385 12.8462 12.0857V14.9208C12.8462 17.4326 16.4937 18.5755 19.8845 18.5755C23.2753 18.5755 26.9228 17.4333 26.9228 14.9215V12.0864C26.9228 9.57458 23.2753 8.43017 19.8845 8.43017V8.42944ZM19.8845 14.2794C16.8059 14.2794 14.3087 13.2974 14.3087 12.0857C14.3087 10.8747 16.8052 9.89194 19.8845 9.89194C22.9638 9.89194 25.4603 10.8747 25.4603 12.0857C25.4603 13.2981 22.9638 14.2794 19.8845 14.2794Z" fill="white"/>
            <defs>
              <linearGradient id="db_paint0_linear" x1="40.5" y1="19.35" x2="0" y2="19.35" gradientUnits="userSpaceOnUse">
                <stop stopColor="#006064"/>
                <stop offset="0.711538" stopColor="#00A1A8"/>
              </linearGradient>
            </defs>
          </svg>
          {/* <span>Database</span> */}
        </button>
      </div>
      
      {/* No parent selectors needed since we only have database */}
      
      {/* Only Database Form */}
      <DatabaseForm 
        onSave={onSaveDatabase}
        onCancel={onCancel}
      />
      
      {/* Guidance message when no parents are available */}
      {activeEntityType === 'schema' && availableDatabases.length === 0 && (
        <div style={{ 
          padding: '30px', 
          textAlign: 'center', 
          border: '1px dashed #ddd',
          borderRadius: '4px',
          color: '#666',
          backgroundColor: '#fafafa'
        }}>
          You need to create a database first before creating a schema.
        </div>
      )}
      
      {activeEntityType === 'table' && availableSchemas.length === 0 && (
        <div style={{ 
          padding: '30px', 
          textAlign: 'center', 
          border: '1px dashed #ddd',
          borderRadius: '4px',
          color: '#666',
          backgroundColor: '#fafafa'
        }}>
          You need to create a schema first before creating a table.
        </div>
      )}
    </div>
  );
};

// Table Creation Form
const TableForm = ({ onSave, onCancel, parentSchema, initialData }) => {
  const [formData, setFormData] = useState({
    name: initialData ? initialData.name : '',
    description: initialData ? initialData.description : '',
    type: initialData ? initialData.type : 'dimension',
    columns: initialData && initialData.columns ? initialData.columns : []
  });
  
  const [showColumnInput, setShowColumnInput] = useState(false);
  const [columnName, setColumnName] = useState('');
  
  const handleChange = (e) => {
    const { name, value } = e.target;
    setFormData(prev => ({ ...prev, [name]: value }));
  };
  
  const handleSubmit = (e) => {
    e.preventDefault();
    onSave(parentSchema ? { ...formData, parentSchema } : formData);
  };
  
  const addColumn = () => {
    if (columnName.trim()) {
      setFormData(prev => ({
        ...prev,
        columns: [...prev.columns, columnName.trim()]
      }));
      setColumnName('');
      setShowColumnInput(false);
    }
  };
  
  const removeColumn = (index) => {
    const updatedColumns = [...formData.columns];
    updatedColumns.splice(index, 1);
    setFormData(prev => ({ ...prev, columns: updatedColumns }));
  };
  
  return (
    <form onSubmit={handleSubmit}>
      <div style={{ marginBottom: '16px' }}>
        <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
          Table Name
        </label>
        <input
          type="text"
          name="name"
          value={formData.name}
          onChange={handleChange}
          placeholder="Table Name"
          style={{
            width: '100%',
            padding: '10px',
            borderRadius: '4px',
            border: '1px solid #ddd',
            fontSize: '14px'
          }}
          required
        />
      </div>
      
      <div style={{ marginBottom: '16px' }}>
        <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
          Table Description
        </label>
        <textarea
          name="description"
          value={formData.description}
          onChange={handleChange}
          placeholder="Table Description"
          style={{
            width: '100%',
            padding: '10px',
            borderRadius: '4px',
            border: '1px solid #ddd',
            fontSize: '14px',
            minHeight: '60px',
            resize: 'vertical'
          }}
        />
      </div>
      
      <div style={{ marginBottom: '16px' }}>
        <label style={{ display: 'block', marginBottom: '6px', fontWeight: '500', color: '#333' }}>
          Table Type
        </label>
        <select
          name="type"
          value={formData.type}
          onChange={handleChange}
          style={{
            width: '100%',
            padding: '10px',
            borderRadius: '4px',
            border: '1px solid #ddd',
            fontSize: '14px',
            backgroundColor: 'white'
          }}
        >
          <option value="dimension">Dimension</option>
          <option value="fact">Fact</option>
          <option value="stage">Stage</option>
        </select>
      </div>
      
      {/* Columns Section */}
      <div style={{ marginBottom: '24px' }}>
        <div style={{ 
          display: 'flex', 
          justifyContent: 'space-between', 
          alignItems: 'center',
          marginBottom: '10px'
        }}>
          <label style={{ fontWeight: '500', color: '#333' }}>
            Columns
          </label>
          <button
            type="button"
            onClick={() => setShowColumnInput(true)}
            style={{
              padding: '5px 10px',
              backgroundColor: '#1890ff',
              color: 'white',
              border: 'none',
              borderRadius: '4px',
              cursor: 'pointer',
              fontSize: '12px',
              display: 'flex',
              alignItems: 'center',
              gap: '5px'
            }}
          >
            <FaPlus size={10} /> Add Column
          </button>
        </div>
        
        {showColumnInput && (
          <div style={{ 
            display: 'flex', 
            gap: '10px',
            marginBottom: '10px'
          }}>
            <input
              type="text"
              value={columnName}
              onChange={(e) => setColumnName(e.target.value)}
              placeholder="Column name"
              style={{
                flex: 1,
                padding: '8px',
                borderRadius: '4px',
                border: '1px solid #ddd',
                fontSize: '14px'
              }}
            />
            <button
              type="button"
              onClick={addColumn}
              style={{
                padding: '8px 15px',
                backgroundColor: '#52c41a',
                color: 'white',
                border: 'none',
                borderRadius: '4px',
                cursor: 'pointer',
                fontSize: '14px'
              }}
            >
              Add
            </button>
            <button
              type="button"
              onClick={() => {
                setShowColumnInput(false);
                setColumnName('');
              }}
              style={{
                padding: '8px 15px',
                backgroundColor: '#f5f5f5',
                color: '#333',
                border: '1px solid #ddd',
                borderRadius: '4px',
                cursor: 'pointer',
                fontSize: '14px'
              }}
            >
              Cancel
            </button>
          </div>
        )}
        
        {formData.columns.length > 0 ? (
          <div style={{ 
            border: '1px solid #eee', 
            borderRadius: '4px',
            maxHeight: '150px',
            overflowY: 'auto'
          }}>
            {formData.columns.map((column, index) => (
              <div 
                key={index} 
                style={{ 
                  padding: '8px 12px',
                  display: 'flex',
                  justifyContent: 'space-between',
                  alignItems: 'center',
                  borderBottom: index < formData.columns.length - 1 ? '1px solid #f0f0f0' : 'none',
                  backgroundColor: index % 2 === 0 ? '#f9f9f9' : 'white'
                }}
              >
                <span>{column}</span>
                <button
                  type="button"
                  onClick={() => removeColumn(index)}
                  style={{
                    background: 'none',
                    border: 'none',
                    color: '#ff4d4f',
                    cursor: 'pointer',
                    fontSize: '14px'
                  }}
                >
                  ×
                </button>
              </div>
            ))}
          </div>
        ) : (
          <div style={{ 
            padding: '15px', 
            textAlign: 'center', 
            border: '1px dashed #ddd',
            borderRadius: '4px',
            color: '#999',
            backgroundColor: '#fafafa',
            fontSize: '13px'
          }}>
            No columns added yet. Click "Add Column" to create one.
          </div>
        )}
      </div>
      
      <div style={{ display: 'flex', justifyContent: 'flex-end', gap: '10px' }}>
        <button
          type="button"
          onClick={onCancel}
          style={{
            padding: '10px 20px',
            backgroundColor: '#f5f5f5',
            color: '#333',
            border: '1px solid #ddd',
            borderRadius: '4px',
            cursor: 'pointer',
            fontSize: '14px'
          }}
        >
          Cancel
        </button>
        <button
          type="submit"
          style={{
            padding: '10px 20px',
            backgroundColor: '#00a99d',
            color: 'white',
            border: 'none',
            borderRadius: '4px',
            cursor: 'pointer',
            fontSize: '14px'
          }}
        >
          {initialData ? 'Update Table' : 'Add Table'}
        </button>
      </div>
    </form>
  );
};

// Custom edge with animated arrow
const CustomEdge = ({ id, source, target, sourceX, sourceY, targetX, targetY, sourcePosition, targetPosition, style = {}, data, markerEnd }) => {
  const [edgePath, labelX, labelY] = getBezierPath({
    sourceX,
    sourceY,
    sourcePosition,
    targetX,
    targetY,
    targetPosition,
  });

  return (
    <>
      <BaseEdge path={edgePath} markerEnd={markerEnd} style={style} />
      <EdgeLabelRenderer>
        <div
          style={{
            position: 'absolute',
            transform: `translate(-50%, -50%) translate(${labelX}px,${labelY}px)`,
            fontSize: 12,
            pointerEvents: 'all',
            backgroundColor: style.stroke || '#1a1a1a',
            padding: '2px 5px',
            borderRadius: '4px',
            color: 'white',
            fontWeight: 500,
            boxShadow: '0 0 5px rgba(0,0,0,0.3)',
            display: 'flex',
            alignItems: 'center',
            gap: '4px'
          }}
          className="nodrag nopan"
        >
          <FaArrowRight size={10} />
          {data?.label || 'connects to'}
        </div>
      </EdgeLabelRenderer>
    </>
  );
};

// Create a unique ID for edges that should be animated
const SERVICE_EDGE_CLASS = 'service-db-connection';

// Add the CSS for the animated connections directly to the stylesheet
// This avoids potential issues with React's rendering cycle
const style = document.createElement('style');
style.textContent = `
  @keyframes dashdraw {
    from {
      stroke-dashoffset: 10;
    }
    to {
      stroke-dashoffset: 0;
    }
  }
  .${SERVICE_EDGE_CLASS} .react-flow__edge-path {
    animation: dashdraw 0.5s linear infinite;
  }
`;
// Add the style element only once
if (!document.getElementById('react-flow-animation')) {
  style.id = 'react-flow-animation';
  document.head.appendChild(style);
}

// Hierarchical Edge for connecting service to databases
const HierarchicalEdge = ({ id, source, target, sourceX, sourceY, targetX, targetY, style = {} }) => {
  // Calculate the path points for a stepped edge with orthogonal lines
  const [path] = getSmoothStepPath({
    sourceX,
    sourceY,
    sourcePosition: Position.Bottom,  // Always start from bottom
    targetX,
    targetY,
    targetPosition: Position.Top,     // Always connect to top
    borderRadius: 20,                 // Smoother corners for longer connections
    stepSize: 30,                     // Increased step size for longer connections
  });

  // Enhanced style for hierarchical connections with animated dotted line
  const hierarchyStyle = {
    strokeWidth: 2,
    stroke: '#00a99d',
    strokeDasharray: '5, 5', // Create dotted/dashed line effect
    ...style,
  };

  return (
    <g className={SERVICE_EDGE_CLASS}>
      <BaseEdge
        id={id}
        path={path}
        style={hierarchyStyle}
        markerEnd={{
          type: 'arrowclosed',
          width: 16,
          height: 16,
          color: '#00a99d',
          strokeWidth: 2,
        }}
      />
    </g>
  );
};

// Custom node types
const ServiceNode = ({ data = {} }) => {
  // Safety check for undefined data
  if (!data) {
    console.error('ServiceNode received undefined data');
    data = {}; // Provide a default empty object
  }
  
  return (
    <div style={{
      padding: '20px',
      borderRadius: '12px',
      background: 'linear-gradient(145deg, #222, #2a2a2a)',
      boxShadow: '0 10px 30px rgba(0, 0, 0, 0.3), 0 0 0 1px rgba(0, 169, 157, 0.3) inset',
      border: '2px solid #00a99d',
      color: '#fff',
      width: '240px',
      minHeight: '110px',
      display: 'flex',
      flexDirection: 'column',
      position: 'relative',
      fontFamily: 'Inter, sans-serif',
      justifyContent: 'space-between'
    }}>
      <div style={{
        display: 'flex',
        flexDirection: 'column',
        alignItems: 'center',
        marginBottom: '16px'
      }}>
        {/* Logo and text container */}
        <div style={{
          display: 'flex',
          alignItems: 'center',
          justifyContent: 'center',
          marginBottom: '8px'
        }}>
          <img 
            src={planPlusLogo} 
            alt="Logo" 
            width="85" 
            style={{ 
              marginRight: '2px', 
              verticalAlign: 'middle' 
            }} 
          />
          <span style={{ 
            fontWeight: '600', 
            fontSize: '20px',
            color: '#fff',
            marginLeft: '2px',
            marginTop: '2px'
          }}>
            Plan+
          </span>
        </div>
        
        {/* Connected databases count */}
        <div style={{ 
          fontSize: '12px',
          color: '#aaa',
          marginTop: '2px',
          textAlign: 'center'
        }}>
          Connected Databases: {data.databases || 0}
        </div>
      </div>
      
      <div style={{
        display: 'flex',
        justifyContent: 'center',
        alignItems: 'center',
        marginTop: '8px'
      }}>
        <button
          style={{
            background: 'linear-gradient(145deg, rgba(0, 169, 157, 0.2), rgba(0, 169, 157, 0.1))',
            border: '1px solid rgba(0, 169, 157, 0.4)',
            borderRadius: '6px',
            padding: '8px 16px',
            color: '#00a99d',
            fontSize: '12px',
            fontWeight: '500',
            cursor: 'pointer',
            display: 'flex',
            alignItems: 'center',
            gap: '6px',
            boxShadow: '0 2px 5px rgba(0, 0, 0, 0.2)',
            transition: 'all 0.2s ease',
            width: '100%',
            justifyContent: 'center'
          }}
          onClick={() => data.onToggle && data.onToggle(data.id)}
        >
          {data.expanded ? (
            <>
              <span>Collapse Connections</span>
              <FaChevronUp size={10} />
            </>
          ) : (
            <>
              <span>Expand Connections</span>
              <FaChevronDown size={10} />
            </>
          )}
        </button>
      </div>
      
      {/* Handles for connections - explicit id for targeting */}
      <Handle
        id="bottom"
        type="source"
        position={Position.Bottom}
        style={{ 
          background: '#00a99d', 
          width: '8px', 
          height: '8px',
          border: '2px solid #00a99d',
          boxShadow: '0 0 5px rgba(0, 169, 157, 0.5)',
          bottom: '-5px' // Move it slightly outside the node for better connection
        }}
      />
    </div>
  );
};

const DatabaseNode = ({ data = {} }) => {
  // Safety check for undefined data
  if (!data) {
    console.error('DatabaseNode received undefined data');
    data = {}; // Provide a default empty object
  }
  
  // Enhanced styles for handles used in hierarchy connections
  
  // Debug: Log the data being received by the component
  // console.log('DatabaseNode data:', data);
  
  // Assign different colors to different databases
  const colorMap = {
    'my_dwh': '#00a99d',   // First database - teal
    'my_dwh2': '#1890ff',  // Second database - blue
    'default': '#9c27b0'   // Default - purple
  };
  
  // Get the database slug from the id (format: db-{slug})
  // Add a safety check for undefined data.id
  // Use data.id if available, otherwise fall back to the node's id
  const nodeId = (data && data.id) ? data.id : 'default';
  const dbSlug = (nodeId && typeof nodeId === 'string' && nodeId.startsWith && nodeId.startsWith('db-')) 
    ? nodeId.substring(3) 
    : 'default';
  
  // Use the color from the map or default to a fallback color
  const borderColor = colorMap[dbSlug] || '#ff9800';
  const handleColor = borderColor;
  const bgColor = '#1a1a1a';
  
  // Create a safe gradient ID by removing any special characters
  const safeDbSlug = (dbSlug && typeof dbSlug === 'string') 
    ? dbSlug.replace(/[^a-zA-Z0-9]/g, '_') 
    : 'default';
  const gradientId = `db_paint0_linear_${safeDbSlug}`;
  
  // Database SVG icon
  const DatabaseIcon = () => (
    <svg width="30" height="28" viewBox="0 0 41 39" fill="none" xmlns="http://www.w3.org/2000/svg">
      <rect width="40.5" height="38.7" rx="3.6" fill={`url(#${gradientId})`}/>
      <path d="M19.8845 24.789C17.0714 24.789 14.2465 24.0672 12.912 22.5097C12.8725 22.683 12.8462 22.8607 12.8462 23.0493V25.8844C12.8462 28.3962 16.4937 29.5392 19.8845 29.5392C23.2753 29.5392 26.9228 28.3962 26.9228 25.8844V23.0493C26.9228 22.8607 26.8964 22.6837 26.8569 22.5104C25.5217 24.068 22.6976 24.7904 19.8837 24.7904L19.8845 24.789ZM19.8845 19.3083C17.0714 19.3083 14.2465 18.5865 12.9127 17.0289C12.8706 17.2053 12.8486 17.3858 12.8469 17.5671V20.4022C12.8469 22.9133 16.4937 24.0563 19.8845 24.0563C23.2753 24.0563 26.9228 22.9133 26.9228 20.4015V17.5657C26.9228 17.3778 26.8964 17.2001 26.8569 17.0268C25.5217 18.5843 22.6976 19.3068 19.8837 19.3068L19.8845 19.3083ZM19.8845 8.42944C16.4937 8.42944 12.8462 9.57385 12.8462 12.0857V14.9208C12.8462 17.4326 16.4937 18.5755 19.8845 18.5755C23.2753 18.5755 26.9228 17.4333 26.9228 14.9215V12.0864C26.9228 9.57458 23.2753 8.43017 19.8845 8.43017V8.42944ZM19.8845 14.2794C16.8059 14.2794 14.3087 13.2974 14.3087 12.0857C14.3087 10.8747 16.8052 9.89194 19.8845 9.89194C22.9638 9.89194 25.4603 10.8747 25.4603 12.0857C25.4603 13.2981 22.9638 14.2794 19.8845 14.2794Z" fill="white"/>
      <defs>
        <linearGradient id={gradientId} x1="40.5" y1="19.35" x2="0" y2="19.35" gradientUnits="userSpaceOnUse">
          <stop stopColor={borderColor}/>
          <stop offset="0.711538" stopColor={handleColor}/>
        </linearGradient>
      </defs>
    </svg>
  );
  
  return (
    <div className="database-node" style={{ 
      padding: '10px', 
      borderRadius: '5px', 
      background: bgColor, 
      border: `1px solid ${borderColor}`,
      width: '180px',
      position: 'relative',
      boxShadow: '0 0 10px rgba(0, 0, 0, 0.3)',
      color: '#ffffff'
    }}>
      {/* Connection handles */}
      <Handle 
        type="source" 
        position={Position.Right} 
        id="right" 
        style={{ background: handleColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      <Handle 
        type="target" 
        position={Position.Left} 
        id="left" 
        style={{ background: handleColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      <Handle 
        type="source" 
        position={Position.Bottom} 
        id="bottom" 
        style={{ background: handleColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      <Handle 
        type="target" 
        position={Position.Top} 
        id="top" 
        style={{ background: handleColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      
      <div style={{ 
        fontWeight: 'bold', 
        marginBottom: '5px', 
        display: 'flex', 
        alignItems: 'center'
      }}>
        <span style={{ marginRight: '8px' }}>
          <DatabaseIcon />
        </span>
        {/* Display the database name from the API response */}
        <span title={data.name || data.label || 'Unknown Database'} style={{ overflow: 'hidden', textOverflow: 'ellipsis', whiteSpace: 'nowrap', maxWidth: '120px' }}>
          {data.name || data.label || 'Unknown Database'}
        </span>
      </div>
      
      {/* <div style={{ fontSize: '0.8em', color: '#aaa', marginBottom: '5px' }}>
        {`${data.schemas} schemas • ${data.tables} tables`}
      </div>
      
      <div style={{ fontSize: '0.7em', color: '#888', marginBottom: '10px' }}>
        Connection: {dbSlug}
      </div> */}
      
      {/* View Details Button */}
      <div style={{ display: 'flex', justifyContent: 'center' }}>
        <button
          onClick={(event) => {
            // Stop propagation to prevent the node click handler from being triggered
            event.stopPropagation();
            // Call the viewDetails function passed in data with safety checks
            if (data && data.onViewDetails && data.id) {
              data.onViewDetails(data.id, data.name || data.label || 'Unknown Database');
            } else {
              console.warn('Cannot view details: missing required data properties');
            }
          }}
          style={{
            padding: '4px 8px',
            backgroundColor: borderColor,
            color: 'white',
            border: 'none',
            borderRadius: '3px',
            cursor: 'pointer',
            fontSize: '11px',
            width: '100%',
            display: 'flex',
            alignItems: 'center',
            justifyContent: 'center',
            gap: '5px',
            transition: 'all 0.2s'
          }}
        >
          <svg width="12" height="12" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
            <path d="M12 4.5C7 4.5 2.73 7.61 1 12C2.73 16.39 7 19.5 12 19.5C17 19.5 21.27 16.39 23 12C21.27 7.61 17 4.5 12 4.5ZM12 17C9.24 17 7 14.76 7 12C7 9.24 9.24 7 12 7C14.76 7 17 9.24 17 12C17 14.76 14.76 17 12 17ZM12 9C10.34 9 9 10.34 9 12C9 13.66 10.34 15 12 15C13.66 15 15 13.66 15 12C15 10.34 13.66 9 12 9Z" fill="white"/>
          </svg>
          View Data Mapping
        </button>
      </div>
    </div>
  );
};

const SchemaNode = ({ data }) => {
  // Use teal color for schemas under Dbtez
  const isDbtezSchema = data.dbId === 'db4';
  const isSalesSchema = data.label === 'Sales';
  
  // Use dark theme colors
  const bgColor = '#1a1a1a';
  const borderColor = isSalesSchema ? '#fa8c16' : (isDbtezSchema ? '#00a99d' : '#52c41a');
  const handleColor = isSalesSchema ? '#fa8c16' : (isDbtezSchema ? '#00a99d' : '#52c41a');
  
  // Schema SVG icon
  const SchemaIcon = () => (
    <svg width="30" height="28" viewBox="0 0 42 39" fill="none" xmlns="http://www.w3.org/2000/svg">
      <rect x="0.6875" y="0.136963" width="40.95" height="38.4259" rx="3.6" fill="url(#schema_paint0_linear)"/>
      <path fillRule="evenodd" clipRule="evenodd" d="M35.28 25.3874H32.8947C32.7256 23.869 32.0612 22.4376 30.9499 21.2794C29.4071 19.6721 27.226 18.7868 24.8092 18.7868H21.7639V11.5874H27.2722C27.572 11.5874 27.8149 11.3441 27.8149 11.0447V6.95504C27.8149 6.65523 27.572 6.41235 27.2722 6.41235H15.2479C14.9478 6.41235 14.7049 6.65523 14.7049 6.95504V11.0447C14.7049 11.3441 14.9478 11.5874 15.2479 11.5874H20.7469V18.7868H17.6905C13.3466 18.7868 10.1185 21.5175 9.63615 25.4032C7.87769 25.5619 6.49976 27.0378 6.49976 28.8374C6.49976 30.7424 8.04467 32.2874 9.94976 32.2874C11.8552 32.2874 13.3998 30.7424 13.3998 28.8374C13.3998 27.1713 12.2188 25.7817 10.6487 25.4584C11.0914 22.0881 13.8538 19.8039 17.6905 19.8039H20.7469V25.3874H17.2251C17.0347 25.3874 16.9226 25.6088 17.0309 25.771L21.0657 32.1814C21.1595 32.3225 21.3603 32.3225 21.4545 32.1814L25.4892 25.771C25.5976 25.6088 25.4851 25.3874 25.295 25.3874H21.7639V19.8039H24.8092C26.9469 19.8039 28.8671 20.5781 30.2157 21.9836C31.1428 22.9492 31.7079 24.1322 31.8728 25.3874H29.4695C29.1687 25.3874 28.9248 25.6313 28.9248 25.9321V31.7426C28.9248 32.0434 29.1687 32.2874 29.4695 32.2874H35.28C35.5808 32.2874 35.8248 32.0434 35.8248 31.7426V25.9321C35.8248 25.6313 35.5808 25.3874 35.28 25.3874Z" fill="white"/>
      <defs>
        <linearGradient id="schema_paint0_linear" x1="0.6875" y1="19.3499" x2="41.6375" y2="19.3499" gradientUnits="userSpaceOnUse">
          <stop stopColor="#FF9D2C"/>
          <stop offset="1" stopColor="#CD750F"/>
        </linearGradient>
      </defs>
    </svg>
  );
  
  return (
    <div className="schema-node" style={{ 
      padding: '10px', 
      borderRadius: '5px', 
      background: bgColor, 
      border: `1px solid ${borderColor}`,
      width: '160px',
      position: 'relative',
      boxShadow: isDbtezSchema ? '0 0 10px rgba(0, 169, 157, 0.3)' : 'none',
      color: '#ffffff'
    }}>
      {/* Connection handles */}
      <Handle 
        type="source" 
        position={Position.Right} 
        id="right" 
        style={{ background: handleColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      <Handle 
        type="target" 
        position={Position.Left} 
        id="left" 
        style={{ background: handleColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      <Handle 
        type="source" 
        position={Position.Bottom} 
        id="bottom" 
        style={{ background: handleColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      <Handle 
        type="target" 
        position={Position.Top} 
        id="top" 
        style={{ background: handleColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      
      <div style={{ 
        fontWeight: 'bold', 
        marginBottom: '5px', 
        display: 'flex', 
        alignItems: 'center'
      }}>
        <span style={{ marginRight: '8px' }}>
          <SchemaIcon />
        </span>
        {data.label}
      </div>
      
      <div style={{ fontSize: '0.8em', color: '#aaa' }}>
        {isDbtezSchema ? 'Schema' : `${data.tables} tables`}
      </div>
    </div>
  );
};

const TableNode = ({ data }) => {
  const tableType = data.type || (data.isFact ? 'fact' : 'dimension');
  
  // Determine colors based on table type
  let borderColor, background, labelBg, labelColor;
  
  switch(tableType) {
    case 'fact':
      borderColor = '#fa8c16'; // Orange for fact tables
      background = '#1a1a1a';
      labelBg = '#fa8c16';
      labelColor = 'white';
      break;
    case 'stage':
      borderColor = '#1890ff'; // Blue for stage tables
      background = '#1a1a1a';
      labelBg = '#1890ff';
      labelColor = 'white';
      break;
    case 'dimension':
    default:
      borderColor = '#52c41a'; // Green for dimension tables
      background = '#1a1a1a';
      labelBg = '#52c41a';
      labelColor = 'white';
  }
  
  // Table SVG icon
  const TableIcon = () => (
    <svg width="30" height="28" viewBox="0 0 42 39" fill="none" xmlns="http://www.w3.org/2000/svg">
      <rect x="0.825195" width="40.5" height="38.7" rx="3.6" fill="url(#table_paint0_linear)"/>
      <path d="M17.1504 20.7083H24.9896V25.2531H17.1504V20.7083ZM17.1504 31.0499H24.9896V26.5051H17.1504V31.0499ZM17.1504 19.4563H24.9896V14.9053H17.1504V19.4563ZM26.0349 19.4563H33.8741V14.9053H26.0349V19.4563ZM8.26592 19.4563H16.1052V14.9053H8.26592V19.4563ZM8.26592 25.2531H16.1052V20.7083H8.26592V25.2531ZM26.061 25.2531H33.9002V20.7083H26.061V25.2531ZM32.3324 7.6499H9.81809C9.40227 7.6499 9.00348 7.84776 8.70946 8.19996C8.41543 8.55215 8.25024 9.02983 8.25024 9.52791V13.6533H16.5494C16.5754 13.6497 16.6017 13.6497 16.6278 13.6533H16.7062H25.4443H25.5227C25.5488 13.6497 25.5751 13.6497 25.6011 13.6533H33.9002V9.52791C33.9002 9.02983 33.7351 8.55215 33.441 8.19996C33.147 7.84776 32.7482 7.6499 32.3324 7.6499ZM8.25024 29.1719C8.25024 29.67 8.41543 30.1477 8.70946 30.4998C8.85504 30.6742 9.02788 30.8126 9.2181 30.9069C9.40832 31.0013 9.6122 31.0499 9.81809 31.0499H16.0895V26.5051H8.25024V29.1719ZM26.0453 31.0499H32.3167C32.7325 31.0499 33.1313 30.852 33.4254 30.4998C33.7194 30.1477 33.8846 29.67 33.8846 29.1719V26.5051H26.0453V31.0499Z" fill="white"/>
      <defs>
        <linearGradient id="table_paint0_linear" x1="41.3252" y1="19.35" x2="0.825195" y2="19.35" gradientUnits="userSpaceOnUse">
          <stop stopColor="#659667"/>
          <stop offset="1" stopColor="#81C784"/>
        </linearGradient>
      </defs>
    </svg>
  );
  
  return (
    <div className="table-node" style={{ 
      padding: '10px', 
      borderRadius: '5px', 
      background: background, 
      border: `1px solid ${borderColor}`,
      width: '150px',
      position: 'relative',
      boxShadow: `0 0 10px ${borderColor}33`,
      color: '#ffffff'
    }}>
      {/* Connection handles */}
      <Handle 
        type="source" 
        position={Position.Right} 
        id="right" 
        style={{ background: borderColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      <Handle 
        type="target" 
        position={Position.Left} 
        id="left" 
        style={{ background: borderColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      <Handle 
        type="source" 
        position={Position.Bottom} 
        id="bottom" 
        style={{ background: borderColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      <Handle 
        type="target" 
        position={Position.Top} 
        id="top" 
        style={{ background: borderColor, width: '10px', height: '10px' }} 
        isConnectable={true}
      />
      
      <div style={{ fontWeight: 'bold', marginBottom: '5px', display: 'flex', alignItems: 'center' }}>
        <div style={{ display: 'flex', alignItems: 'center', flex: 1 }}>
          <span style={{ marginRight: '8px' }}>
            <TableIcon />
          </span>
          <span style={{ maxWidth: '90px', overflow: 'hidden', textOverflow: 'ellipsis', whiteSpace: 'nowrap' }}>
            {data.label}
          </span>
          <span style={{ 
            fontSize: '10px', 
            background: labelBg, 
            color: labelColor, 
            padding: '1px 4px', 
            borderRadius: '3px', 
            marginLeft: '5px',
            verticalAlign: 'middle'
          }}>
            {tableType.toUpperCase()}
          </span>
        </div>
      </div>
      <div style={{ fontSize: '0.8em', color: '#aaa' }}>
        {data.columns && data.columns.length > 0 ? 
          `${data.columns.length} columns` : 
          tableType === 'fact' ? 'Fact Table' : 
          tableType === 'stage' ? 'Stage Table' : 
          'Dimension Table'}
      </div>
    </div>
  );
};

// Generate data for InfiniteCanvas using mockApiData from DataflowCanvas
const generateMockData = () => {
  // Get unique database slugs from schemas
  const uniqueDatabases = [...new Set(mockApiData.schemas.map(schema => schema.database))];
  
  // Create databases from unique database slugs
  const databases = uniqueDatabases.map((dbSlug, index) => {
    // Find all schemas for this database
    const dbSchemas = mockApiData.schemas.filter(schema => schema.database === dbSlug);
    
    // Find all tables for this database
    const dbTables = mockApiData.tables.filter(table => table.database === dbSlug);
    
    // Get the database name from the first schema (if available)
    const dbName = dbSchemas.length > 0 && dbSchemas[0].databaseName 
      ? dbSchemas[0].databaseName 
      : `Database ${dbSlug}`;
    
    return { 
      id: `db-${dbSlug}`, 
      name: dbName, // This will be displayed in the UI
      slug: dbSlug,
      schemas: dbSchemas.length, 
      tables: dbTables.length 
    };
  });
  
  // If no databases were found, add a default one
  if (databases.length === 0) {
    databases.push({ 
      id: 'db-default', 
      name: 'Default Database', 
      slug: 'default',
      schemas: mockApiData.schemas.length, 
      tables: mockApiData.tables.length 
    });
  }

  // Create schemas from mockApiData
  const schemas = mockApiData.schemas.map((schema, index) => {
    // Find the database for this schema
    const dbSlug = schema.database || 'default';
    const db = databases.find(db => db.slug === dbSlug) || databases[0];
    
    return {
      id: `schema-${schema.slug}`,
      dbId: db.id,
      name: schema.name,
      slug: schema.slug,
      tables: mockApiData.tables.filter(table => table.schema === schema.slug).length
    };
  });

  // Create tables from mockApiData
  const tables = mockApiData.tables.map((table, index) => {
    // Find the schema this table belongs to
    const schemaSlug = table.schema;
    const schemaObj = schemas.find(s => s.slug === schemaSlug);
    
    if (!schemaObj) {
      console.warn(`Schema not found for table ${table.name} (${table.slug})`);
      return null;
    }
    
    return {
      id: `table-${table.slug}`,
      schemaId: schemaObj.id,
      name: table.name,
      slug: table.slug,
      isFact: table.type === 'fact',
      type: table.type,
      columns: table.columns,
      database: table.database
    };
  }).filter(Boolean); // Remove any null entries

  return { databases, schemas, tables };
};

const InfiniteCanvas = () => {
  const reactFlowWrapper = useRef(null);
  const [reactFlowInstance, setReactFlowInstance] = useState(null);
  const { fitView, setViewport, getViewport } = useReactFlow();
  
  // State to track which elements are expanded
  const [expandedDatabases, setExpandedDatabases] = useState({});
  const [expandedSchemas, setExpandedSchemas] = useState({});
  const [selectedTable, setSelectedTable] = useState(null);
  
  // State for infinite canvas
  const [scale, setScale] = useState(1);
  const [position, setPosition] = useState({ x: 0, y: 0 });
  
  // State for connection mode
  const [isConnectionMode, setIsConnectionMode] = useState(false);
  const [connectionSource, setConnectionSource] = useState(null);
  const [connectionType, setConnectionType] = useState('default'); // default, reference, dependency
  
  // State for entity creation modals
  const [showAddEntityModal, setShowAddEntityModal] = useState(false);
  const [activeEntityType, setActiveEntityType] = useState('database'); // 'database', 'schema', or 'table'
  const [selectedDatabaseForSchema, setSelectedDatabaseForSchema] = useState(null);
  const [selectedSchemaForTable, setSelectedSchemaForTable] = useState(null);
  
  // Initialize with mock data
  const mockData = generateMockData();
  
  // State for storing created entities
  const [databases, setDatabases] = useState(() => {
    // Try to load databases from localStorage
    try {
      const savedDatabases = localStorage.getItem('databases');
      if (savedDatabases) {
        return JSON.parse(savedDatabases);
      }
    } catch (error) {
      console.error('Error loading databases from localStorage:', error);
    }
    
    // If no saved databases, use the databases from our API
    return mockData.databases.map(db => ({
      ...db,
      description: db.description || `Database ${db.name}`,
      type: db.type || 'PostgreSQL'
    }));
  });
  
  // Track if we've already fetched the database data
  const [hasFetchedDatabases, setHasFetchedDatabases] = useState(false);
  
  // Fetch real database data from API - only once
  useEffect(() => {
    // Skip if we've already fetched the data
    if (hasFetchedDatabases) {
      console.log('Skipping database fetch - already loaded');
      return;
    }
    
    const fetchDatabasesFromAPI = async () => {
      try {
        // API configuration
        const API_BASE_URL = 'https://sandbox.kezel.io/api';
        const token = "abdhsg"; // Replace with your actual token
        const orgSlug = "sN05Pjv11qvH"; // Replace with your actual org slug
        
        console.log('Fetching databases from API...');
        
        const response = await axios.post(
          `${API_BASE_URL}/qbt_database_list_get`,
          {
            token: token,
            org: orgSlug,
          },
          {
            headers: {
              'Content-Type': 'application/json'
            }
          }
        );
        
        console.log('API Response:', response.data);
        
        if (response.data && response.data.items && Array.isArray(response.data.items)) {
          // Map API response to our database format
          const apiDatabases = response.data.items.map((db, index) => ({
            id: `db-${db.con || index}`,
            name: db.name, // Use the name from API
            slug: db.con,
            description: db.description || `Database ${db.name}`,
            type: db.database_type || 'PostgreSQL',
            schemas: 0, // We'll update these later if needed
            tables: 0
          }));
          
          console.log('Mapped databases from API:', apiDatabases);
          
          // Update state with API data
          setDatabases(apiDatabases);
          
          // Mark that we've fetched the data
          setHasFetchedDatabases(true);
        }
      } catch (error) {
        console.error('Error fetching databases from API:', error);
      }
    };
    
    // Call the function to fetch databases
    fetchDatabasesFromAPI();
  }, [hasFetchedDatabases]); // Only re-run if hasFetchedDatabases changes
  
  const [schemas, setSchemas] = useState(() => {
    // Try to load schemas from localStorage
    try {
      const savedSchemas = localStorage.getItem('schemas');
      if (savedSchemas) {
        return JSON.parse(savedSchemas);
      }
    } catch (error) {
      console.error('Error loading schemas from localStorage:', error);
    }
    return [];
  });
  
  const [tables, setTables] = useState(() => {
    // Try to load tables from localStorage
    try {
      const savedTables = localStorage.getItem('tables');
      if (savedTables) {
        return JSON.parse(savedTables);
      }
    } catch (error) {
      console.error('Error loading tables from localStorage:', error);
    }
    return [];
  });
  
  // Save databases, schemas, and tables to localStorage whenever they change
  useEffect(() => {
    localStorage.setItem('databases', JSON.stringify(databases));
  }, [databases]);
  
  useEffect(() => {
    localStorage.setItem('schemas', JSON.stringify(schemas));
  }, [schemas]);
  
  useEffect(() => {
    localStorage.setItem('tables', JSON.stringify(tables));
  }, [tables]);
  
  // Handler for creating a new database with nested schemas and tables
  const handleCreateDatabase = (formData) => {
    // Create the database
    const dbId = `db-${Date.now()}`;
    const newDatabase = {
      id: dbId,
      name: formData.name,
      description: formData.description,
      url: formData.url,
      key: formData.key,
      type: formData.type,
      schemas: formData.schemas.length,
      tables: formData.schemas.reduce((count, schema) => count + schema.tables.length, 0)
    };
    
    setDatabases(prev => [...prev, newDatabase]);
    
    // Add the new database node to the canvas
    const dbNode = {
      id: dbId,
      type: 'database',
      data: { 
        label: newDatabase.name,
        name: newDatabase.name, // Include the name explicitly
        schemas: newDatabase.schemas,
        tables: newDatabase.tables,
        expanded: false,
        onToggle: toggleDatabaseExpansion,
        onViewDetails: handleViewDataFlow // Add the function to handle redirection
      },
      position: { 
        x: Math.random() * 300, 
        y: Math.random() * 300 
      },
    };
    
    const newNodes = [dbNode];
    const newEdges = [];
    
    // Process schemas
    const createdSchemas = formData.schemas.map((schemaData, schemaIndex) => {
      const schemaId = `schema-${Date.now()}-${schemaIndex}`;
      const newSchema = {
        id: schemaId,
        dbId: dbId,
        name: schemaData.name,
        description: schemaData.description,
        slug: schemaData.name.toLowerCase().replace(/\s+/g, '_'),
        tables: schemaData.tables.length
      };
      
      // Create schema node (positioned relative to database)
      const schemaNode = {
        id: schemaId,
        type: 'schema',
        data: { 
          label: schemaData.name, 
          tables: schemaData.tables.length,
          dbId: dbId,
          expanded: false,
          onToggle: toggleSchemaExpansion
        },
        position: { 
          x: dbNode.position.x + 300, 
          y: dbNode.position.y + (schemaIndex * 150)
        },
      };
      
      newNodes.push(schemaNode);
      
      // Create edge from database to schema
      newEdges.push({
        id: `e-${dbId}-${schemaId}`,
        source: dbId,
        target: schemaId,
        type: 'custom',
        animated: true,
        style: { stroke: '#00a99d', strokeWidth: 2 },
        markerEnd: {
          type: 'arrowclosed',
          width: 20,
          height: 20,
          color: '#00a99d',
        },
        data: {
          label: 'contains'
        }
      });
      
      // Process tables for this schema
      const schemaTables = schemaData.tables.map((tableData, tableIndex) => {
        const tableId = `table-${Date.now()}-${schemaIndex}-${tableIndex}`;
        const newTable = {
          id: tableId,
          schemaId: schemaId,
          name: tableData.name,
          description: tableData.description,
          slug: tableData.name.toLowerCase().replace(/\s+/g, '_'),
          type: tableData.type,
          isFact: tableData.type === 'fact',
          columns: tableData.columns || []
        };
        
        // Create table node (positioned relative to schema)
        const tableNode = {
          id: tableId,
          type: 'table',
          data: { 
            label: tableData.name, 
            type: tableData.type,
            isFact: tableData.type === 'fact',
            schemaId: schemaId,
            columns: tableData.columns || []
          },
          position: { 
            x: schemaNode.position.x + 300, 
            y: schemaNode.position.y + (tableIndex * 100)
          },
        };
        
        newNodes.push(tableNode);
        
        // Create edge from schema to table
        newEdges.push({
          id: `e-${schemaId}-${tableId}`,
          source: schemaId,
          target: tableId,
          type: 'custom',
          animated: true,
          style: { 
            stroke: tableData.type === 'fact' ? '#fa8c16' : 
                   tableData.type === 'stage' ? '#1890ff' : '#52c41a', 
            strokeWidth: 2 
          },
          markerEnd: {
            type: 'arrowclosed',
            width: 20,
            height: 20,
            color: tableData.type === 'fact' ? '#fa8c16' : 
                  tableData.type === 'stage' ? '#1890ff' : '#52c41a',
          },
          data: {
            label: 'contains'
          }
        });
        
        return newTable;
      });
      
      // Add tables to state
      setTables(prev => [...prev, ...schemaTables]);
      
      return newSchema;
    });
    
    // Add schemas to state
    setSchemas(prev => [...prev, ...createdSchemas]);
    
    // Add all nodes and edges to the canvas
    setNodes(nodes => [...nodes, ...newNodes]);
    setEdges(edges => [...edges, ...newEdges]);
    
    // Auto-expand the database if it has schemas
    if (formData.schemas.length > 0) {
      setTimeout(() => {
        toggleDatabaseExpansion(dbId);
        
        // Auto-expand schemas if they have tables
        formData.schemas.forEach((schema, index) => {
          if (schema.tables.length > 0) {
            const schemaId = `schema-${Date.now()}-${index}`;
            toggleSchemaExpansion(schemaId);
          }
        });
      }, 500);
    }
    
    setShowAddEntityModal(false);
    
    // Fit view to show all the new nodes
    setTimeout(() => {
      fitView();
    }, 1000);
  };
  
  // Handler for creating a new schema
  const handleCreateSchema = (formData) => {
    const parentDb = formData.parentDatabase;
    
    const newSchema = {
      id: `schema-${Date.now()}`,
      dbId: parentDb.id,
      name: formData.name,
      description: formData.description,
      slug: formData.name.toLowerCase().replace(/\s+/g, '_'),
      tables: 0
    };
    
    setSchemas(prev => [...prev, newSchema]);
    
    // Update the parent database's schema count
    setDatabases(prev => prev.map(db => 
      db.id === parentDb.id 
        ? { ...db, schemas: db.schemas + 1 } 
        : db
    ));
    
    // If the parent database is expanded, add the schema node to the canvas
    if (expandedDatabases[parentDb.id]) {
      const dbNode = nodes.find(n => n.id === parentDb.id);
      
      const newNode = {
        id: newSchema.id,
        type: 'schema',
        data: { 
          label: newSchema.name, 
          tables: newSchema.tables,
          dbId: parentDb.id,
          expanded: false,
          onToggle: toggleSchemaExpansion
        },
        position: { 
          x: dbNode.position.x + 150, 
          y: dbNode.position.y + 150 + (schemas.filter(s => s.dbId === parentDb.id).length * 100)
        },
      };
      
      setNodes(nodes => [...nodes, newNode]);
      
      // Add an edge from the database to the schema
      const newEdge = {
        id: `e-${parentDb.id}-${newSchema.id}`,
        source: parentDb.id,
        target: newSchema.id,
        type: 'custom',
        animated: true,
        style: { stroke: '#00a99d', strokeWidth: 2 },
        markerEnd: {
          type: 'arrowclosed',
          width: 20,
          height: 20,
          color: '#00a99d',
        },
        data: {
          label: 'contains'
        }
      };
      
      setEdges(edges => [...edges, newEdge]);
    }
    
    setShowAddEntityModal(false);
    setSelectedDatabaseForSchema(null);
    
    // Fit view to show the new node
    setTimeout(() => {
      fitView();
    }, 10);
  };
  
  // Handler for creating a new table
  const handleCreateTable = (formData) => {
    const parentSchema = formData.parentSchema;
    
    const newTable = {
      id: `table-${Date.now()}`,
      schemaId: parentSchema.id,
      name: formData.name,
      description: formData.description,
      slug: formData.name.toLowerCase().replace(/\s+/g, '_'),
      type: formData.type,
      isFact: formData.type === 'fact',
      columns: []
    };
    
    setTables(prev => [...prev, newTable]);
    
    // Update the parent schema's table count
    setSchemas(prev => prev.map(schema => 
      schema.id === parentSchema.id 
        ? { ...schema, tables: schema.tables + 1 } 
        : schema
    ));
    
    // Update the parent database's table count
    const parentDb = databases.find(db => db.id === parentSchema.dbId);
    if (parentDb) {
      setDatabases(prev => prev.map(db => 
        db.id === parentDb.id 
          ? { ...db, tables: db.tables + 1 } 
          : db
      ));
    }
    
    // If the parent schema is expanded, add the table node to the canvas
    if (expandedSchemas[parentSchema.id]) {
      const schemaNode = nodes.find(n => n.id === parentSchema.id);
      
      const newNode = {
        id: newTable.id,
        type: 'table',
        data: { 
          label: newTable.name, 
          type: newTable.type,
          isFact: newTable.isFact,
          schemaId: parentSchema.id,
          columns: newTable.columns
        },
        position: { 
          x: schemaNode.position.x + 150, 
          y: schemaNode.position.y + 150 + (tables.filter(t => t.schemaId === parentSchema.id).length * 100)
        },
      };
      
      setNodes(nodes => [...nodes, newNode]);
      
      // Add an edge from the schema to the table
      const newEdge = {
        id: `e-${parentSchema.id}-${newTable.id}`,
        source: parentSchema.id,
        target: newTable.id,
        type: 'custom',
        animated: true,
        style: { 
          stroke: newTable.isFact ? '#fa8c16' : '#52c41a', 
          strokeWidth: 2 
        },
        markerEnd: {
          type: 'arrowclosed',
          width: 20,
          height: 20,
          color: newTable.isFact ? '#fa8c16' : '#52c41a',
        },
        data: {
          label: 'contains'
        }
      };
      
      setEdges(edges => [...edges, newEdge]);
    }
    
    setShowAddEntityModal(false);
    setSelectedSchemaForTable(null);
    
    // Fit view to show the new node
    setTimeout(() => {
      fitView();
    }, 10);
  };
  
  // Define node types (memoized to prevent unnecessary re-renders)
  const nodeTypes = useMemo(() => ({
    service: ServiceNode,
    database: DatabaseNode,
    schema: SchemaNode,
    table: TableNode,
  }), []);
  
  // Define edge types
  const edgeTypes = useMemo(() => ({
    custom: CustomEdge,
    hierarchical: HierarchicalEdge,
  }), []);

    // Function to handle redirection to DataFlow view
  const handleViewDataFlow = (dbId, dbName) => {
    // Safety check for undefined parameters
    if (!dbId || !dbName) {
      console.error('Invalid database ID or name provided to handleViewDataFlow');
      return;
    }
    
    // Find the database object to get the slug
    const database = databases.find(db => db.id === dbId);
    if (!database) {
      console.error(`Database with ID ${dbId} not found`);
      return;
    }
    
    // Get the database slug, with a fallback
    const dbSlug = database.slug || (dbId.startsWith('db-') ? dbId.substring(3) : 'default');
    
    console.log(`Viewing data flow for database: ${dbName} (${dbSlug})`);
    
    // Selectively clear cached data for the current database only
    try {
      // Only clear the cache for the current database
      if (dbSlug === 'my_dwh2') {
        console.log('Setting empty schemas for service 2 database');
        // For service 2, we want to ensure it always shows as empty
        localStorage.setItem(`schemas_${dbSlug}`, JSON.stringify([]));
        localStorage.setItem(`tables_${dbSlug}`, JSON.stringify([]));
      } else {
        // For other databases like MyDataWarehouseDB, we want to preserve their data
        // but ensure we're not using stale data
        console.log(`Preserving schemas for database: ${dbName}`);
        
        // We don't need to clear anything for MyDataWarehouseDB
        // This ensures its schemas will be loaded from the API or mock data
      }
    } catch (error) {
      console.error('Error managing cached data in localStorage:', error);
    }
    
    // Store the selected database info in localStorage for the DataFlow component to use
    localStorage.setItem('selectedDatabase', JSON.stringify({
      id: dbId,
      name: dbName,
      slug: dbSlug, // Include the slug for API calls
      isEmpty: dbSlug === 'my_dwh2', // Flag to indicate if this database has no schemas
      timestamp: Date.now() // Add timestamp to ensure uniqueness
    }));
    
    // Set the current database slug in mockData.js
    if (typeof window.setCurrentDbSlug === 'function') {
      // First clear any existing data
      window.setCurrentDbSlug(null);
      
      // Then set the new database slug
      setTimeout(() => {
        window.setCurrentDbSlug(dbSlug);
      }, 50);
    } else {
      console.warn('window.setCurrentDbSlug function is not available');
    }
    
    // Check if this database should have empty schemas
    const isEmpty = dbSlug === 'my_dwh2';
    
    // Trigger an event that App.jsx can listen to
    const event = new CustomEvent('viewDataFlow', { 
      detail: { 
        databaseId: dbId, 
        databaseName: dbName,
        databaseSlug: dbSlug,
        isEmpty: isEmpty, // Flag to indicate if this database has no schemas
        timestamp: Date.now() // Add timestamp to ensure uniqueness
      }
    });
    window.dispatchEvent(event);
    
    // No alert needed for MVP - the view will change automatically
  };

  // Create a service node for "Plan Plus"
  const serviceNode = {
    id: 'service-plan-plus',
    type: 'service',
    data: {
      id: 'service-plan-plus',
      // label: 'Qubit Service: Plan+',
      databases: databases.length,
      expanded: true,
      onToggle: (id) => {
        // Toggle expansion of the service node
        setNodes(nodes => nodes.map(node => {
          if (node.id === id) {
            return { ...node, data: { ...node.data, expanded: !node.data.expanded } };
          }
          return node;
        }));
      }
    },
    position: { x: 400, y: 70 } // Positioned even higher on the canvas for maximum spacing
  };

  // Initialize with database nodes as children of the service
  const databaseNodes = databases
    .filter(db => db && db.id) // Filter out any invalid database objects
    .map((db, index) => {
      // Create a safe ID for the database
      const safeId = db.id;
      
      return {
        id: safeId,
        type: 'database',
        data: { 
          id: safeId, // Explicitly include id in data to ensure it's available
          label: db.name || `Database ${index + 1}`,
          name: db.name, // Include the name from the API response
          schemas: db.schemas || 0,
          tables: db.tables || 0,
          expanded: false,
          onToggle: (id) => toggleDatabaseExpansion(id),
          onViewDetails: handleViewDataFlow // Add the function to handle redirection
        },
        // Position databases in a grid layout with significantly increased vertical spacing
        position: { 
          x: 250 + (index % 2) * 400, // 2 columns
          y: 400 + Math.floor(index / 2) * 250 // further increased vertical spacing (400px) from service node
        },
      };
    });
    
  // Combine the service node with database nodes
  const initialNodes = [
    serviceNode, 
    ...databaseNodes
  ];

  // Create edges from service node to each database with custom hierarchical connections
  const initialEdges = databases
    .filter(db => db && db.id)
    .map((db) => ({
      id: `edge-service-to-${db.id}`,
      source: 'service-plan-plus',
      target: db.id,
      type: 'hierarchical', // Using our custom hierarchical edge
      animated: true,
      style: { 
        stroke: '#00a99d', 
        strokeWidth: 2,
        strokeDasharray: '5, 5', // Create dotted/dashed line effect
      },
      // No need for handles as they're built into the edge type
    }));
  
  const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes);
  const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges);
  
  // Update nodes when databases change (e.g., after API fetch)
  useEffect(() => {
    console.log('Databases updated, updating nodes:', databases);
    
    // Create a parent service node for "Plan Plus"
    const serviceNode = {
      id: 'service-plan-plus',
      type: 'service',
      data: {
        id: 'service-plan-plus',
        // label: 'Qubit Service: Plan+',
        databases: databases.length,
        expanded: true,
        onToggle: (id) => {
          // Toggle expansion of the service node
          setNodes(nodes => nodes.map(node => {
            if (node.id === id) {
              return { ...node, data: { ...node.data, expanded: !node.data.expanded } };
            }
            return node;
          }));
        }
      },
      position: { x: 400, y: 70 } // Positioned even higher on the canvas for maximum spacing
    };
    
    // Map database nodes as children of the service node
    const databaseNodes = databases
      .filter(db => db && db.id)
      .map((db, index) => {
        const existingNode = nodes.find(node => node.id === db.id);
        
        // If node already exists, update its data
        if (existingNode) {
          return {
            ...existingNode,
            data: {
              ...existingNode.data,
              label: db.name || `Database ${index + 1}`,
              name: db.name, // Make sure name is set from the API response
              schemas: db.schemas || 0,
              tables: db.tables || 0,
            },
            // Update position with significantly increased vertical spacing
            position: {
              x: 250 + (index % 2) * 400, // 2 columns
              y: 400 + Math.floor(index / 2) * 250 // further increased vertical spacing (400px) from service node
            }
          };
        }
        
        // Otherwise create a new node
        return {
          id: db.id,
          type: 'database',
          data: { 
            id: db.id,
            label: db.name || `Database ${index + 1}`,
            name: db.name, // Include the name from the API response
            schemas: db.schemas || 0,
            tables: db.tables || 0,
            expanded: false,
            onToggle: (id) => toggleDatabaseExpansion(id),
            onViewDetails: handleViewDataFlow
          },
          // Position databases in a grid layout with significantly increased vertical spacing
          position: { 
            x: 250 + (index % 2) * 400,
            y: 400 + Math.floor(index / 2) * 250 // further increased vertical spacing (400px) from service node
          },
        };
      });
    
    // Create edges from service node to each database with custom hierarchical connections
    const serviceEdges = databases
      .filter(db => db && db.id)
      .map((db) => ({
        id: `edge-service-to-${db.id}`,
        source: 'service-plan-plus',
        target: db.id,
        type: 'hierarchical', // Using our custom hierarchical edge
        animated: true, 
        style: { 
          stroke: '#00a99d', 
          strokeWidth: 2,
          strokeDasharray: '5, 5', // Create dotted/dashed line effect
        },
        // No need for handles as they're built into the edge type
      }));
    
    // Combine service node with database nodes
    const updatedNodes = [
      serviceNode, 
      ...databaseNodes
    ];
    
    // Only update if we have nodes to show
    if (updatedNodes.length > 0) {
      setNodes(updatedNodes);
      setEdges(serviceEdges);
    }
  }, [databases]);

  // Track viewport changes using the onMove callback instead of event listeners
  const onMove = useCallback((event, viewport) => {
    setScale(viewport.zoom);
    setPosition({ x: viewport.x, y: viewport.y });
  }, []);

  const onConnect = useCallback(
    (params) => {
      // Create a custom edge with the selected connection type
      const edgeColor = connectionType === 'reference' ? '#00a99d' : 
                        connectionType === 'dependency' ? '#ff4d4f' : '#722ed1';
      
      const edgeLabel = connectionType === 'reference' ? 'references' : 
                        connectionType === 'dependency' ? 'depends on' : 'connects to';
      
      const newEdge = {
        ...params,
        id: `e-custom-${params.source}-${params.target}`,
        type: 'custom',
        animated: true,
        style: { 
          stroke: edgeColor,
          strokeWidth: 2
        },
        markerEnd: {
          type: 'arrowclosed',
          width: 20,
          height: 20,
          color: edgeColor,
        },
        data: {
          label: edgeLabel
        }
      };
      
      // Add the new edge to the edges array
      setEdges(eds => [...eds, newEdge]);
      
      // If in connection mode, exit it after creating a connection
      if (isConnectionMode) {
        setIsConnectionMode(false);
        setConnectionSource(null);
      }
      
      console.log(`Connection created: ${params.source} -> ${params.target} (${connectionType})`);
    },
    [setEdges, connectionType, isConnectionMode]
  );

  const onInit = (instance) => {
    setReactFlowInstance(instance);
    
    // Auto-expand the first database and its schemas on load
    setTimeout(() => {
      // Find the first database in our list
      if (databases.length > 0) {
        const firstDb = databases[0];
        if (firstDb && firstDb.id) {
          console.log(`Auto-expanding first database: ${firstDb.id}`);
          
          // Expand the first database
          toggleDatabaseExpansion(firstDb.id);
          
          // After expanding the database, also expand all its schemas
          setTimeout(() => {
            const schemas = mockData.schemas.filter(schema => schema && schema.dbId === firstDb.id);
            schemas.forEach(schema => {
              if (schema && schema.id) {
                toggleSchemaExpansion(schema.id);
              }
            });
            
            // Fit view after all expansions
            setTimeout(() => {
              try {
                fitView({ padding: 0.5, maxZoom: 0.7 });
              } catch (error) {
                console.error('Error calling fitView:', error);
              }
            }, 100);
          }, 100);
        }
      }
    }, 100);
  };

  // Handle node click for expanding/collapsing
  const onNodeClick = (event, node) => {
    // Only handle expansion/collapse if not in connection mode
    if (!isConnectionMode) {
      if (node.type === 'database') {
        toggleDatabaseExpansion(node.id);
      } else if (node.type === 'schema') {
        toggleSchemaExpansion(node.id);
      } else if (node.type === 'table') {
        showTableDetails(node.id);
      }
    }
  };
  
  // Handle connection creation
  const onConnectStart = (event, params) => {
    console.log('Connection started:', params);
    setConnectionSource(params);
  };
  
  // Handle connection completion
  const onConnectEnd = (event) => {
    console.log('Connection ended');
  };

  // Toggle database expansion to show/hide schemas
  const toggleDatabaseExpansion = (dbId) => {
    // Safety check for undefined or invalid dbId
    if (!dbId) {
      console.error('Invalid database ID provided to toggleDatabaseExpansion');
      return;
    }
    
    const isExpanded = expandedDatabases[dbId];
    
    // Update the database node to show the correct toggle state
    setNodes(nodes => nodes.map(node => {
      if (node.id === dbId) {
        return {
          ...node,
          data: {
            ...node.data,
            expanded: !isExpanded,
            onToggle: toggleDatabaseExpansion
          }
        };
      }
      return node;
    }));
    
    if (isExpanded) {
      // Collapse: remove all schemas and tables for this database
      setNodes(nodes => nodes.filter(node => {
        if (node.type === 'database') return true;
        if (node.type === 'schema' && node.data) {
          return node.data.dbId !== dbId;
        }
        if (node.type === 'table' && node.data && node.data.schemaId) {
          const schema = mockData.schemas.find(s => s.id === node.data.schemaId);
          return !schema || schema.dbId !== dbId;
        }
        return true;
      }));
      
      // Safely filter edges
      setEdges(edges => edges.filter(edge => {
        if (!edge.source || !edge.target) return true;
        const sourceStr = String(edge.source);
        const targetStr = String(edge.target);
        return !sourceStr.startsWith(dbId) && !targetStr.startsWith(dbId);
      }));
      
      // Update expanded state
      setExpandedDatabases({
        ...expandedDatabases,
        [dbId]: false
      });
      
      // Also collapse any expanded schemas
      const updatedExpandedSchemas = { ...expandedSchemas };
      mockData.schemas.forEach(schema => {
        if (schema && schema.dbId === dbId) {
          updatedExpandedSchemas[schema.id] = false;
        }
      });
      setExpandedSchemas(updatedExpandedSchemas);
      
    } else {
      // Expand: add schema nodes for this database
      const dbNode = nodes.find(n => n.id === dbId);
      
      // Safety check if database node exists
      if (!dbNode) {
        console.error(`Database node with ID ${dbId} not found`);
        return;
      }
      
      const dbSchemas = mockData.schemas.filter(schema => schema && schema.dbId === dbId);
      
      // Check if we have schemas for this database
      if (dbSchemas.length === 0) {
        console.warn(`No schemas found for database ${dbId}`);
      }
      
      const schemaNodes = dbSchemas.map((schema, index) => ({
        id: schema.id,
        type: 'schema',
        data: { 
          id: schema.id, // Include id in data
          label: schema.name || `Schema ${index + 1}`, 
          tables: schema.tables || 0,
          dbId: dbId,
          expanded: false,
          onToggle: toggleSchemaExpansion
        },
        position: { 
          x: dbNode.position.x - 150 + (index * 150), 
          y: dbNode.position.y + 150 
        },
      }));
      
      const schemaEdges = dbSchemas.map(schema => ({
        id: `e-${dbId}-${schema.id}`,
        source: dbId,
        target: schema.id,
        type: 'custom',
        animated: true,
        style: { stroke: '#00a99d', strokeWidth: 2 },
        markerEnd: {
          type: 'arrowclosed',
          width: 20,
          height: 20,
          color: '#00a99d',
        },
        data: {
          label: 'contains'
        }
      }));
      
      setNodes(nodes => [...nodes, ...schemaNodes]);
      setEdges(edges => [...edges, ...schemaEdges]);
      
      // Update expanded state
      setExpandedDatabases({
        ...expandedDatabases,
        [dbId]: true
      });
    }
    
    // Safely call fitView with a delay
    setTimeout(() => {
      try {
        fitView();
      } catch (error) {
        console.error('Error calling fitView:', error);
      }
    }, 10);
  };

  // Toggle schema expansion to show/hide tables
  const toggleSchemaExpansion = (schemaId) => {
    // Safety check for undefined or invalid schemaId
    if (!schemaId) {
      console.error('Invalid schema ID provided to toggleSchemaExpansion');
      return;
    }
    
    const isExpanded = expandedSchemas[schemaId];
    
    // Update the schema node to show the correct toggle state
    setNodes(nodes => nodes.map(node => {
      if (node.id === schemaId) {
        return {
          ...node,
          data: {
            ...node.data,
            expanded: !isExpanded,
            onToggle: toggleSchemaExpansion
          }
        };
      }
      return node;
    }));
    
    if (isExpanded) {
      // Collapse: remove all tables for this schema
      setNodes(nodes => nodes.filter(node => {
        if (node.type !== 'table') return true;
        return !node.data || node.data.schemaId !== schemaId;
      }));
      
      // Safely filter edges
      setEdges(edges => edges.filter(edge => {
        if (!edge.source || !edge.target) return true;
        const sourceStr = String(edge.source);
        const targetStr = String(edge.target);
        return !sourceStr.startsWith(schemaId) && !targetStr.startsWith(schemaId);
      }));
      
      // Update expanded state
      setExpandedSchemas({
        ...expandedSchemas,
        [schemaId]: false
      });
      
    } else {
      // Expand: add table nodes for this schema
      const schemaNode = nodes.find(n => n.id === schemaId);
      
      // Safety check if schema node exists
      if (!schemaNode) {
        console.error(`Schema node with ID ${schemaId} not found`);
        return;
      }
      
      const schemaTables = mockData.tables.filter(table => table && table.schemaId === schemaId);
      
      // Check if we have tables for this schema
      if (schemaTables.length === 0) {
        console.warn(`No tables found for schema ${schemaId}`);
      }
      
      const tableNodes = schemaTables.map((table, index) => ({
        id: table.id,
        type: 'table',
        data: { 
          id: table.id, // Include id in data
          label: table.name || `Table ${index + 1}`, 
          isFact: table.isFact || false,
          type: table.type || (table.isFact ? 'fact' : 'dimension'),
          schemaId: schemaId,
          columns: table.columns || []
        },
        position: { 
          x: schemaNode.position.x - 200 + (index * 150), 
          y: schemaNode.position.y + 150 
        },
      }));
      
      const tableEdges = schemaTables.map(table => ({
        id: `e-${schemaId}-${table.id}`,
        source: schemaId,
        target: table.id,
        type: 'custom',
        animated: true,
        style: { 
          stroke: table.isFact ? '#fa8c16' : '#52c41a', 
          strokeWidth: 2 
        },
        markerEnd: {
          type: 'arrowclosed',
          width: 20,
          height: 20,
          color: table.isFact ? '#fa8c16' : '#52c41a',
        },
        data: {
          label: table.isFact ? 'fact table' : 'dimension'
        }
      }));
      
      setNodes(nodes => [...nodes, ...tableNodes]);
      setEdges(edges => [...edges, ...tableEdges]);
      
      // Update expanded state
      setExpandedSchemas({
        ...expandedSchemas,
        [schemaId]: true
      });
    }
    
    setTimeout(() => {
      fitView();
    }, 10);
  };

  // Show table details when a table is clicked
  const showTableDetails = (tableId) => {
    // Look for the table in both the state tables and mock data
    const table = tables.find(t => t.id === tableId) || mockData.tables.find(t => t.id === tableId);
    if (!table) return;
    
    setSelectedTable(table);
    
    // Format the table type
    const tableType = table.type || (table.isFact ? 'fact' : 'dimension');
    
    // Format the columns list
    const columnsText = table.columns && table.columns.length > 0 
      ? `Columns:\n${table.columns.join('\n')}` 
      : 'No columns defined';
    
    // In a real application, you would show a modal with detailed information
    alert(`Table Details: ${table.name}\nType: ${tableType.toUpperCase()}\n\n${columnsText}\n\nIn a real application, this would show a modal with columns, relationships, and metrics.`);
  };

  // Delete a custom connection
  const deleteConnection = (edgeId) => {
    setEdges(edges => edges.filter(edge => edge.id !== edgeId));
  };
  
  // Get node label by ID
  const getNodeLabel = (nodeId) => {
    const node = nodes.find(n => n.id === nodeId);
    return node ? node.data.label : nodeId;
  };
  
  // Reset the canvas to show only databases
  const resetView = () => {
    setNodes(initialNodes);
    setEdges(initialEdges);
    setExpandedDatabases({});
    setExpandedSchemas({});
    setSelectedTable(null);
    setIsConnectionMode(false);
    setConnectionSource(null);
    
    setTimeout(() => {
      fitView();
    }, 10);
  };

  return (
    <div style={{ width: '100%', height: '80vh', background: '#121212' }} ref={reactFlowWrapper}>
      <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onConnect={onConnect}
        onConnectStart={onConnectStart}
        onConnectEnd={onConnectEnd}
        onInit={onInit}
        onNodeClick={onNodeClick}
        onMove={onMove}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        minZoom={0.05} // Allow zooming out much further
        maxZoom={8}    // Allow zooming in closer
        defaultViewport={{ x: 0, y: 0, zoom: 0.7 }}
        fitView
        fitViewOptions={{ padding: 0.5, maxZoom: 0.7 }}
        style={{ background: '#121212' }}
        className={isConnectionMode ? 'connection-mode' : ''}
        connectionMode="loose"
        connectionLineStyle={{
          stroke: connectionType === 'reference' ? '#00a99d' : 
                 connectionType === 'dependency' ? '#ff4d4f' : '#52c41a',
          strokeWidth: 3,
          strokeDasharray: '5,5'
        }}
        connectionLineType="bezier"
        proOptions={{ hideAttribution: true }}
        defaultMarkerColor="#00a99d"
      >
        <Controls 
          showInteractive={true}
          position="bottom-right"
          style={{ 
            background: '#1a1a1a', 
            border: '1px solid #333',
            borderRadius: '4px',
            boxShadow: '0 0 10px rgba(0,0,0,0.5)'
          }}
        />
        <MiniMap 
          nodeStrokeWidth={3}
          zoomable
          pannable
          style={{ 
            background: '#1a1a1a', 
            border: '1px solid #333',
            borderRadius: '4px'
          }}
          maskColor="rgba(18, 18, 18, 0.7)"
        />
        <Background 
          variant="dots" 
          gap={12} 
          size={1} 
          color="#333"
        />
        
        {/* Add Button */}
        <Panel position="top-right">
          <button
            onClick={() => setShowAddEntityModal(true)}
            style={{
              background: 'transparent',
              border: 'none',
              padding: 0,
              cursor: 'pointer'
            }}
            title="Add DataSource"
          >
            <svg width="99" height="36" viewBox="0 0 99 36" fill="none" xmlns="http://www.w3.org/2000/svg">
              <rect x="0.5" y="0.5" width="98" height="35" rx="7.5" fill="#3EA29A"/>
              <rect x="0.5" y="0.5" width="98" height="35" rx="7.5" stroke="#27B5AA"/>
              <path fillRule="evenodd" clipRule="evenodd" d="M17.3448 10.0172C20.4387 9.67462 23.561 9.67462 26.6548 10.0172C28.3678 10.2092 29.7498 11.5582 29.9508 13.2772C30.3175 16.4152 30.3175 19.5852 29.9508 22.7232C29.7498 24.4422 28.3678 25.7912 26.6548 25.9832C23.561 26.3257 20.4387 26.3257 17.3448 25.9832C15.6318 25.7912 14.2498 24.4422 14.0488 22.7232C13.6823 19.5855 13.6823 16.4158 14.0488 13.2782C14.1505 12.443 14.5312 11.6667 15.1293 11.075C15.7273 10.4833 16.5077 10.1109 17.3438 10.0182M21.9998 13.0072C22.1987 13.0072 22.3895 13.0862 22.5302 13.2268C22.6708 13.3675 22.7498 13.5583 22.7498 13.7572V17.2502H26.2428C26.4417 17.2502 26.6325 17.3292 26.7732 17.4698C26.9138 17.6105 26.9928 17.8013 26.9928 18.0002C26.9928 18.1991 26.9138 18.3898 26.7732 18.5305C26.6325 18.6712 26.4417 18.7502 26.2428 18.7502H22.7498V22.2432C22.7498 22.4421 22.6708 22.6328 22.5302 22.7735C22.3895 22.9142 22.1987 22.9932 21.9998 22.9932C21.8009 22.9932 21.6102 22.9142 21.4695 22.7735C21.3289 22.6328 21.2498 22.4421 21.2498 22.2432V18.7502H17.7568C17.5579 18.7502 17.3672 18.6712 17.2265 18.5305C17.0859 18.3898 17.0068 18.1991 17.0068 18.0002C17.0068 17.8013 17.0859 17.6105 17.2265 17.4698C17.3672 17.3292 17.5579 17.2502 17.7568 17.2502H21.2498V13.7572C21.2498 13.5583 21.3289 13.3675 21.4695 13.2268C21.6102 13.0862 21.8009 13.0072 21.9998 13.0072Z" fill="white"/>
              <path d="M38.9787 23H37.348L41.0121 12.8182H42.7869L46.451 23H44.8203L41.9418 14.6676H41.8622L38.9787 23ZM39.2521 19.0128H44.5419V20.3054H39.2521V19.0128ZM50.5376 23.1491C49.9212 23.1491 49.371 22.9917 48.8871 22.6768C48.4065 22.3587 48.0286 21.9062 47.7536 21.3196C47.4818 20.7296 47.3459 20.022 47.3459 19.1967C47.3459 18.3714 47.4834 17.6655 47.7585 17.0788C48.0369 16.4922 48.4181 16.0431 48.902 15.7315C49.3859 15.42 49.9344 15.2642 50.5476 15.2642C51.0215 15.2642 51.4027 15.3438 51.6911 15.5028C51.9827 15.6586 52.2081 15.8409 52.3672 16.0497C52.5296 16.2585 52.6555 16.4425 52.745 16.6016H52.8345V12.8182H54.321V23H52.8693V21.8118H52.745C52.6555 21.9742 52.5263 22.1598 52.3572 22.3686C52.1915 22.5774 51.9628 22.7597 51.6712 22.9155C51.3795 23.0713 51.0017 23.1491 50.5376 23.1491ZM50.8658 21.8814C51.2933 21.8814 51.6546 21.7687 51.9496 21.5433C52.2479 21.3146 52.4732 20.9981 52.6257 20.5938C52.7815 20.1894 52.8594 19.7187 52.8594 19.1818C52.8594 18.6515 52.7831 18.1875 52.6307 17.7898C52.4782 17.392 52.2545 17.0821 51.9595 16.8601C51.6645 16.638 51.3 16.527 50.8658 16.527C50.4183 16.527 50.0455 16.643 49.7472 16.875C49.4489 17.107 49.2235 17.4235 49.071 17.8246C48.9219 18.2256 48.8473 18.678 48.8473 19.1818C48.8473 19.6922 48.9235 20.1513 49.076 20.5589C49.2285 20.9666 49.4538 21.2898 49.7521 21.5284C50.0537 21.7637 50.425 21.8814 50.8658 21.8814ZM59.3013 23.1491C58.6848 23.1491 58.1346 22.9917 57.6507 22.6768C57.1702 22.3587 56.7923 21.9062 56.5172 21.3196C56.2454 20.7296 56.1096 20.022 56.1096 19.1967C56.1096 18.3714 56.2471 17.6655 56.5222 17.0788C56.8006 16.4922 57.1818 16.0431 57.6657 15.7315C58.1496 15.42 58.6981 15.2642 59.3113 15.2642C59.7852 15.2642 60.1664 15.3438 60.4547 15.5028C60.7464 15.6586 60.9718 15.8409 61.1309 16.0497C61.2933 16.2585 61.4192 16.4425 61.5087 16.6016H61.5982V12.8182H63.0847V23H61.633V21.8118H61.5087C61.4192 21.9742 61.29 22.1598 61.1209 22.3686C60.9552 22.5774 60.7265 22.7597 60.4348 22.9155C60.1432 23.0713 59.7653 23.1491 59.3013 23.1491ZM59.6294 21.8814C60.057 21.8814 60.4183 21.7687 60.7132 21.5433C61.0115 21.3146 61.2369 20.9981 61.3894 20.5938C61.5452 20.1894 61.623 19.7187 61.623 19.1818C61.623 18.6515 61.5468 18.1875 61.3944 17.7898C61.2419 17.392 61.0182 17.0821 60.7232 16.8601C60.4282 16.638 60.0636 16.527 59.6294 16.527C59.182 16.527 58.8091 16.643 58.5108 16.875C58.2125 17.107 57.9872 17.4235 57.8347 17.8246C57.6855 18.2256 57.611 18.678 57.611 19.1818C57.611 19.6922 57.6872 20.1513 57.8397 20.5589C57.9921 20.9666 58.2175 21.2898 58.5158 21.5284C58.8174 21.7637 59.1886 21.8814 59.6294 21.8814Z" fill="white"/>
              <path d="M69.3439 17.9091C69.3439 16.6629 69.508 15.5161 69.8361 14.4688C70.1642 13.4214 70.6432 12.4553 71.2729 11.5703H72.6351C72.3899 11.8984 72.1612 12.3011 71.949 12.7784C71.7369 13.2557 71.5513 13.7794 71.3922 14.3494C71.2331 14.9162 71.1088 15.5045 71.0194 16.1143C70.9299 16.7209 70.8851 17.3191 70.8851 17.9091C70.8851 18.6979 70.963 19.4967 71.1188 20.3054C71.2746 21.1141 71.485 21.8648 71.7502 22.5575C72.0153 23.2502 72.3103 23.8153 72.6351 24.2528H71.2729C70.6432 23.3679 70.1642 22.4018 69.8361 21.3544C69.508 20.3071 69.3439 19.1586 69.3439 17.9091ZM75.1818 23H73.5511L77.2152 12.8182H78.9901L82.6541 23H81.0234L78.1449 14.6676H78.0653L75.1818 23ZM75.4553 19.0128H80.745V20.3054H75.4553V19.0128ZM86.8638 17.9091C86.8638 19.1586 86.6998 20.3071 86.3716 21.3544C86.0435 22.4018 85.5646 23.3679 84.9348 24.2528H83.5726C83.8179 23.9247 84.0466 23.522 84.2587 23.0447C84.4708 22.5675 84.6564 22.0455 84.8155 21.4787C84.9746 20.9086 85.0989 20.3187 85.1884 19.7088C85.2779 19.099 85.3226 18.4991 85.3226 17.9091C85.3226 17.1236 85.2447 16.3265 85.089 15.5178C84.9332 14.709 84.7227 13.9583 84.4576 13.2656C84.1924 12.5729 83.8974 12.0078 83.5726 11.5703H84.9348C85.5646 12.4553 86.0435 13.4214 86.3716 14.4688C86.6998 15.5161 86.8638 16.6629 86.8638 17.9091Z" fill="white" fillOpacity="0.5"/>
            </svg>
          </button>
        </Panel>
        
        {/* Minimal control panel */}
        <Panel position="top-left">
          <div style={{ 
            background: 'rgba(26, 26, 26, 0.8)', 
            padding: '8px', 
            borderRadius: '5px',
            boxShadow: '0 0 10px rgba(0,0,0,0.3)',
            border: '1px solid #333',
            color: '#ffffff',
            backdropFilter: 'blur(5px)',
            maxWidth: '200px'
          }}>
            <div style={{ display: 'flex', gap: '5px', marginBottom: '5px' }}>
              <button 
                onClick={resetView}
                style={{
                  padding: '5px 8px',
                  background: 'linear-gradient(45deg, #00a99d, #1890ff)',
                  color: 'white',
                  border: 'none',
                  borderRadius: '3px',
                  cursor: 'pointer',
                  fontSize: '12px',
                  flex: 1,
                  boxShadow: '0 0 5px rgba(0, 169, 157, 0.3)'
                }}
                className="glow-effect"
              >
                Reset View
              </button>
              <button 
                onClick={() => {
                  if (reactFlowInstance) {
                    const viewport = reactFlowInstance.getViewport();
                    reactFlowInstance.setViewport({ 
                      x: viewport.x, 
                      y: viewport.y, 
                      zoom: Math.max(0.05, viewport.zoom * 0.5) 
                    });
                  }
                }}
                style={{
                  padding: '5px 8px',
                  background: 'linear-gradient(45deg, #52c41a, #00a99d)',
                  color: 'white',
                  border: 'none',
                  borderRadius: '3px',
                  cursor: 'pointer',
                  fontSize: '12px',
                  boxShadow: '0 0 5px rgba(82, 196, 26, 0.3)'
                }}
                className="glow-effect"
              >
                Zoom Out
              </button>
            </div>
            
            <div style={{ 
              fontSize: '11px',
              color: '#aaa',
              marginTop: '5px'
            }}>
              <div style={{ display: 'flex', justifyContent: 'space-between' }}>
                <span>Zoom: {Math.round(scale * 70)}%</span>
                <span>Click nodes to expand</span>
              </div>
            </div>
          </div>
        </Panel>
        
        {/* Dbtez Control Panel */}
        <Panel position="bottom-right">
          <div style={{ 
            background: 'rgba(26, 26, 26, 0.8)', 
            padding: '8px', 
            borderRadius: '5px',
            boxShadow: '0 0 10px rgba(0,0,0,0.3)',
            border: '1px solid #00a99d',
            color: '#ffffff',
            backdropFilter: 'blur(5px)',
            display: 'flex',
            flexDirection: 'column',
            gap: '5px',
            marginBottom: '60px',
            marginRight: '10px'
          }}>
            <div style={{ 
              display: 'flex', 
              alignItems: 'center', 
              justifyContent: 'space-between',
              marginBottom: '5px'
            }}>
              <span style={{ 
                display: 'flex', 
                alignItems: 'center', 
                fontSize: '12px',
                color: '#00a99d',
                fontWeight: 'bold'
              }}>
                <span style={{ 
                  display: 'inline-flex',
                  alignItems: 'center',
                  justifyContent: 'center',
                  width: '18px', 
                  height: '18px', 
                  background: 'linear-gradient(45deg, #00a99d, #52c41a)', 
                  borderRadius: '4px',
                  color: 'white',
                  marginRight: '5px',
                  fontSize: '10px'
                }}>
                  <FaDatabase />
                </span>
                Dbtez Data Warehouse
              </span>
            </div>
            
            <div style={{ display: 'flex', gap: '5px' }}>
              <button 
                onClick={() => toggleDatabaseExpansion('db4')}
                style={{
                  padding: '4px 0',
                  background: 'linear-gradient(45deg, #00a99d, #52c41a)',
                  color: 'white',
                  border: 'none',
                  borderRadius: '3px',
                  cursor: 'pointer',
                  fontSize: '14px',
                  fontWeight: 'bold',
                  width: '30px',
                  display: 'flex',
                  alignItems: 'center',
                  justifyContent: 'center',
                  boxShadow: '0 0 5px rgba(0, 169, 157, 0.3)'
                }}
                className="glow-effect"
                title="Expand/Collapse Database"
              >
                {expandedDatabases['db4'] ? '−' : '+'}
              </button>
              
              <button 
                onClick={() => {
                  if (!expandedDatabases['db4']) {
                    toggleDatabaseExpansion('db4');
                    setTimeout(() => {
                      toggleSchemaExpansion('schema10');
                    }, 100);
                  } else if (!expandedSchemas['schema10']) {
                    toggleSchemaExpansion('schema10');
                  } else {
                    // If both are expanded, collapse schema
                    toggleSchemaExpansion('schema10');
                  }
                }}
                style={{
                  padding: '4px 0',
                  background: 'linear-gradient(45deg, #fa8c16, #faad14)',
                  color: 'white',
                  border: 'none',
                  borderRadius: '3px',
                  cursor: 'pointer',
                  fontSize: '14px',
                  fontWeight: 'bold',
                  width: '30px',
                  display: 'flex',
                  alignItems: 'center',
                  justifyContent: 'center',
                  boxShadow: '0 0 5px rgba(250, 140, 22, 0.3)'
                }}
                className="glow-effect"
                title="Expand/Collapse Sales Schema"
              >
                {expandedSchemas['schema10'] ? '−' : '+'}
              </button>
            </div>
          </div>
        </Panel>
        
        {/* Connection Mode Indicator */}
        {isConnectionMode && (
          <Panel position="top-center">
            <div style={{
              background: '#fffbe6',
              padding: '8px 15px',
              borderRadius: '20px',
              boxShadow: '0 2px 8px rgba(0,0,0,0.15)',
              display: 'flex',
              alignItems: 'center',
              gap: '8px'
            }}>
              <span style={{ 
                width: '10px', 
                height: '10px', 
                borderRadius: '50%', 
                background: connectionType === 'reference' ? '#1890ff' : 
                           connectionType === 'dependency' ? '#ff4d4f' : '#722ed1',
                display: 'inline-block',
                animation: 'pulse 1.5s infinite'
              }}></span>
              <span>
                <strong>Connection Mode Active</strong> - 
                Drag from a connection point (dot) to another node
              </span>
            </div>
          </Panel>
        )}
        
        {/* Connections Panel */}
        <Panel position="bottom-left">
          {edges.some(edge => edge.id.startsWith('e-custom')) && (
            <div style={{
              background: 'white',
              padding: '10px',
              borderRadius: '5px',
              boxShadow: '0 0 10px rgba(0,0,0,0.1)',
              maxWidth: '300px',
              maxHeight: '200px',
              overflowY: 'auto'
            }}>
              <h4 style={{ margin: '0 0 8px 0', fontSize: '14px' }}>Custom Connections</h4>
              <ul style={{ 
                listStyle: 'none', 
                padding: 0, 
                margin: 0,
                fontSize: '12px'
              }}>
                {edges
                  .filter(edge => edge.id.startsWith('e-custom'))
                  .map(edge => (
                    <li key={edge.id} style={{ 
                      padding: '5px', 
                      borderBottom: '1px solid #f0f0f0',
                      display: 'flex',
                      justifyContent: 'space-between',
                      alignItems: 'center'
                    }}>
                      <div>
                        <span style={{ 
                          display: 'inline-block', 
                          width: '8px', 
                          height: '8px', 
                          borderRadius: '50%', 
                          background: edge.style?.stroke || '#722ed1',
                          marginRight: '5px'
                        }}></span>
                        <strong>{getNodeLabel(edge.source)}</strong>
                        <span style={{ margin: '0 5px' }}>→</span>
                        <strong>{getNodeLabel(edge.target)}</strong>
                        <div style={{ 
                          fontSize: '10px', 
                          color: '#666',
                          marginTop: '2px',
                          marginLeft: '13px'
                        }}>
                          {edge.label || 'Connected to'}
                        </div>
                      </div>
                      <button
                        onClick={() => deleteConnection(edge.id)}
                        style={{
                          background: 'transparent',
                          border: 'none',
                          color: '#ff4d4f',
                          cursor: 'pointer',
                          fontSize: '12px',
                          padding: '2px 5px'
                        }}
                      >
                        ✕
                      </button>
                    </li>
                  ))}
              </ul>
            </div>
          )}
        </Panel>
      </ReactFlow>
      
      {/* Entity Creation Modal */}
      <Modal
        isOpen={showAddEntityModal}
        title="Add DataSource"
        onClose={() => {
          setShowAddEntityModal(false);
          setActiveEntityType('database');
          setSelectedDatabaseForSchema(null);
          setSelectedSchemaForTable(null);
        }}
      >
        <EntitySelector 
          activeEntityType={activeEntityType}
          setActiveEntityType={setActiveEntityType}
          databases={databases}
          schemas={schemas}
          selectedDatabaseForSchema={selectedDatabaseForSchema}
          setSelectedDatabaseForSchema={setSelectedDatabaseForSchema}
          selectedSchemaForTable={selectedSchemaForTable}
          setSelectedSchemaForTable={setSelectedSchemaForTable}
          onSaveDatabase={handleCreateDatabase}
          onSaveSchema={handleCreateSchema}
          onSaveTable={handleCreateTable}
          onCancel={() => {
            setShowAddEntityModal(false);
            setActiveEntityType('database');
            setSelectedDatabaseForSchema(null);
            setSelectedSchemaForTable(null);
          }}
        />
      </Modal>
    </div>
  );
};

export default function InfiniteCanvasWrapper() {
  return (
    <div style={{ width: '100%', height: '100%', background: '#121212' }}>
      <ReactFlowProvider>
        <InfiniteCanvas />
      </ReactFlowProvider>
    </div>
  );
}