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

// Import icons from react-icons
import { 
  FaDatabase, 
  FaTable, 
  FaPlus, 
  FaTimes, 
  FaKey, 
  FaLink,
  FaLayerGroup,
  FaColumns,
  FaProjectDiagram,
  FaChevronDown,
  FaChevronRight,
  FaServer,
  FaCloud,
  FaHdd
} from 'react-icons/fa';
import { CustomDatabaseIcon, CustomDocumentIcon, CustomDimensionIcon } from './CustomIcons';
import { Breadcrumbs, Link, Typography, Menu, MenuItem, ListItemIcon, ListItemText } from '@mui/material';
import AddTableModal from './AddTableModal';

// Custom styles for ER Diagram
const generateERStyles = () => {
  return `
    .react-flow__node {
      z-index: 1;
      margin: 20px;
    }
    
    .er-database-wrapper {
      z-index: -2;
      pointer-events: all;
      border-radius: 25px;
      cursor: grab;
      user-select: none;
      transition: all 0.2s ease;
      padding: 100px;
      border: 3px solid rgba(0, 169, 157, 0.4);
      background: rgba(0, 169, 157, 0.03);
      box-shadow: 0 0 30px rgba(0, 169, 157, 0.1);
    }
    
    .er-database-wrapper:hover {
      border-color: rgba(0, 169, 157, 0.6);
      background: rgba(0, 169, 157, 0.05);
      box-shadow: 0 0 40px rgba(0, 169, 157, 0.15);
    }
    
    .er-database-wrapper:active {
      cursor: grabbing;
    }
    
    .er-database-label {
      position: absolute;
      top: 20px;
      left: 20px;
      background: rgba(0, 169, 157, 0.95);
      color: white;
      padding: 12px 16px;
      border-radius: 8px;
      font-size: 14px;
      font-weight: bold;
      box-shadow: 0 4px 12px rgba(0, 169, 157, 0.3);
      display: flex;
      align-items: center;
      gap: 10px;
      z-index: 15;
    }
    
    .er-schema-group {
      z-index: -1;
      pointer-events: all;
      border-radius: 20px;
      cursor: grab;
      user-select: none;
      transition: all 0.2s ease;
      padding: 80px;
      border: 2px dashed rgba(138, 43, 226, 0.3);
      background: rgba(138, 43, 226, 0.05);
      position: relative;
      overflow: visible;
    }
    
    .er-schema-group:hover {
      border-color: rgba(138, 43, 226, 0.5);
      background: rgba(138, 43, 226, 0.08);
      box-shadow: 0 0 20px rgba(138, 43, 226, 0.1);
    }
    
    .er-schema-group:active {
      cursor: grabbing;
    }
    
    .er-schema-label {
      position: absolute;
      top: 15px;
      left: 15px;
      background: rgba(138, 43, 226, 0.9);
      color: white;
      padding: 8px 12px;
      border-radius: 6px;
      font-size: 12px;
      font-weight: bold;
      box-shadow: 0 2px 8px rgba(138, 43, 226, 0.3);
      display: flex;
      align-items: center;
      gap: 8px;
      z-index: 10;
    }
    
    .er-table-node {
      background: white;
      border: 2px solid #e1e5e9;
      border-radius: 10px;
      box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);
      width: 260px;
      min-height: 200px;
      transition: all 0.2s ease;
    }
    
    .er-table-node:hover {
      border-color: #8a2be2;
      box-shadow: 0 4px 16px rgba(138, 43, 226, 0.2);
    }
    
    .er-table-header {
      background: linear-gradient(135deg, #8a2be2, #9932cc);
      color: white;
      padding: 14px 16px;
      border-radius: 8px 8px 0 0;
      display: flex;
      align-items: center;
      gap: 10px;
      font-weight: bold;
      font-size: 15px;
    }
    
    .er-table-header.stage {
      background: linear-gradient(135deg, #00a99d, #52c41a);
    }
    
    .er-table-header.fact {
      background: linear-gradient(135deg, #fa8c16, #faad14);
    }
    
    .er-table-header.dimension {
      background: linear-gradient(135deg, #52c41a, #73d13d);
    }
    
    .er-column-list {
      padding: 0;
      margin: 0;
      list-style: none;
      max-height: 300px;
      overflow-y: auto;
    }
    
    .er-column-item {
      padding: 10px 16px;
      border-bottom: 1px solid #f0f0f0;
      display: flex;
      align-items: center;
      gap: 10px;
      font-size: 13px;
      transition: background-color 0.2s ease;
    }
    
    .er-column-item:hover {
      background-color: #f8f9fa;
    }
    
    .er-column-item:last-child {
      border-bottom: none;
    }
    
    .er-column-name {
      flex: 1;
      font-weight: 500;
      color: #333;
    }
    
    .er-column-type {
      color: #666;
      font-size: 11px;
      background: #f5f5f5;
      padding: 2px 6px;
      border-radius: 3px;
    }
    
    .er-primary-key {
      color: #faad14;
    }
    
    .er-foreign-key {
      color: #8a2be2;
    }
    
    .er-relationship-edge {
      stroke: #8a2be2;
      stroke-width: 2;
    }
    
    .er-relationship-edge .react-flow__edge-path {
      stroke: #8a2be2;
      stroke-width: 2;
    }
    
    .er-relationship-edge-animated {
      stroke: #8a2be2;
      stroke-width: 2;
      stroke-dasharray: 5;
      animation: dashdraw 0.5s linear infinite;
    }
    
    @keyframes dashdraw {
      to {
        stroke-dashoffset: -10;
      }
    }
    
    /* Enhanced arrow markers */
    .react-flow__edge.selected .react-flow__edge-path {
      stroke: #faad14 !important;
      stroke-width: 3 !important;
    }
    
    .react-flow__edge:hover .react-flow__edge-path {
      stroke: #52c41a !important;
      stroke-width: 3 !important;
    }
    
    /* Custom arrow styling */
    .er-custom-arrow {
      pointer-events: none;
      z-index: 10;
    }
    
    /* Cardinality indicators */
    .er-cardinality-indicator {
      background: rgba(255, 255, 255, 0.95) !important;
      border: 1px solid #8a2be2 !important;
      border-radius: 3px !important;
      padding: 2px 6px !important;
      font-size: 10px !important;
      font-weight: bold !important;
      color: #8a2be2 !important;
      box-shadow: 0 1px 3px rgba(0,0,0,0.2) !important;
    }
    
    /* Relationship direction indicators */
    .er-direction-arrow {
      fill: #8a2be2;
      stroke: #8a2be2;
      stroke-width: 1;
    }
    
    .er-direction-arrow.selected {
      fill: #faad14;
      stroke: #faad14;
    }
    
    .er-relationship-label {
      background: rgba(138, 43, 226, 0.9);
      color: white;
      padding: 4px 8px;
      border-radius: 4px;
      font-size: 10px;
      font-weight: bold;
    }
    
    .er-relationship-label:hover + .relationship-tooltip {
      opacity: 1 !important;
    }
    
    .er-add-button {
      position: fixed;
      bottom: 30px;
      right: 30px;
      width: 60px;
      height: 60px;
      border-radius: 50%;
      background: linear-gradient(135deg, #8a2be2, #9932cc);
      border: none;
      color: white;
      font-size: 24px;
      cursor: pointer;
      box-shadow: 0 4px 16px rgba(138, 43, 226, 0.3);
      transition: all 0.2s ease;
      z-index: 1000;
      display: flex;
      align-items: center;
      justify-content: center;
    }
    
    .er-add-button:hover {
      transform: scale(1.1);
      box-shadow: 0 6px 20px rgba(138, 43, 226, 0.4);
    }
    
    .er-add-button:active {
      transform: scale(0.95);
    }
    
    .er-breadcrumb-dropdown {
      display: inline-flex;
      align-items: center;
      cursor: pointer;
      padding: 4px 8px;
      border-radius: 4px;
      transition: background-color 0.2s ease;
    }
    
    .er-breadcrumb-dropdown:hover {
      background-color: rgba(0, 169, 157, 0.1);
    }
    
    .er-breadcrumb-dropdown .dropdown-arrow {
      margin-left: 4px;
      font-size: 12px;
      transition: transform 0.2s ease;
    }
    
    .er-breadcrumb-dropdown.open .dropdown-arrow {
      transform: rotate(180deg);
    }
    
    .er-service-menu {
      background: rgba(26, 26, 26, 0.95) !important;
      border: 1px solid rgba(0, 169, 157, 0.3) !important;
      border-radius: 8px !important;
      backdrop-filter: blur(10px) !important;
      box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3) !important;
    }
    
    .er-service-menu .MuiMenuItem-root {
      color: #ffffff !important;
      padding: 12px 16px !important;
      border-radius: 4px !important;
      margin: 4px 8px !important;
      transition: all 0.2s ease !important;
    }
    
    .er-service-menu .MuiMenuItem-root:hover {
      background-color: rgba(0, 169, 157, 0.1) !important;
    }
    
    .er-service-menu .MuiListItemIcon-root {
      color: #00a99d !important;
      min-width: 32px !important;
    }
    
    /* Ensure ReactFlow markers are visible */
    .react-flow__edges {
      z-index: 1;
    }
    
    .react-flow__edge {
      pointer-events: all;
    }
    
    .react-flow__edge path {
      stroke: #8a2be2;
      stroke-width: 2;
    }
    
    .react-flow__edge.selected path {
      stroke: #faad14 !important;
      stroke-width: 3 !important;
    }
    
    /* Custom arrows are always visible */
    .er-custom-arrow {
      pointer-events: none;
      z-index: 10;
    }
  `;
};

