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 './ERDiagramCanvas.scss';
import axios from 'axios';

// Import icons from react-icons
import { 
  FaDatabase, 
  FaTable, 
  FaTimes, 
  FaKey, 
  FaLink,
  FaLayerGroup,
  FaColumns,
  FaProjectDiagram,
  FaChevronDown,
  FaChevronRight,
  FaServer,
  FaCloud,
  FaHdd,
  FaEdit
} 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';
// import UpdateTableModal from './UpdateTableModal';



// Custom Table Node Component for Data Entity
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}
      />
      <Handle
        type="source"
        position={Position.Right}
      />
      
      <div className={`er-table-header ${data.table_type || 'default'}`}>
        <div className="er-table-header-content">
          {getTableIcon()}
          <span>{data.name}</span>
        </div>
        {/* <button 
          className="er-table-update-btn"
          onClick={(e) => {
            e.stopPropagation();
            data.onUpdateTable && data.onUpdateTable(data);
          }}
          title="Update Table"
        >
          <FaEdit />
        </button> */}
      </div>
      
      <ul className="er-column-list">
        {data.columns && data.columns.map((column, index) => {
          // Determine the column's role in relationships
          const isSourceKey = column.is_source_key;
          const isForeignKey = column.is_foreign_key;
          const isPrimaryKey = column.is_primary_key;
          
          // Create CSS classes for highlighting - ONLY based on API relationship data
          let columnClasses = "er-column-item";
          if (isSourceKey) columnClasses += " er-source-key";
          if (isForeignKey) columnClasses += " er-foreign-key-highlight";
          
          // Debug logging for CSS classes - only for actual relationships
          if (isSourceKey || isForeignKey) {
            console.log(`🎨 Column ${column.name} classes: ${columnClasses}`);
            console.log(`🎨 isSourceKey: ${isSourceKey}, isForeignKey: ${isForeignKey}`);
            console.log(`Column data:`, column);
          }
          
          return (
            <li key={index} className={columnClasses}>
              {isPrimaryKey && <FaKey className="er-primary-key" />}
              {isForeignKey && <FaLink className="er-foreign-key" />}
              {isSourceKey && !isPrimaryKey && <FaKey className="er-source-key-icon" />}
              <span className="er-column-name">{column.name}</span>
              <span className="er-column-type">{column.data_type}</span>
              {/* Show relationship info on hover */}
              {column.relationship_info && column.relationship_info.length > 0 && (
                <div className="er-relationship-tooltip">
                  {column.relationship_info.map((rel, relIndex) => (
                    <div key={relIndex} className="er-relationship-info">
                      {isSourceKey ? `→ ${rel.targetTable}.${rel.targetColumn}` : `← ${rel.sourceTable}.${rel.sourceColumn}`}
                    </div>
                  ))}
                </div>
              )}
            </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={{ 
        '--db-width': `${data.width || autoWidth}px`,
        '--db-height': `${data.height || autoHeight}px`
      }}
    >
      <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={{ 
        '--schema-width': `${data.width || autoWidth}px`,
        '--schema-height': `${data.height || autoHeight}px`
      }}
    >
      <div className="er-schema-label">
        <FaLayerGroup />
        <span>{data.name}</span>
        <span style={{ fontSize: '10px', opacity: 0.8 }}>
          ({tableCount} tables)
        </span>
      </div>
      {tableCount === 0 && (
        <div className="er-schema-empty">
          <FaTable style={{ fontSize: '24px', marginBottom: '8px', opacity: 0.5 }} />
          <br />
          No tables in this schema
        </div>
      )}
    </div>
  );
};

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

// Edge types for Data Entity
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">Data Entity</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 Data Entity 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([]);

  // Update Table Modal state
  const [isUpdateTableModalOpen, setIsUpdateTableModalOpen] = useState(false);
  const [selectedTableForUpdate, setSelectedTableForUpdate] = useState(null);

  // 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`,
    DATASOURCE_KEY_LIST: `${API_BASE_URL}/qbt_datasource_key_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",
            tbl: "accounts", // Add tbl field for consistency
            table_type: "dimension",
            columns: [
              { name: "account_id", col: "account_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
              { name: "account_number", col: "account_number", data_type: "VARCHAR(20)", is_primary_key: false, is_foreign_key: false },
              { name: "account_name", col: "account_name", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false },
              { name: "account_type", col: "account_type", data_type: "VARCHAR(50)", is_primary_key: false, is_foreign_key: false },
              { name: "balance", col: "balance", data_type: "DECIMAL(15,2)", is_primary_key: false, is_foreign_key: false },
              { name: "created_date", col: "created_date", data_type: "TIMESTAMP", is_primary_key: false, is_foreign_key: false }
            ]
          },
          {
            id: 2,
            name: "transactions",
            tbl: "transactions", // Add tbl field for consistency
            table_type: "fact",
            columns: [
              { name: "transaction_id", col: "transaction_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
              { name: "account_id", col: "account_id", data_type: "INTEGER", is_primary_key: false, is_foreign_key: true },
              { name: "transaction_date", col: "transaction_date", data_type: "DATE", is_primary_key: false, is_foreign_key: false },
              { name: "amount", col: "amount", data_type: "DECIMAL(12,2)", is_primary_key: false, is_foreign_key: false }
            ]
          }
        ]
      },
      {
        sch: "SALES_MART",
        name: "Sales Mart",
        tables: [
          {
            id: 3,
            name: "customers",
            tbl: "customers", // Add tbl field for consistency
            table_type: "dimension",
            columns: [
              { name: "customer_id", col: "customer_id", data_type: "INTEGER", is_primary_key: true, is_foreign_key: false },
              { name: "customer_name", col: "customer_name", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false },
              { name: "email", col: "email", data_type: "VARCHAR(100)", is_primary_key: false, is_foreign_key: false }
            ]
          }
        ]
      }
    ],
    relationships: [
      {
        id: "relationship-1",
        sourceTable: "accounts",
        sourceSchema: "FINANCE_MART",
        sourceColumn: "account_id",
        sourceKeyName: "pk_accounts",
        targetTable: "transactions",
        targetSchema: "FINANCE_MART",
        targetColumn: "account_id",
        targetKeyName: "fk_transactions_account",
        relationship_type: "1:N"
      }
    ]
  };

  // 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 relationships/foreign keys from the API
  const fetchRelationships = async (dbSlug) => {
    try {
      console.log(`Fetching relationships for database: ${dbSlug}`);
      const response = await axios.post(
        ENDPOINTS.DATASOURCE_KEY_LIST,
        {
          token: token,
          org: orgSlug,
          con: dbSlug
        },
        {
          headers: {
            'Content-Type': 'application/json'
          }
        }
      );
      
      console.log(`Relationships response for database ${dbSlug}:`, response.data);
      
      let relationships = [];
      if (response.data.status === 200 && Array.isArray(response.data.items)) {
        relationships = response.data.items.map((item, index) => {
          // Extract source and destination information
          const source = item.source?.[0];
          const destination = item.destination?.[0];
          
          if (source && destination) {
            return {
              id: `relationship-${index}`,
              sourceTable: source.table_slug,
              sourceSchema: source.schema_slug,
              sourceColumn: source.column_slug,
              sourceKeyName: source.key_name,
              targetTable: destination.table_slug,
              targetSchema: destination.schema_slug,
              targetColumn: destination.column_slug,
              targetKeyName: destination.key_name,
              relationship_type: '1:N' // Default relationship type, can be enhanced later
            };
          }
          return null;
        }).filter(Boolean); // Remove null entries
      }
      
      console.log(`Number of relationships found for database ${dbSlug}: ${relationships.length}`);
      console.log('Parsed relationships:', relationships);
      return relationships;
    } catch (error) {
      console.error(`Error fetching relationships for database ${dbSlug}:`, error);
      // Return empty array on error to prevent breaking the diagram
      return [];
    }
  };

  // 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;
    }
  };




  // Fetch databases and generate Data Entity 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 Data Entity:', 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);
            
            // Fetch relationships for the selected database
            const relationships = await fetchRelationships(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,
              relationships: relationships
            };
            
            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 Data Entity 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 Data Entity 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]);

  // Regenerate diagram when selectedDatabase changes (including after adding new table)
  useEffect(() => {
    if (selectedDatabase && selectedDatabase.schemas) {
      console.log('Selected database changed, regenerating diagram...');
      generateERDiagram(selectedDatabase);
    }
  }, [selectedDatabase]);

  // 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);
      
      // Find the target schema
      const targetSchema = availableSchemas.find(schema => schema.sch === tableData.schema);
      if (!targetSchema) {
        throw new Error('Target schema not found');
      }

      // Prepare API payload
      const apiPayload = {
        token: token,
        org: orgSlug,
        con: selectedDatabase?.slug || "my_dwh", // Use selected database slug
        sch: tableData.schema,
        name: tableData.name,
        external_name: null,
        table_type: tableData.table_type,
        description: tableData.description,
        columns: tableData.columns.map(col => ({
          column_name: col.name,
          data_type: col.data_type,
          is_nullable: col.is_nullable
        })),
        keys: tableData.keys.map(key => ({
          key_name: key.name,
          key_type: key.key_type,
          key_columns: [{
            column_name: key.column_name,
            sequence: key.sequence || 1
          }]
        }))
      };

      console.log('API Payload:', apiPayload);

      // Make API call to create the table
      const response = await axios.post(ENDPOINTS.TABLE_CREATE, apiPayload, {
        headers: {
          'Content-Type': 'application/json'
        }
      });

      console.log('API Response:', response.data);

      if (response.data.status === 200) {
        // API call successful, create new table object for local state
        const newTable = {
          id: response.data.id || `new-table-${Date.now()}`,
          name: tableData.name,
          tbl: response.data.tbl || tableData.name.toLowerCase(), // Table slug from API or fallback
          description: tableData.description,
          table_type: tableData.table_type,
          columns: tableData.columns,
          schema: tableData.schema,
          database: selectedDatabase?.slug || 'Unknown',
          created_at: new Date().toISOString()
        };
        
        console.log('Table created successfully via API:', newTable);
        
        // Update the current database structure immediately
        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);
          
          console.log(`Added table "${newTable.name}" to schema "${tableData.schema}"`);
          console.log('Updated database structure:', updatedDatabase);
          
          // Update the selected database state
          setSelectedDatabase(updatedDatabase);
          
          // Update databases list
          setDatabases(prevDatabases => {
            return prevDatabases.map(db => {
              if (db.id === selectedDatabase?.id) {
                return updatedDatabase;
              }
              return db;
            });
          });
          
          // Regenerate the Data Entity with the new table immediately
          console.log('Scheduling diagram regeneration...');
          setTimeout(() => {
            console.log('Regenerating diagram with updated database...');
            generateERDiagram(updatedDatabase);
          }, 100);
        }
        
        // Update existing tables list for the modal
        setExistingTables(prev => [...prev, newTable]);
        
        // Close the modal immediately after successful creation
        setIsAddTableModalOpen(false);
        
        console.log('Table added successfully and modal closed');
        
        // Force a small delay to ensure all state updates are processed and then fit view
        setTimeout(() => {
          if (fitView) {
            fitView({ 
              padding: 0.1, 
              includeHiddenNodes: false,
              minZoom: 0.15,
              maxZoom: 0.8,
              duration: 800
            });
          }
        }, 500);
        
      } else {
        throw new Error(response.data.message || 'Failed to create table');
      }
      
    } catch (error) {
      console.error('Error adding table:', error);
      
      // Always close the modal on error to prevent it from being stuck
      setIsAddTableModalOpen(false);
      
      // Check if it's an API error
      if (error.response) {
        const errorMessage = error.response.data?.message || `API Error: ${error.response.status}`;
        throw new Error(errorMessage);
      } else if (error.request) {
        throw new Error('Network error: Unable to connect to the API');
      } else {
        throw new Error(error.message || 'Unknown error occurred');
      }
    }
  };

  // Handle update table action
  const handleUpdateTable = (tableData) => {
    console.log('Opening update modal for table:', tableData);
    setSelectedTableForUpdate(tableData);
    setIsUpdateTableModalOpen(true);
  };

  // Handle table update submission
  const handleUpdateTableSubmit = async (updatedTableData) => {
    try {
      console.log('Updating table with data:', updatedTableData);
      
      // The API call is handled in the UpdateTableModal component
      // After successful update, we need to refresh the diagram
      
      // Close the modal
      setIsUpdateTableModalOpen(false);
      setSelectedTableForUpdate(null);
      
      // Refresh the diagram by re-fetching data
      // You can implement a more efficient update by modifying the existing nodes
      // For now, let's trigger a refresh of the current database
      if (selectedDatabase) {
        generateERDiagram(selectedDatabase);
      }
      
      console.log('Table updated successfully');
      
    } catch (error) {
      console.error('Error updating table:', error);
      // Error handling is done in the UpdateTableModal component
    }
  };

  // Generate Data Entity with Database Wrapper structure
  const generateERDiagram = (database) => {
    console.log('🔄 Starting Data Entity generation...');
    console.log('Database received:', database?.name);
    console.log('Total schemas:', database?.schemas?.length);
    console.log('Relationships received:', database?.relationships?.length);
    console.log('Relationship details:', database?.relationships);
    
    // Debug each relationship in detail
    if (database?.relationships && database.relationships.length > 0) {
      database.relationships.forEach((rel, index) => {
        console.log(`🔗 Relationship ${index + 1}:`);
        console.log(`   Source: ${rel.sourceSchema}.${rel.sourceTable}.${rel.sourceColumn}`);
        console.log(`   Target: ${rel.targetSchema}.${rel.targetTable}.${rel.targetColumn}`);
        console.log(`   Type: ${rel.relationship_type}`);
      });
    }
    
    const newNodes = [];
    const newEdges = []; // Will be populated with relationship edges

    // Process the selected database
    console.log('Generating Data Entity 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,
      tableNames: s.tables?.map(t => t.name) || []
    })));
    
    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 tableSpacingCalc = 330; // Horizontal spacing between tables
            const tableRowSpacingCalc = 320; // Vertical spacing between table rows
            
            // 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(1000, requiredWidth);
            const schemaHeight = Math.max(850, 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}`}`;
            
            // Debug table structure
            if (tableIndex === 0) {
              console.log(`📋 Sample table structure:`, table);
              console.log(`📋 Sample column structure:`, table.columns?.[0]);
            }
            
            // 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 based on relationships
            const enhancedColumns = (table.columns || []).map((col, colIndex) => {
              // Get table identifier - use table_slug from API or fallback to name
              const tableSlug = table.tbl || table.table_slug || table.name;
              // Get schema identifier
              const schemaSlug = schemaLayout.schema.sch || schemaLayout.schema.schema_slug || schemaLayout.schema.name;
              // Get column identifier - use column_slug from the column data
              const columnSlug = col.col || col.column_slug || col.name;
              
              // Check if this column is a foreign key (target/destination of a relationship)
              const isForeignKey = database.relationships?.some(rel => {
                const matchesTable = rel.targetTable === tableSlug;
                const matchesSchema = rel.targetSchema === schemaSlug;
                const matchesColumn = rel.targetColumn === columnSlug;
                
                return matchesTable && matchesSchema && matchesColumn;
              }) || false;
              
              // Check if this column is a source key (source of a relationship)
              const isSourceKey = database.relationships?.some(rel => {
                const matchesTable = rel.sourceTable === tableSlug;
                const matchesSchema = rel.sourceSchema === schemaSlug;
                const matchesColumn = rel.sourceColumn === columnSlug;
                
                return matchesTable && matchesSchema && matchesColumn;
              }) || false;
              
              // Debug logging for relationship matching
              if (database.relationships && database.relationships.length > 0) {
                // Only log for the first few columns to avoid spam
                if (colIndex < 3) {
                  console.log(`🔍 Checking column: table=${tableSlug}, schema=${schemaSlug}, column=${columnSlug}`);
                  console.log(`🔍 Available relationships:`, database.relationships.map(r => 
                    `${r.sourceSchema}.${r.sourceTable}.${r.sourceColumn} -> ${r.targetSchema}.${r.targetTable}.${r.targetColumn}`
                  ));
                }
                
                if (isForeignKey || isSourceKey) {
                  console.log(`🔗 MATCH FOUND! Column ${tableSlug}.${columnSlug} - isSourceKey: ${isSourceKey}, isForeignKey: ${isForeignKey}`);
                  
                  // Find the matching relationship for more details
                  const matchingRel = database.relationships.find(rel => 
                    (rel.targetTable === tableSlug && rel.targetSchema === schemaSlug && rel.targetColumn === columnSlug) ||
                    (rel.sourceTable === tableSlug && rel.sourceSchema === schemaSlug && rel.sourceColumn === columnSlug)
                  );
                  console.log(`🔗 Matching relationship:`, matchingRel);
                }
              }
              
              return {
                ...col,
                is_primary_key: col.is_primary_key || colIndex === 0, // Use existing or first column as primary key
                is_foreign_key: isForeignKey,
                is_source_key: isSourceKey,
                relationship_info: database.relationships?.filter(rel => {
                  const isTarget = rel.targetTable === tableSlug && rel.targetSchema === schemaSlug && rel.targetColumn === columnSlug;
                  const isSource = rel.sourceTable === tableSlug && rel.sourceSchema === schemaSlug && rel.sourceColumn === columnSlug;
                  return isTarget || isSource;
                }) || []
              };
            });

            newNodes.push({
              id: tableId,
              type: 'erTable',
              position: { x: tableX, y: tableY },
              data: {
                ...table,
                columns: enhancedColumns,
                schema: schemaLayout.schema.sch,
                database: database.name
                // onUpdateTable: handleUpdateTable
              },
              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 relationship edges based on API data
    if (database.relationships && database.relationships.length > 0) {
      console.log('Creating relationship edges:', database.relationships.length);
      
      database.relationships.forEach((relationship, index) => {
        // Find source and target table nodes by matching table names and schemas
        const sourceNode = newNodes.find(node => 
          node.type === 'erTable' && 
          node.data?.schema === relationship.sourceSchema &&
          (node.data?.name === relationship.sourceTable || node.data?.tbl === relationship.sourceTable)
        );
        
        const targetNode = newNodes.find(node => 
          node.type === 'erTable' && 
          node.data?.schema === relationship.targetSchema &&
          (node.data?.name === relationship.targetTable || node.data?.tbl === relationship.targetTable)
        );
        
        if (sourceNode && targetNode) {
          const edgeId = `edge-${relationship.id || index}`;
          
          console.log(`Creating edge: ${sourceNode.id} -> ${targetNode.id}`);
          console.log(`Relationship: ${relationship.sourceTable}.${relationship.sourceColumn} -> ${relationship.targetTable}.${relationship.targetColumn}`);
          
          newEdges.push({
            id: edgeId,
            source: sourceNode.id,
            target: targetNode.id,
            type: 'erRelationship',
            data: {
              relationship_type: relationship.relationship_type || '1:N',
              source_column: relationship.sourceColumn,
              target_column: relationship.targetColumn,
              source_key_name: relationship.sourceKeyName,
              target_key_name: relationship.targetKeyName
            },
            style: {
              stroke: '#8a2be2',
              strokeWidth: 2,
            },
            markerEnd: {
              type: MarkerType.ArrowClosed,
              color: '#8a2be2',
            },
            animated: false,
            selectable: true
          });
        } else {
          console.warn(`Could not find nodes for relationship: ${relationship.sourceTable} -> ${relationship.targetTable}`);
          console.warn(`Source node found: ${!!sourceNode}, Target node found: ${!!targetNode}`);
          console.warn(`Looking for source ID: ${sourceTableId}, target ID: ${targetTableId}`);
        }
      });
      
      console.log(`Created ${newEdges.length} relationship edges out of ${database.relationships.length} relationships`);
    } else {
      console.log('No relationships found in database data');
    }

    console.log('📊 Setting nodes and edges...');
    console.log('New nodes count:', newNodes.length);
    console.log('New edges count:', newEdges.length);
    
    setNodes(newNodes);
    setEdges(newEdges);
    
    console.log('✅ Data Entity generation completed');
    
    // 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);
  };

  // onConnect removed - no connections allowed between tables

  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 Data Entity...</h3>
            <p>Fetching schema for {selectedService?.name} - {selectedDataSource?.name}</p>
          </div>
        </div>
      </div>
    );
  }

  return (
    <div className="er-diagram-container">
      {/* Breadcrumb Header */}
      <div className="er-breadcrumb-header">
        <HierarchicalBreadcrumb
          selectedService={selectedService}
          selectedDataSource={selectedDataSource}
          onServiceChange={setSelectedService}
          onDataSourceChange={setSelectedDataSource}
        />
      </div>
      
      {/* ReactFlow Container */}
      <div className="er-reactflow-container">
        <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        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(255, 255, 255, 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>
        

        </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}
      >
        <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.3453 10.0167C20.4392 9.67413 23.5614 9.67413 26.6553 10.0167C28.3683 10.2087 29.7503 11.5577 29.9513 13.2767C30.3179 16.4147 30.3179 19.5847 29.9513 22.7227C29.7503 24.4417 28.3683 25.7907 26.6553 25.9827C23.5614 26.3252 20.4392 26.3252 17.3453 25.9827C15.6323 25.7907 14.2503 24.4417 14.0493 22.7227C13.6828 19.585 13.6828 16.4153 14.0493 13.2777C14.151 12.4426 14.5317 11.6662 15.1298 11.0745C15.7278 10.4828 16.5082 10.1104 17.3443 10.0177M22.0003 13.0067C22.1992 13.0067 22.39 13.0857 22.5307 13.2264C22.6713 13.367 22.7503 13.5578 22.7503 13.7567V17.2497H26.2433C26.4422 17.2497 26.633 17.3287 26.7737 17.4694C26.9143 17.61 26.9933 17.8008 26.9933 17.9997C26.9933 18.1986 26.9143 18.3894 26.7737 18.53C26.633 18.6707 26.4422 18.7497 26.2433 18.7497H22.7503V22.2427C22.7503 22.4416 22.6713 22.6324 22.5307 22.773C22.39 22.9137 22.1992 22.9927 22.0003 22.9927C21.8014 22.9927 21.6106 22.9137 21.47 22.773C21.3293 22.6324 21.2503 22.4416 21.2503 22.2427V18.7497H17.7573C17.5584 18.7497 17.3676 18.6707 17.227 18.53C17.0863 18.3894 17.0073 18.1986 17.0073 17.9997C17.0073 17.8008 17.0863 17.61 17.227 17.4694C17.3676 17.3287 17.5584 17.2497 17.7573 17.2497H21.2503V13.7567C21.2503 13.5578 21.3293 13.367 21.47 13.2264C21.6106 13.0857 21.8014 13.0067 22.0003 13.0067Z" 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>

      {/* Add Table Modal */}
      <AddTableModal
        open={isAddTableModalOpen}
        onClose={() => setIsAddTableModalOpen(false)}
        onAddTable={handleAddTable}
        schemas={availableSchemas}
        existingTables={existingTables}
        position="bottom-right"
      />

      {/* Update Table Modal */}
      {/* <UpdateTableModal
        open={isUpdateTableModalOpen}
        onClose={() => {
          setIsUpdateTableModalOpen(false);
          setSelectedTableForUpdate(null);
        }}
        onUpdateTable={handleUpdateTableSubmit}
        tableData={selectedTableForUpdate}
        schemas={availableSchemas}
        existingTables={existingTables}
        position="center"
      /> */}
  
    </div>
  );
};

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

export default ERDiagramCanvas;