// Custom Table Node Component for ER Diagram
const ERTableNode = ({ data, id }) => {
  const getTableIcon = () => {
    switch(data.table_type) {
      case 'stage':
        return <CustomDatabaseIcon width="16" height="16" />;
      case 'fact':
        return <CustomDocumentIcon width="16" height="16" />;
      case 'dimension':
        return <CustomDimensionIcon width="16" height="16" />;
      default:
        return <FaTable />;
    }
  };

  return (
    <div className="er-table-node">
      <Handle
        type="target"
        position={Position.Left}
        style={{ background: '#8a2be2', width: 8, height: 8 }}
      />
      <Handle
        type="source"
        position={Position.Right}
        style={{ background: '#8a2be2', width: 8, height: 8 }}
      />
      
      <div className={`er-table-header ${data.table_type || 'default'}`}>
        {getTableIcon()}
        <span>{data.name}</span>
      </div>
      
      <ul className="er-column-list">
        {data.columns && data.columns.map((column, index) => (
          <li key={index} className="er-column-item">
            {column.is_primary_key && <FaKey className="er-primary-key" />}
            {column.is_foreign_key && <FaLink className="er-foreign-key" />}
            <span className="er-column-name">{column.name}</span>
            <span className="er-column-type">{column.data_type}</span>
          </li>
        ))}
      </ul>
    </div>
  );
};

// Custom Edge Component for Relationships
const ERRelationshipEdge = ({ 
  id, 
  sourceX, 
  sourceY, 
  targetX, 
  targetY, 
  sourcePosition, 
  targetPosition, 
  data,
  selected 
}) => {
  const [edgePath, labelX, labelY] = getBezierPath({
    sourceX,
    sourceY,
    sourcePosition,
    targetX,
    targetY,
    targetPosition,
  });

  // Determine relationship type and styling
  const relationshipType = data?.relationship_type || '1:N';
  const isOneToMany = relationshipType === '1:N';
  const isOneToOne = relationshipType === '1:1';
  const isManyToMany = relationshipType === 'N:M';

  // Calculate arrow position and angle
  const dx = targetX - sourceX;
  const dy = targetY - sourceY;
  const angle = Math.atan2(dy, dx) * (180 / Math.PI);
  
  // Arrow position (closer to target)
  const arrowDistance = 20;
  const arrowX = targetX - (arrowDistance * Math.cos(Math.atan2(dy, dx)));
  const arrowY = targetY - (arrowDistance * Math.sin(Math.atan2(dy, dx)));

  // Start arrow position for many-to-many
  const startArrowX = sourceX + (arrowDistance * Math.cos(Math.atan2(dy, dx)));
  const startArrowY = sourceY + (arrowDistance * Math.sin(Math.atan2(dy, dx)));

  return (
    <>
      {/* Main edge path */}
      <path
        d={edgePath}
        className={`er-relationship-edge ${selected ? 'selected' : ''}`}
        style={{
          stroke: selected ? '#faad14' : '#8a2be2',
          strokeWidth: selected ? 3 : 2,
          fill: 'none',
        }}
      />
      
      {/* Custom Arrow at target */}
      <polygon
        points="0,-4 8,0 0,4 2,0"
        fill={selected ? '#faad14' : '#8a2be2'}
        stroke={selected ? '#faad14' : '#8a2be2'}
        strokeWidth="1"
        transform={`translate(${arrowX}, ${arrowY}) rotate(${angle})`}
        className="er-custom-arrow"
      />
      
      {/* Custom Arrow at source for many-to-many */}
      {isManyToMany && (
        <polygon
          points="0,-4 -8,0 0,4 -2,0"
          fill={selected ? '#faad14' : '#8a2be2'}
          stroke={selected ? '#faad14' : '#8a2be2'}
          strokeWidth="1"
          transform={`translate(${startArrowX}, ${startArrowY}) rotate(${angle + 180})`}
          className="er-custom-arrow"
        />
      )}
      
      {/* Relationship cardinality indicators - positioned directly on the line */}
      <EdgeLabelRenderer>
        {/* Source cardinality (near source table) - positioned on the line */}
        <div
          style={{
            position: 'absolute',
            transform: `translate(-50%, -50%) translate(${sourceX + (labelX - sourceX) * 0.25}px,${sourceY + (labelY - sourceY) * 0.25}px)`,
            pointerEvents: 'none',
            color: selected ? '#faad14' : '#8a2be2',
            fontSize: '12px',
            fontWeight: 'bold',
            textShadow: '1px 1px 2px rgba(0,0,0,0.8), -1px -1px 2px rgba(255,255,255,0.8)',
            background: 'rgba(18, 18, 18, 0.8)',
            padding: '2px 4px',
            borderRadius: '2px',
            border: `1px solid ${selected ? '#faad14' : '#8a2be2'}`,
          }}
          className="nodrag nopan"
        >
          {isOneToMany || isOneToOne ? '1' : 'N'}
        </div>
        
        {/* Target cardinality (near target table) - positioned on the line */}
        <div
          style={{
            position: 'absolute',
            transform: `translate(-50%, -50%) translate(${targetX + (labelX - targetX) * 0.25}px,${targetY + (labelY - targetY) * 0.25}px)`,
            pointerEvents: 'none',
            color: selected ? '#faad14' : '#8a2be2',
            fontSize: '12px',
            fontWeight: 'bold',
            textShadow: '1px 1px 2px rgba(0,0,0,0.8), -1px -1px 2px rgba(255,255,255,0.8)',
            background: 'rgba(18, 18, 18, 0.8)',
            padding: '2px 4px',
            borderRadius: '2px',
            border: `1px solid ${selected ? '#faad14' : '#8a2be2'}`,
          }}
          className="nodrag nopan"
        >
          {isOneToOne ? '1' : 'N'}
        </div>
        
        {/* Main relationship label - centered on the line */}
        <div
          style={{
            position: 'absolute',
            transform: `translate(-50%, -50%) translate(${labelX}px,${labelY}px)`,
            pointerEvents: 'all',
            cursor: 'pointer',
          }}
          className="nodrag nopan"
        >
          <div 
            className="er-relationship-label"
            style={{
              background: selected ? 'rgba(250, 173, 20, 0.95)' : 'rgba(138, 43, 226, 0.95)',
              color: 'white',
              padding: '3px 6px',
              borderRadius: '3px',
              fontSize: '11px',
              fontWeight: 'bold',
              boxShadow: '0 2px 6px rgba(0,0,0,0.3)',
              border: `1px solid ${selected ? '#faad14' : '#8a2be2'}`,
              minWidth: '24px',
              textAlign: 'center',
            }}
          >
            {relationshipType}
          </div>
          {/* Column mapping tooltip - shown on hover */}
          {data?.source_column && data?.target_column && (
            <div style={{
              position: 'absolute',
              top: '100%',
              left: '50%',
              transform: 'translateX(-50%)',
              marginTop: '4px',
              background: 'rgba(0, 0, 0, 0.9)',
              color: 'white',
              padding: '2px 6px',
              borderRadius: '2px',
              fontSize: '9px',
              whiteSpace: 'nowrap',
              opacity: 0,
              transition: 'opacity 0.2s ease',
              pointerEvents: 'none',
            }}
            className="relationship-tooltip"
            >
              {data.source_column} → {data.target_column}
            </div>
          )}
        </div>
      </EdgeLabelRenderer>
    </>
  );
};

// Database Wrapper Node Component
const ERDatabaseWrapperNode = ({ data, id }) => {
  // Auto-calculate dimensions based on schema count and content
  const schemaCount = data.schemaCount || 0;
  const totalTables = data.totalTables || 0;
  
  // Dynamic sizing for database wrapper - much larger to contain all schemas
  const autoWidth = Math.max(1400, data.contentWidth + 200); // Base width + content + padding
  const autoHeight = Math.max(1000, data.contentHeight + 200); // Base height + content + padding
  
  return (
    <div className="er-database-wrapper" style={{ 
      width: data.width || autoWidth, 
      height: data.height || autoHeight 
    }}>
      <div className="er-database-label">
        <FaDatabase />
        <span>{data.name}</span>
        <span style={{ fontSize: '12px', opacity: 0.9 }}>
          ({schemaCount} schemas • {totalTables} tables)
        </span>
      </div>
    </div>
  );
};

// Schema Group Node Component with Auto-sizing
const ERSchemaGroupNode = ({ data, id }) => {
  // Auto-calculate dimensions based on table count and content
  const tableCount = data.tableCount || 0;
  const tablesPerRow = Math.ceil(Math.sqrt(tableCount));
  const rows = Math.ceil(tableCount / tablesPerRow);
  
  // Dynamic sizing based on content
  const autoWidth = Math.max(800, tablesPerRow * 320 + 160); // Base width + table width + padding
  const autoHeight = Math.max(600, rows * 280 + 200); // Base height + table height + padding
  
  return (
    <div className="er-schema-group" style={{ 
      width: data.width || autoWidth, 
      height: data.height || autoHeight 
    }}>
      <div className="er-schema-label">
        <FaLayerGroup />
        <span>{data.name}</span>
        <span style={{ fontSize: '10px', opacity: 0.8 }}>
          ({tableCount} tables)
        </span>
      </div>
      {tableCount === 0 && (
        <div style={{
          position: 'absolute',
          top: '50%',
          left: '50%',
          transform: 'translate(-50%, -50%)',
          color: '#999',
          fontSize: '14px',
          fontStyle: 'italic',
          textAlign: 'center'
        }}>
          <FaTable style={{ fontSize: '24px', marginBottom: '8px', opacity: 0.5 }} />
          <br />
          No tables in this schema
        </div>
      )}
    </div>
  );
};

// Node types for ER Diagram
const nodeTypes = {
  erTable: ERTableNode,
  erSchemaGroup: ERSchemaGroupNode,
  erDatabaseWrapper: ERDatabaseWrapperNode,
};

// Edge types for ER Diagram
const edgeTypes = {
  erRelationship: ERRelationshipEdge,
};

// Mock data for services and data sources
const mockServices = [
  {
    id: 'plan-plus',
    name: 'Plan Plus',
    icon: <FaServer />,
    dataSources: [
      { id: 'postgres-main', name: 'PostgreSQL Main', type: 'PostgreSQL', icon: <FaDatabase /> },
      { id: 'mysql-analytics', name: 'MySQL Analytics', type: 'MySQL', icon: <FaDatabase /> },
      { id: 'mongodb-logs', name: 'MongoDB Logs', type: 'MongoDB', icon: <FaHdd /> }
    ]
  },
  {
    id: 'project-1',
    name: 'Project 1',
    icon: <FaCloud />,
    dataSources: [
      { id: 'snowflake-dw', name: 'Snowflake DW', type: 'Snowflake', icon: <FaCloud /> },
      { id: 'redshift-analytics', name: 'Redshift Analytics', type: 'Redshift', icon: <FaServer /> }
    ]
  },
  {
    id: 'project-2',
    name: 'Project 2',
    icon: <FaProjectDiagram />,
    dataSources: [
      { id: 'bigquery-main', name: 'BigQuery Main', type: 'BigQuery', icon: <FaCloud /> },
      { id: 'postgres-backup', name: 'PostgreSQL Backup', type: 'PostgreSQL', icon: <FaDatabase /> }
    ]
  }
];

// Hierarchical Breadcrumb Component
const HierarchicalBreadcrumb = ({ 
  selectedService, 
  selectedDataSource, 
  onServiceChange, 
  onDataSourceChange 
}) => {
  const [serviceMenuAnchor, setServiceMenuAnchor] = useState(null);
  const [dataSourceMenuAnchor, setDataSourceMenuAnchor] = useState(null);

  const handleServiceClick = (event) => {
    setServiceMenuAnchor(event.currentTarget);
  };

  const handleDataSourceClick = (event) => {
    setDataSourceMenuAnchor(event.currentTarget);
  };

  const handleServiceSelect = (service) => {
    onServiceChange(service);
    setServiceMenuAnchor(null);
    // Reset data source when service changes
    if (service.dataSources.length > 0) {
      onDataSourceChange(service.dataSources[0]);
    }
  };

  const handleDataSourceSelect = (dataSource) => {
    onDataSourceChange(dataSource);
    setDataSourceMenuAnchor(null);
  };

  const handleCloseMenus = () => {
    setServiceMenuAnchor(null);
    setDataSourceMenuAnchor(null);
  };

  return (
    <>
      <Breadcrumbs 
        aria-label="breadcrumb"
        separator="›"
        sx={{
          '& .MuiBreadcrumbs-separator': {
            color: '#666',
            margin: '0 8px',
          },
          fontSize: '14px',
        }}
      >
        <Link
          underline="hover"
          color="primary"
          href="#"
          onClick={(e) => e.preventDefault()}
          sx={{ fontWeight: 500 }}
        >
          Qubit
        </Link>
        
        <Typography color="text.secondary">ER Diagram</Typography>
        
        {/* DBTEZ Services Dropdown */}
        <div 
          className={`er-breadcrumb-dropdown ${serviceMenuAnchor ? 'open' : ''}`}
          onClick={handleServiceClick}
        >
          <Typography 
            color="primary" 
            sx={{ fontWeight: 500, cursor: 'pointer' }}
          >
            DBTEZ Services
          </Typography>
          <FaChevronDown className="dropdown-arrow" />
        </div>
        
        {/* Selected Service */}
        {selectedService && (
          <Typography color="text.secondary" sx={{ display: 'flex', alignItems: 'center', gap: '4px' }}>
            {selectedService.icon}
            {selectedService.name}
          </Typography>
        )}
      </Breadcrumbs>

      {/* Services Menu */}
      <Menu
        anchorEl={serviceMenuAnchor}
        open={Boolean(serviceMenuAnchor)}
        onClose={handleCloseMenus}
        className="er-service-menu"
        PaperProps={{
          className: 'er-service-menu'
        }}
        transformOrigin={{ horizontal: 'left', vertical: 'top' }}
        anchorOrigin={{ horizontal: 'left', vertical: 'bottom' }}
      >
        {mockServices.map((service) => (
          <MenuItem 
            key={service.id} 
            onClick={() => handleServiceSelect(service)}
            selected={selectedService?.id === service.id}
          >
            <ListItemIcon>
              {service.icon}
            </ListItemIcon>
            <ListItemText primary={service.name} />
          </MenuItem>
        ))}
      </Menu>

      {/* Data Sources Menu */}
      {selectedService && (
        <Menu
          anchorEl={dataSourceMenuAnchor}
          open={Boolean(dataSourceMenuAnchor)}
          onClose={handleCloseMenus}
          className="er-service-menu"
          PaperProps={{
            className: 'er-service-menu'
          }}
          transformOrigin={{ horizontal: 'left', vertical: 'top' }}
          anchorOrigin={{ horizontal: 'left', vertical: 'bottom' }}
        >
          {selectedService.dataSources.map((dataSource) => (
            <MenuItem 
              key={dataSource.id} 
              onClick={() => handleDataSourceSelect(dataSource)}
              selected={selectedDataSource?.id === dataSource.id}
            >
              <ListItemIcon>
                {dataSource.icon}
              </ListItemIcon>
              <ListItemText 
                primary={dataSource.name} 
                secondary={dataSource.type}
              />
            </MenuItem>
          ))}
        </Menu>
      )}
    </>
  );
};

// Main ER Diagram Canvas Component
const ERDiagramCanvasContent = () => {
  const [nodes, setNodes, onNodesChange] = useNodesState([]);
  const [edges, setEdges, onEdgesChange] = useEdgesState([]);
  const [isLoading, setIsLoading] = useState(true);
  const [databases, setDatabases] = useState([]);
  const [selectedDatabase, setSelectedDatabase] = useState(null);
  const [showAddMenu, setShowAddMenu] = useState(false);
  
  // ReactFlow instance for programmatic control
  const { fitView } = useReactFlow();
  
  // Service and Data Source selection state
  const [selectedService, setSelectedService] = useState(mockServices[0]); // Default to Plan Plus
  const [selectedDataSource, setSelectedDataSource] = useState(mockServices[0].dataSources[0]); // Default to first data source
  const [isUsingMockData, setIsUsingMockData] = useState(false);
  const [apiError, setApiError] = useState(null);
  const [dataSourceMenuAnchor, setDataSourceMenuAnchor] = useState(null);
  
  // Add Table Modal state
  const [isAddTableModalOpen, setIsAddTableModalOpen] = useState(false);
  const [availableSchemas, setAvailableSchemas] = useState([]);
  const [existingTables, setExistingTables] = useState([]);

  // API Configuration
  const API_BASE_URL = 'https://sandbox.kezel.io/api';
  const token = "abdhsg";
  const orgSlug = "sN05Pjv11qvH";

  // API endpoints
  const ENDPOINTS = {
    DATABASE_LIST: `${API_BASE_URL}/qbt_database_list_get`,
    SCHEMA_LIST: `${API_BASE_URL}/qbt_schema_list_get`,
    TABLE_LIST: `${API_BASE_URL}/qbt_table_list_get`,
    COLUMN_LIST: `${API_BASE_URL}/qbt_column_list_get`,
    SCHEMA_CREATE: `${API_BASE_URL}/qbt_schema_create`,
    SCHEMA_DELETE: `${API_BASE_URL}/qbt_schema_delete`,
    SCHEMA_UPDATE: `${API_BASE_URL}/qbt_schema_update`,
    TABLE_CREATE: `${API_BASE_URL}/qbt_table_create`,
    TABLE_UPDATE: `${API_BASE_URL}/qbt_table_update`,
    TABLE_DELETE: `${API_BASE_URL}/qbt_table_delete`,
    COLUMN_CREATE: `${API_BASE_URL}/qbt_column_create`,
    COLUMN_UPDATE: `${API_BASE_URL}/qbt_column_update`,
    COLUMN_DELETE: `${API_BASE_URL}/qbt_column_delete`
  };

  // Mock database data representing ONE database with multiple schemas (fallback only)
  // const mockDatabaseData = {
  //   id: 1,
  //   name: "Sample Database",
  //   slug: "sample_database",
  //   description: "Sample Database Structure (Mock Data)",
  //   service: selectedService?.name || "Unknown Service",
  //   dataSource: selectedDataSource?.name || "Unknown DataSource",
  //   schemas: [
  //     {
  //       sch: "FINANCE_MART",
  //       tables: [
  //         {
  //           id: 1,
  //           name: "accounts",
  //           table_type: "dimension",
  //           columns: [
  //             { name: "account_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "account_number", data_type: "VARCHAR(20)", is_primary_key: false, is_foreign_key: false },
  //             { name: "account_name", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false },
  //             { name: "account_type", data_type: "VARCHAR(50)", is_primary_key: false, is_foreign_key: false },
  //             { name: "balance", data_type: "DECIMAL(15,2)", is_primary_key: false, is_foreign_key: false },
  //             { name: "created_date", data_type: "TIMESTAMP", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         },
  //         {
  //           id: 2,
  //           name: "transactions",
  //           table_type: "fact",
  //           columns: [
  //             { name: "transaction_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "account_id", data_type: "INTEGER", is_primary_key: false, is_foreign_key: true },
  //             { name: "transaction_date", data_type: "DATE", is_primary_key: false, is_foreign_key: false },
  //             { name: "amount", data_type: "DECIMAL(12,2)", is_primary_key: false, is_foreign_key: false },
  //             { name: "transaction_type", data_type: "VARCHAR(20)", is_primary_key: false, is_foreign_key: false },
  //             { name: "description", data_type: "VARCHAR(255)", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         },
  //         {
  //           id: 3,
  //           name: "customers",
  //           table_type: "dimension",
  //           columns: [
  //             { name: "customer_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "first_name", data_type: "VARCHAR(50)", is_primary_key: false, is_foreign_key: false },
  //             { name: "last_name", data_type: "VARCHAR(50)", is_primary_key: false, is_foreign_key: false },
  //             { name: "email", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false },
  //             { name: "phone", data_type: "VARCHAR(20)", is_primary_key: false, is_foreign_key: false },
  //             { name: "registration_date", data_type: "TIMESTAMP", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         },
  //         {
  //           id: 4,
  //           name: "financial_reports",
  //           table_type: "fact",
  //           columns: [
  //             { name: "report_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "account_id", data_type: "INTEGER", is_primary_key: false, is_foreign_key: true },
  //             { name: "report_date", data_type: "DATE", is_primary_key: false, is_foreign_key: false },
  //             { name: "balance_amount", data_type: "DECIMAL(15,2)", is_primary_key: false, is_foreign_key: false },
  //             { name: "profit_loss", data_type: "DECIMAL(12,2)", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         }
  //       ]
  //     },
  //     {
  //       sch: "Schema100",
  //       tables: [
  //         {
  //           id: 5,
  //           name: "products",
  //           table_type: "dimension",
  //           columns: [
  //             { name: "product_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "product_name", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false },
  //             { name: "category", data_type: "VARCHAR(50)", is_primary_key: false, is_foreign_key: false },
  //             { name: "price", data_type: "DECIMAL(10,2)", is_primary_key: false, is_foreign_key: false },
  //             { name: "stock_quantity", data_type: "INTEGER", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         },
  //         {
  //           id: 6,
  //           name: "orders",
  //           table_type: "fact",
  //           columns: [
  //             { name: "order_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "customer_id", data_type: "INTEGER", is_primary_key: false, is_foreign_key: true },
  //             { name: "product_id", data_type: "INTEGER", is_primary_key: false, is_foreign_key: true },
  //             { name: "order_date", data_type: "DATE", is_primary_key: false, is_foreign_key: false },
  //             { name: "quantity", data_type: "INTEGER", is_primary_key: false, is_foreign_key: false },
  //             { name: "total_amount", data_type: "DECIMAL(12,2)", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         },
  //         {
  //           id: 7,
  //           name: "inventory",
  //           table_type: "stage",
  //           columns: [
  //             { name: "inventory_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "product_id", data_type: "INTEGER", is_primary_key: false, is_foreign_key: true },
  //             { name: "warehouse_location", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false },
  //             { name: "stock_level", data_type: "INTEGER", is_primary_key: false, is_foreign_key: false },
  //             { name: "last_updated", data_type: "TIMESTAMP", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         },
  //         {
  //           id: 8,
  //           name: "suppliers",
  //           table_type: "dimension",
  //           columns: [
  //             { name: "supplier_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "supplier_name", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false },
  //             { name: "contact_email", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false },
  //             { name: "phone_number", data_type: "VARCHAR(20)", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         }
  //       ]
  //     },
  //     {
  //       sch: "New_sch",
  //       tables: [
  //         {
  //           id: 9,
  //           name: "analytics_summary",
  //           table_type: "fact",
  //           columns: [
  //             { name: "summary_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "date", data_type: "DATE", is_primary_key: false, is_foreign_key: false },
  //             { name: "total_revenue", data_type: "DECIMAL(15,2)", is_primary_key: false, is_foreign_key: false },
  //             { name: "total_orders", data_type: "INTEGER", is_primary_key: false, is_foreign_key: false },
  //             { name: "unique_customers", data_type: "INTEGER", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         },
  //         {
  //           id: 10,
  //           name: "user_sessions",
  //           table_type: "stage",
  //           columns: [
  //             { name: "session_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "customer_id", data_type: "INTEGER", is_primary_key: false, is_foreign_key: true },
  //             { name: "session_start", data_type: "TIMESTAMP", is_primary_key: false, is_foreign_key: false },
  //             { name: "session_end", data_type: "TIMESTAMP", is_primary_key: false, is_foreign_key: false },
  //             { name: "pages_viewed", data_type: "INTEGER", is_primary_key: false, is_foreign_key: false },
  //             { name: "device_type", data_type: "VARCHAR(50)", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         },
  //         {
  //           id: 11,
  //           name: "reports",
  //           table_type: "dimension",
  //           columns: [
  //             { name: "report_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
  //             { name: "report_name", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false },
  //             { name: "report_type", data_type: "VARCHAR(50)", is_primary_key: false, is_foreign_key: false },
  //             { name: "created_by", data_type: "VARCHAR(50)", is_primary_key: false, is_foreign_key: false },
  //             { name: "created_date", data_type: "TIMESTAMP", is_primary_key: false, is_foreign_key: false }
  //           ]
  //         }
  //       ]
  //     }
  //   ]
  // };

  // API Functions for fetching real data
  const fetchDatabases = async () => {
    try {
      const response = await axios.post(
        ENDPOINTS.DATABASE_LIST,
        {
          token: token,
          org: orgSlug,
        },
        {
          headers: {
            'Content-Type': 'application/json'
          }
        }
      );
      
      console.log('Database list response:', response.data);
      const databases = response.data.items || [];
      console.log(`Found ${databases.length} databases:`, databases);
      
      return databases.filter(db => db.con); // Filter out databases without connection slug
    } catch (error) {
      console.error('Error fetching databases:', error);
      throw error;
    }
  };

  const fetchSchemas = async (dbSlug) => {
    try {
      console.log(`Fetching schemas for database slug: ${dbSlug}`);
      const response = await axios.post(
        ENDPOINTS.SCHEMA_LIST,
        {
          token: token,
          org: orgSlug,
          con: dbSlug
        },
        {
          headers: {
            'Content-Type': 'application/json'
          }
        }
      );
      
      console.log(`Schema list for database ${dbSlug}:`, response.data);
      
      let schemas = [];
      if (Array.isArray(response.data.items)) {
        schemas = response.data.items.map(item => ({
          name: item.name,
          sch: item.sch, // Use 'sch' as the slug
          description: item.description || "",
          created_at: item.created_at,
          is_validated: item.is_validated,
          database: dbSlug
        }));
      }
      
      console.log(`Number of schemas found for database ${dbSlug}: ${schemas.length}`);
      console.log('Schema details:', schemas.map(s => ({ name: s.name, sch: s.sch })));
      return schemas;
    } catch (error) {
      console.error(`Error fetching schemas for database ${dbSlug}:`, error);
      throw error;
    }
  };

  const fetchTables = async (dbSlug, schemaSlug) => {
    try {
      console.log(`Fetching tables for database: ${dbSlug}, schema: ${schemaSlug}`);
      const response = await axios.post(
        ENDPOINTS.TABLE_LIST,
        {
          token: token,
          org: orgSlug,
          con: dbSlug,
          sch: schemaSlug
        },
        {
          headers: {
            'Content-Type': 'application/json'
          }
        }
      );
      
      console.log(`Table list for schema ${schemaSlug}:`, response.data);
      
      let tables = [];
      if (Array.isArray(response.data.items)) {
        tables = response.data.items.map(item => ({
          id: item.id || `${dbSlug}-${schemaSlug}-${item.name}`,
          name: item.name,
          tbl: item.tbl, // Table slug
          description: item.description || "",
          table_type: item.table_type || "dimension", // Default to dimension if not specified
          created_at: item.created_at,
          database: dbSlug,
          schema: schemaSlug
        }));
      }
      
      console.log(`Number of tables found for schema ${schemaSlug}: ${tables.length}`);
      return tables;
    } catch (error) {
      console.error(`Error fetching tables for schema ${schemaSlug}:`, error);
      throw error;
    }
  };

  const fetchColumns = async (dbSlug, schemaSlug, tableSlug) => {
    try {
      console.log(`Fetching columns for database: ${dbSlug}, schema: ${schemaSlug}, table: ${tableSlug}`);
      const response = await axios.post(
        ENDPOINTS.COLUMN_LIST,
        {
          token: token,
          org: orgSlug,
          con: dbSlug,
          sch: schemaSlug,
          tbl: tableSlug
        },
        {
          headers: {
            'Content-Type': 'application/json'
          }
        }
      );
      
      console.log(`Column list for table ${tableSlug}:`, response.data);
      
      let columns = [];
      if (Array.isArray(response.data.items)) {
        columns = response.data.items.map((item, index) => ({
          name: item.name,
          col: item.col, // Column slug
          data_type: item.data_type || "VARCHAR(255)",
          description: item.description || "",
          is_primary_key: index === 0, // Assume first column is primary key
          is_foreign_key: item.name.toLowerCase().includes('_id') && index > 0, // Simple heuristic for foreign keys
          is_nullable: item.is_nullable || false,
          database: dbSlug,
          schema: schemaSlug,
          table: tableSlug
        }));
      }
      
      console.log(`Number of columns found for table ${tableSlug}: ${columns.length}`);
      return columns;
    } catch (error) {
      console.error(`Error fetching columns for table ${tableSlug}:`, error);
      throw error;
    }
  };

  // Function to fetch complete database structure with schemas, tables, and columns
  const fetchCompleteDatabase = async (dbSlug) => {
    try {
      console.log(`Fetching complete structure for database: ${dbSlug}`);
      
      // Fetch schemas
      const schemas = await fetchSchemas(dbSlug);
      
      // Fetch tables and columns for each schema
      const schemasWithTables = await Promise.all(
        schemas.map(async (schema) => {
          try {
            const tables = await fetchTables(dbSlug, schema.sch);
            
            // Fetch columns for each table
            const tablesWithColumns = await Promise.all(
              tables.map(async (table) => {
                try {
                  const columns = await fetchColumns(dbSlug, schema.sch, table.tbl);
                  return {
                    ...table,
                    columns: columns
                  };
                } catch (columnError) {
                  console.warn(`Error fetching columns for table ${table.name}:`, columnError);
                  return {
                    ...table,
                    columns: [] // Return table with empty columns if column fetch fails
                  };
                }
              })
            );
            
            return {
              ...schema,
              tables: tablesWithColumns
            };
          } catch (tableError) {
            console.warn(`Error fetching tables for schema ${schema.sch}:`, tableError);
            return {
              ...schema,
              tables: [] // Return schema with empty tables if table fetch fails
            };
          }
        })
      );
      
      console.log(`Complete database structure for ${dbSlug}:`, schemasWithTables);
      return schemasWithTables;
    } catch (error) {
      console.error(`Error fetching complete database structure for ${dbSlug}:`, error);
      throw error;
    }
  };

  // Generate dummy ER relationship data based on single database structure
  const generateDummyERData = (database) => {
    const erData = [];
    
    if (database.schemas && database.schemas.length > 0) {
      database.schemas.forEach(schema => {
        if (schema.tables && schema.tables.length > 0) {
            // Create relationships between tables in the same schema
            for (let i = 0; i < schema.tables.length - 1; i++) {
              const sourceTable = schema.tables[i];
              const targetTable = schema.tables[i + 1];
              
              // Create a dummy relationship with varied types
              const relationshipTypes = ['1:N', '1:1', 'N:M'];
              const randomType = relationshipTypes[Math.floor(Math.random() * relationshipTypes.length)];
              
              erData.push({
                source_column_set: [
                  { 
                    table_id: sourceTable.id,
                    column_name: sourceTable.columns?.[0]?.name || 'id',
                    table_name: sourceTable.name,
                    schema_name: schema.sch
                  }
                ],
                destination_column_set: [
                  {
                    table_id: targetTable.id,
                    column_name: `${sourceTable.name}_id`,
                    table_name: targetTable.name,
                    schema_name: schema.sch
                  }
                ],
                relationship_type: randomType
              });
          }
        }
      });
    }
    
    return erData;
  };

  // Fetch databases and generate ER diagram using real API
  useEffect(() => {
    const fetchERData = async () => {
      try {
        setIsLoading(true);
        
        console.log(`Fetching ER data for service: ${selectedService?.name}, data source: ${selectedDataSource?.name}`);
        
        try {
          // Fetch databases from API
          const databases = await fetchDatabases();
          console.log('Fetched databases:', databases);
          
          if (databases && databases.length > 0) {
            // Select database based on selectedDataSource or use first available
            const targetDatabase = databases.find(db => 
              selectedDataSource?.name === db.name ||
              selectedDataSource?.name === db.con ||
              selectedDataSource?.slug === db.con
            ) || databases[0]; // Fallback to first database if no match
            
            console.log('Selected database for ER diagram:', targetDatabase);
            console.log('Available databases:', databases.map(db => ({ name: db.name, con: db.con })));
            
            // Fetch complete structure (schemas, tables, columns) for the selected database
            const schemasWithTables = await fetchCompleteDatabase(targetDatabase.con);
            
            // Create the complete database object using actual API data
            const completeDatabase = {
              id: targetDatabase.id,
              name: targetDatabase.name,
              slug: targetDatabase.con,
              description: targetDatabase.description || `Database: ${targetDatabase.name}`,
              service: selectedService?.name,
              dataSource: selectedDataSource?.name,
              schemas: schemasWithTables
            };
            
            console.log('Complete database structure:', completeDatabase);
            console.log('Schemas in database:', completeDatabase.schemas.map(s => ({
              name: s.name,
              sch: s.sch,
              tableCount: s.tables?.length || 0,
              tables: s.tables?.map(t => t.name) || []
            })));
            
            setDatabases([completeDatabase]); // Wrap in array for compatibility
            setSelectedDatabase(completeDatabase); // Set selected database for modal
            generateERDiagram(completeDatabase);
            setIsUsingMockData(false);
            setApiError(null);
            console.log('Successfully fetched and generated ER diagram from API data');
            
          } else {
            throw new Error('No databases found from API');
          }
        } catch (apiError) {
          // Handle specific error types
          let errorMessage = 'Unknown API error';
          if (apiError.code === 'ERR_NETWORK' || apiError.message.includes('CORS')) {
            errorMessage = 'CORS error - API not accessible from browser';
            console.warn('CORS error detected, using mock data:', apiError.message);
          } else if (apiError.response?.status === 403) {
            errorMessage = '403 Forbidden - Check API token and permissions';
            console.warn('403 Forbidden error, using mock data. Check your API token and permissions.');
          } else if (apiError.response?.status === 401) {
            errorMessage = '401 Unauthorized - Invalid API token';
            console.warn('401 Unauthorized error, using mock data. Check your API token.');
          } else {
            errorMessage = `API Error: ${apiError.message}`;
            console.warn('API error, using mock data:', apiError.message);
          }
          
          setApiError(errorMessage);
          setIsUsingMockData(true);
          
          // Fallback to mock data
          const singleDbData = {
            ...mockDatabaseData,
            service: selectedService?.name,
            dataSource: selectedDataSource?.name
          };
          
          setDatabases([singleDbData]); // Wrap in array for compatibility
          setSelectedDatabase(singleDbData); // Set selected database for modal
          generateERDiagram(singleDbData);
          console.log('Using mock data for ER diagram due to API error');
        }
      } catch (error) {
        console.error('Unexpected error in fetchERData:', error);
        setApiError('Unexpected error occurred');
        setIsUsingMockData(true);
        
        // Fallback to mock data
        const singleDbData = {
          ...mockDatabaseData,
          service: selectedService?.name,
          dataSource: selectedDataSource?.name
        };
        
        setDatabases([singleDbData]);
        setSelectedDatabase(singleDbData); // Set selected database for modal
        generateERDiagram(singleDbData);
      } finally {
        setIsLoading(false);
      }
    };

    // Only fetch if we have selected service and data source
    if (selectedService && selectedDataSource) {
      fetchERData();
    }
  }, [selectedService, selectedDataSource]); // Re-fetch when service or data source changes

  // Auto-fit view when nodes are updated
  useEffect(() => {
    if (nodes.length > 0 && !isLoading) {
      const timer = setTimeout(() => {
        fitView({ 
          padding: 0.1, 
          includeHiddenNodes: false,
          minZoom: 0.15,
          maxZoom: 0.8,
          duration: 600
        });
      }, 300);
      
      return () => clearTimeout(timer);
    }
  }, [nodes, isLoading, fitView]);

  // Auto-alignment layout calculator
  const calculateOptimalLayout = (databaseData) => {
    const layouts = [];
    let totalWidth = 0;
    let maxHeight = 0;

    databaseData.forEach((db) => {
      if (db.schemas && db.schemas.length > 0) {
        const dbLayout = { schemas: [], totalWidth: 0, maxHeight: 0 };
        
        db.schemas.forEach((schema) => {
          if (schema.tables && schema.tables.length > 0) {
            const tableCount = schema.tables.length;
            const tablesPerRow = Math.min(4, Math.ceil(Math.sqrt(tableCount))); // Max 4 tables per row
            const rows = Math.ceil(tableCount / tablesPerRow);
            
            // Calculate optimal schema dimensions
            const schemaWidth = Math.max(800, tablesPerRow * 320 + 160);
            const schemaHeight = Math.max(600, rows * 280 + 200);
            
            dbLayout.schemas.push({
              schema,
              width: schemaWidth,
              height: schemaHeight,
              tablesPerRow,
              rows
            });
            
            dbLayout.totalWidth += schemaWidth + 100; // Add spacing
            dbLayout.maxHeight = Math.max(dbLayout.maxHeight, schemaHeight);
          }
        });
        
        layouts.push(dbLayout);
        totalWidth = Math.max(totalWidth, dbLayout.totalWidth);
        maxHeight += dbLayout.maxHeight + 100; // Add vertical spacing
      }
    });

    return { layouts, totalWidth, maxHeight };
  };

  // Handle adding a new table
  const handleAddTable = async (tableData) => {
    try {
      console.log('Adding new table:', tableData);
      
      // Here you would typically make an API call to create the table
      // For now, we'll add it to the current diagram
      
      // Find the target schema
      const targetSchema = availableSchemas.find(schema => schema.sch === tableData.schema);
      if (!targetSchema) {
        throw new Error('Target schema not found');
      }
      
      // Create a new table object
      const newTable = {
        id: `new-table-${Date.now()}`,
        name: tableData.name,
        description: tableData.description,
        table_type: tableData.table_type,
        columns: tableData.columns,
        schema: tableData.schema,
        database: selectedDatabase?.name || 'Unknown'
      };
      
      // Add the table to existing tables list
      setExistingTables(prev => [...prev, newTable]);
      
      // Update the current database structure and regenerate the diagram
      const updatedDatabase = { ...selectedDatabase };
      const schemaIndex = updatedDatabase.schemas.findIndex(s => s.sch === tableData.schema);
      
      if (schemaIndex !== -1) {
        if (!updatedDatabase.schemas[schemaIndex].tables) {
          updatedDatabase.schemas[schemaIndex].tables = [];
        }
        updatedDatabase.schemas[schemaIndex].tables.push(newTable);
        
        // Regenerate the ER diagram with the new table
        generateERDiagram(updatedDatabase);
        
        console.log('Table added successfully:', newTable);
      }
      
    } catch (error) {
      console.error('Error adding table:', error);
      throw error;
    }
  };

  // Generate ER diagram with Database Wrapper structure
  const generateERDiagram = (database) => {
    const newNodes = [];
    const newEdges = [];
    
    // Generate dummy relationships
    const relationships = generateDummyERData(database);

    // Process the selected database
    console.log('Generating ER diagram for database:', database?.name);
    console.log('Total schemas in database:', database?.schemas?.length);
    console.log('Schema details:', database?.schemas?.map(s => ({ name: s.name, sch: s.sch, tableCount: s.tables?.length || 0 })));
    
    if (database && database.schemas && database.schemas.length > 0) {
      // Calculate database wrapper dimensions
      let totalSchemaWidth = 0;
      let maxSchemaHeight = 0;
      let totalTables = 0;
      
      // Calculate layout for all schemas
      const schemaLayouts = database.schemas.map(schema => {
          const tableCount = (schema.tables && schema.tables.length) || 0;
          totalTables += tableCount;
          
          if (tableCount > 0) {
            const tablesPerRow = Math.min(3, Math.ceil(Math.sqrt(tableCount))); // Max 3 tables per row for better fit
            const rows = Math.ceil(tableCount / tablesPerRow);
            
            // Calculate schema dimensions with proper padding for tables
            const tableWidth = 260; // Width of each table node
            const tableHeight = 280; // Height of each table node (including spacing)
            const tableSpacingCalc = 330; // Horizontal spacing between tables
            const tableRowSpacingCalc = 320; // Vertical spacing between table rows
            const schemaPadding = 200; // Padding around the schema content
            
            // Calculate required width and height based on table layout
            const tableStartX = 90; // Starting X position within schema
            const tableStartY = 120; // Starting Y position within schema
            const requiredWidth = tableStartX + (tablesPerRow * tableSpacingCalc) + 100; // Extra margin
            const requiredHeight = tableStartY + (rows * tableRowSpacingCalc) + 100; // Extra margin
            
            const schemaWidth = Math.max(900, requiredWidth);
            const schemaHeight = Math.max(650, requiredHeight);
            
            console.log(`Schema "${schema.name || schema.sch}" layout: ${tableCount} tables, ${tablesPerRow} per row, ${rows} rows`);
            console.log(`Required dimensions: ${requiredWidth}x${requiredHeight}, Final: ${schemaWidth}x${schemaHeight}`);
            
            totalSchemaWidth += schemaWidth + 150; // Add spacing between schemas
            maxSchemaHeight = Math.max(maxSchemaHeight, schemaHeight);
            
            return {
              schema,
              width: schemaWidth,
              height: schemaHeight,
              tablesPerRow,
              rows,
              tableCount
            };
          } else {
            // Handle schemas with no tables - show empty schema container
            const schemaWidth = 900; // Minimum width for empty schema
            const schemaHeight = 400; // Minimum height for empty schema
            
            totalSchemaWidth += schemaWidth + 150; // Add spacing between schemas
            maxSchemaHeight = Math.max(maxSchemaHeight, schemaHeight);
            
            return {
              schema,
              width: schemaWidth,
              height: schemaHeight,
              tablesPerRow: 0,
              rows: 0,
              tableCount: 0
            };
          }
        });

        console.log('Schema layouts created:', schemaLayouts.length);
        console.log('Schema layout details:', schemaLayouts.map(sl => ({ 
          name: sl.schema.name || sl.schema.sch,
          sch: sl.schema.sch, 
          tableCount: sl.tableCount,
          width: sl.width,
          height: sl.height
        })));

        // Create Database Wrapper Node
        const databaseWrapperId = `database-${database.slug}`;
        const databaseWrapperWidth = Math.max(2000, totalSchemaWidth + 300);
        const databaseWrapperHeight = Math.max(1400, maxSchemaHeight + 400);
        
        newNodes.push({
          id: databaseWrapperId,
          type: 'erDatabaseWrapper',
          position: { x: 50, y: 50 },
          data: {
            name: database.name,
            schemaCount: database.schemas.length,
            totalTables: totalTables,
            width: databaseWrapperWidth,
            height: databaseWrapperHeight,
            contentWidth: totalSchemaWidth,
            contentHeight: maxSchemaHeight
          },
          draggable: true,
          selectable: false,
          style: {
            width: databaseWrapperWidth,
            height: databaseWrapperHeight,
            zIndex: -2
          }
        });

        // Create Schema Group Nodes within Database Wrapper
        let schemaXOffset = 180; // Start position within database wrapper
        const schemaYOffset = 180; // Y position within database wrapper

        schemaLayouts.forEach((schemaLayout, schemaIndex) => {
          const schemaGroupId = `schema-${database.slug}-${schemaLayout.schema.sch}`;

          // Create schema group node
          newNodes.push({
            id: schemaGroupId,
            type: 'erSchemaGroup',
            position: { x: schemaXOffset, y: schemaYOffset },
            data: {
              name: schemaLayout.schema.name || schemaLayout.schema.sch, // Use schema name, fallback to sch
              sch: schemaLayout.schema.sch, // Keep sch for identification
              tableCount: schemaLayout.tableCount,
              width: schemaLayout.width,
              height: schemaLayout.height
            },
            draggable: true,
            selectable: false,
            parentNode: databaseWrapperId,
            extent: 'parent',
            style: {
              width: schemaLayout.width,
              height: schemaLayout.height,
              zIndex: -1
            }
          });

          // Create table nodes within schema (only if tables exist)
          if (schemaLayout.schema.tables && schemaLayout.schema.tables.length > 0) {
            const tableSpacing = 330; // Horizontal spacing between tables
            const tableRowSpacing = 320; // Vertical spacing between table rows
            const tableStartX = 90; // Starting X position within schema (accounting for schema padding)
            const tableStartY = 120; // Starting Y position within schema (accounting for schema padding and label)
            let currentRow = 0;
            let currentCol = 0;

            schemaLayout.schema.tables.forEach((table, tableIndex) => {
            const tableId = `table-${table.id || `${database.slug}-${schemaLayout.schema.sch}-${table.name}`}`;
            
            // Calculate table position within schema (relative to schema, not absolute)
            const tableX = tableStartX + (currentCol * tableSpacing);
            const tableY = tableStartY + (currentRow * tableRowSpacing);
            
            // Debug table positioning
            if (tableIndex === 0) {
              console.log(`Schema "${schemaLayout.schema.name || schemaLayout.schema.sch}" dimensions: ${schemaLayout.width}x${schemaLayout.height}`);
              console.log(`Tables per row: ${schemaLayout.tablesPerRow}, Total rows: ${schemaLayout.rows}`);
              console.log(`Table count: ${schemaLayout.tableCount}`);
            }
            
            // Ensure table position is within schema bounds
            const maxTableX = schemaLayout.width - 280; // Table width + some margin
            const maxTableY = schemaLayout.height - 300; // Table height + some margin
            
            if (tableX > maxTableX || tableY > maxTableY) {
              console.warn(`Table "${table.name}" position (${tableX}, ${tableY}) exceeds schema bounds (${schemaLayout.width}x${schemaLayout.height})`);
            }
            
            // Add primary key and foreign key indicators to columns
            const enhancedColumns = (table.columns || []).map((col, colIndex) => ({
              ...col,
              is_primary_key: colIndex === 0, // First column as primary key
              is_foreign_key: relationships.some(rel => 
                rel.destination_column_set.some(dest => 
                  dest.column_name === col.name && dest.table_name === table.name
                )
              )
            }));

            newNodes.push({
              id: tableId,
              type: 'erTable',
              position: { x: tableX, y: tableY },
              data: {
                ...table,
                columns: enhancedColumns,
                schema: schemaLayout.schema.sch,
                database: database.name
              },
              draggable: true,
              parentNode: schemaGroupId,
              extent: 'parent'
            });

              // Move to next position
              currentCol++;
              if (currentCol >= schemaLayout.tablesPerRow) {
                currentCol = 0;
                currentRow++;
              }
            });
          } else {
            // Schema has no tables - show empty state message
            console.log(`Schema "${schemaLayout.schema.name || schemaLayout.schema.sch}" (${schemaLayout.schema.sch}) has no tables`);
          }

          schemaXOffset += schemaLayout.width + 150; // Move to next schema position
        });
    }

    // Create edges for relationships
    relationships.forEach((rel, index) => {
      const sourceTableId = `table-${rel.source_column_set[0]?.table_id || `${rel.source_column_set[0]?.schema_name}-${rel.source_column_set[0]?.table_name}`}`;
      const targetTableId = `table-${rel.destination_column_set[0]?.table_id || `${rel.destination_column_set[0]?.schema_name}-${rel.destination_column_set[0]?.table_name}`}`;

      // Check if both nodes exist
      const sourceExists = newNodes.some(node => node.id === sourceTableId);
      const targetExists = newNodes.some(node => node.id === targetTableId);

      if (sourceExists && targetExists) {
        newEdges.push({
          id: `relationship-${index}`,
          type: 'erRelationship',
          source: sourceTableId,
          target: targetTableId,
          data: {
            relationship_type: rel.relationship_type,
            source_column: rel.source_column_set[0]?.column_name,
            target_column: rel.destination_column_set[0]?.column_name
          },
          style: {
            stroke: '#8a2be2',
            strokeWidth: 2
          },
          animated: false
        });
      }
    });

    setNodes(newNodes);
    setEdges(newEdges);
    
    // Update available schemas and existing tables for the modal
    if (database && database.schemas) {
      setAvailableSchemas(database.schemas);
      
      // Collect all existing tables from all schemas
      const allTables = [];
      database.schemas.forEach(schema => {
        if (schema.tables) {
          schema.tables.forEach(table => {
            allTables.push({
              ...table,
              schema: schema.sch,
              schemaName: schema.name || schema.sch
            });
          });
        }
      });
      setExistingTables(allTables);
    }
    
    // Auto-fit view after a short delay to ensure nodes are rendered
    setTimeout(() => {
      if (newNodes.length > 0) {
        fitView({ 
          padding: 0.1, 
          includeHiddenNodes: false,
          minZoom: 0.15,
          maxZoom: 0.8,
          duration: 800 // Smooth animation
        });
      }
    }, 300);
  };

  const onConnect = useCallback(
    (params) => setEdges((eds) => addEdge({
      ...params,
      type: 'erRelationship',
      data: { 
        relationship_type: '1:N',
        source_column: 'id',
        target_column: 'foreign_key_id'
      },
      style: {
        stroke: '#8a2be2',
        strokeWidth: 2
      },
      animated: true
    }, eds)),
    [setEdges]
  );

  const handleAddClick = () => {
    setShowAddMenu(!showAddMenu);
  };

  const handleDataSourceSelect = (dataSource) => {
    setSelectedDataSource(dataSource);
    setDataSourceMenuAnchor(null);
  };

  const handleCloseMenus = () => {
    setDataSourceMenuAnchor(null);
  };

  if (isLoading) {
    return (
      <div style={{
        width: '100vw',
        height: '100vh',
        display: 'flex',
        flexDirection: 'column',
        background: '#121212',
        color: '#fff',
        minWidth: '1200px',
        minHeight: '800px'
      }}>
        {/* Breadcrumb Header */}
        <div style={{ padding: '20px 20px 0 20px' }}>
          <HierarchicalBreadcrumb
            selectedService={selectedService}
            selectedDataSource={selectedDataSource}
            onServiceChange={setSelectedService}
            onDataSourceChange={setSelectedDataSource}
          />
        </div>
        
        {/* Loading Content */}
        <div style={{
          flex: 1,
          display: 'flex',
          alignItems: 'center',
          justifyContent: 'center'
        }}>
          <div style={{ textAlign: 'center' }}>
            <FaProjectDiagram style={{ fontSize: '48px', color: '#8a2be2', marginBottom: '16px' }} />
            <h3>Loading ER Diagram...</h3>
            <p>Fetching schema for {selectedService?.name} - {selectedDataSource?.name}</p>
          </div>
        </div>
      </div>
    );
  }

  return (
    <div style={{ width: '100vw', height: '100vh', position: 'relative', display: 'flex', flexDirection: 'column', minWidth: '1200px', minHeight: '800px' }}>
      <style dangerouslySetInnerHTML={{ __html: generateERStyles() }} />
      
      {/* Breadcrumb Header */}
      <div style={{ padding: '0 20px 20px 20px', flexShrink: 0 }}>
        <HierarchicalBreadcrumb
          selectedService={selectedService}
          selectedDataSource={selectedDataSource}
          onServiceChange={setSelectedService}
          onDataSourceChange={setSelectedDataSource}
        />
      </div>
      
      {/* ReactFlow Container */}
      <div style={{ flex: 1, position: 'relative', minHeight: '700px' }}>
        <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onConnect={onConnect}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        connectionLineStyle={{
          stroke: '#8a2be2',
          strokeWidth: 2,
        }}
        connectionLineType="bezier"
        fitView
        fitViewOptions={{
          padding: 0.15,
          includeHiddenNodes: false,
          minZoom: 0.2,
          maxZoom: 1.0,
        }}
        style={{
          background: '#121212',
        }}
        defaultViewport={{ x: 0, y: 0, zoom: 0.8 }}
        minZoom={0.2}
        maxZoom={2.0}
        attributionPosition="bottom-left"
      >
        <Controls 
          style={{
            background: 'rgba(26, 26, 26, 0.8)',
            border: '1px solid rgba(138, 43, 226, 0.3)',
            borderRadius: '8px'
          }}
        />
        <MiniMap 
          style={{
            background: 'rgba(26, 26, 26, 0.8)',
            border: '1px solid rgba(138, 43, 226, 0.3)',
            borderRadius: '8px'
          }}
          nodeColor={(node) => {
            if (node.type === 'erTable') {
              switch(node.data.table_type) {
                case 'stage': return '#00a99d';
                case 'fact': return '#fa8c16';
                case 'dimension': return '#52c41a';
                default: return '#8a2be2';
              }
            }
            return '#666';
          }}
        />
        <Background 
          color="#333" 
          gap={20} 
          size={1}
          style={{ background: '#121212' }}
        />
        
        <Panel position="top-left">
          <div style={{
            background: 'rgba(26, 26, 26, 0.9)',
            padding: '12px 16px',
            borderRadius: '8px',
            border: '1px solid rgba(138, 43, 226, 0.3)',
            color: '#fff',
            fontSize: '14px',
            backdropFilter: 'blur(5px)',
            minWidth: '280px'
          }}>
            <div style={{ display: 'flex', alignItems: 'center', gap: '8px', marginBottom: '8px' }}>
              <FaProjectDiagram style={{ color: '#8a2be2' }} />
              <strong>Entity Relationship Diagram</strong>
            </div>
            <div style={{ fontSize: '12px', opacity: 0.8, borderTop: '1px solid rgba(138, 43, 226, 0.2)', paddingTop: '8px' }}>
              <div style={{ marginBottom: '4px' }}>
                Tables: {nodes.filter(n => n.type === 'erTable').length} • 
                Relationships: {edges.length} • 
                Schemas: {nodes.filter(n => n.type === 'erSchemaGroup').length}
              </div>
              {isUsingMockData && (
                <div style={{ 
                  fontSize: '11px', 
                  color: '#faad14', 
                  display: 'flex', 
                  alignItems: 'center', 
                  gap: '4px',
                  marginTop: '4px'
                }}>
                  <span style={{ 
                    width: '6px', 
                    height: '6px', 
                    background: '#faad14', 
                    borderRadius: '50%' 
                  }}></span>
                  Using mock data - {apiError}
                </div>
              )}
            </div>
          </div>
        </Panel>
        
        {/* Relationship Legend Panel */}
        <Panel position="bottom-left">
          <div style={{
            background: 'rgba(26, 26, 26, 0.9)',
            padding: '12px 16px',
            borderRadius: '8px',
            border: '1px solid rgba(138, 43, 226, 0.3)',
            color: '#fff',
            fontSize: '12px',
            backdropFilter: 'blur(5px)',
            minWidth: '200px'
          }}>
            <div style={{ display: 'flex', alignItems: 'center', gap: '8px', marginBottom: '8px' }}>
              <FaLink style={{ color: '#8a2be2' }} />
              <strong>Relationship Types</strong>
            </div>
            <div style={{ fontSize: '11px', lineHeight: '1.4' }}>
              <div style={{ display: 'flex', alignItems: 'center', gap: '8px', marginBottom: '4px' }}>
                <span style={{ color: '#8a2be2' }}>→</span>
                <span><strong>1:N</strong> - One to Many</span>
              </div>
              <div style={{ display: 'flex', alignItems: 'center', gap: '8px', marginBottom: '4px' }}>
                <span style={{ color: '#8a2be2' }}>→</span>
                <span><strong>1:1</strong> - One to One</span>
              </div>
              <div style={{ display: 'flex', alignItems: 'center', gap: '8px', marginBottom: '4px' }}>
                <span style={{ color: '#8a2be2' }}>↔</span>
                <span><strong>N:M</strong> - Many to Many</span>
              </div>
              <div style={{ fontSize: '10px', opacity: 0.7, marginTop: '8px', borderTop: '1px solid rgba(138, 43, 226, 0.2)', paddingTop: '6px' }}>
                • Arrows show relationship direction<br/>
                • Numbers show cardinality<br/>
                • Click edges to select them
              </div>
            </div>
          </div>
        </Panel>
        </ReactFlow>
      </div>

      {/* Data Sources Menu */}
      {selectedService && (
        <Menu
          anchorEl={dataSourceMenuAnchor}
          open={Boolean(dataSourceMenuAnchor)}
          onClose={handleCloseMenus}
          className="er-service-menu"
          PaperProps={{
            className: 'er-service-menu'
          }}
          transformOrigin={{ horizontal: 'left', vertical: 'top' }}
          anchorOrigin={{ horizontal: 'left', vertical: 'bottom' }}
        >
          {selectedService.dataSources.map((dataSource) => (
            <MenuItem 
              key={dataSource.id} 
              onClick={() => handleDataSourceSelect(dataSource)}
              selected={selectedDataSource?.id === dataSource.id}
            >
              <ListItemIcon>
                {dataSource.icon}
              </ListItemIcon>
              <ListItemText 
                primary={dataSource.name} 
                secondary={dataSource.type}
              />
            </MenuItem>
          ))}
        </Menu>
      )}

      {/* Add Table Button */}
      <button
        className="er-add-button"
        onClick={() => setIsAddTableModalOpen(true)}
        title="Add New Table"
        disabled={!availableSchemas.length}
      >
        <FaPlus />
      </button>

      {/* Add Table Modal */}
      <AddTableModal
        open={isAddTableModalOpen}
        onClose={() => setIsAddTableModalOpen(false)}
        onAddTable={handleAddTable}
        schemas={availableSchemas}
        existingTables={existingTables}
      />
  
    </div>
  );
};

// Wrapper component with ReactFlowProvider
const ERDiagramCanvas = () => {
  return (
    <ReactFlowProvider>
      <ERDiagramCanvasContent />
    </ReactFlowProvider>
  );
};

export default ERDiagramCanvas;