src/indexing/extract_db_summary_snowflake.py

import json
import os
import datetime
import statistics
from collections import defaultdict
import re
from typing import Dict, Any, List, Tuple, Optional, Union
import argparse
import glob

import snowflake.connector

SNOWFLAKE_USER = os.environ.get("SNOWFLAKE_USER", "")
SNOWFLAKE_PASSWORD = os.environ.get("SNOWFLAKE_PASSWORD", "")
SNOWFLAKE_ACCOUNT = os.environ.get("SNOWFLAKE_ACCOUNT", "")

def connect_snowflake():
    if not SNOWFLAKE_USER or not SNOWFLAKE_PASSWORD or not SNOWFLAKE_ACCOUNT:
        raise ValueError(
            "Snowflake credentials not set. Export SNOWFLAKE_USER, SNOWFLAKE_PASSWORD, SNOWFLAKE_ACCOUNT."
        )
    return snowflake.connector.connect(
        user=SNOWFLAKE_USER,
        password=SNOWFLAKE_PASSWORD,
        account=SNOWFLAKE_ACCOUNT
    )


class LocalTablePatternAnalyzer:
    """
    Analyzes table patterns and groups similar tables for local JSON files
    """
    
    def __init__(self):
        self.table_groups = {}
        self.representative_tables = {}
        self.group_metadata = {}
        
    def analyze_and_group_tables(self, tables_dict: Dict[str, List[str]], table_column_info: Dict[str, Dict[str, List[str]]] = None) -> Dict[str, Any]:
        """
        Analyze table patterns and group similar tables locally
        
        Args:
            tables_dict: Dictionary of schema_key -> list of table names
            table_column_info: Dictionary of schema_key -> table_name -> list of column names
            
        Returns:
            Dictionary containing grouped table analysis
        """
        print("Analyzing table patterns and grouping similar tables...")
        
        grouped_analysis = {
            'filtered_tables': {},
            'group_info': {},
            'table_column_inf': {}
        }
        
        for schema_key, table_list in tables_dict.items():
            if not table_list:
                continue
                
            print(f"Processing schema: {schema_key}")
            db_name, schema_name = schema_key.split('.', 1)
            
            # Group tables by pattern
            pattern_groups = self._group_tables_by_pattern(table_list, db_name, schema_name)
            
            # Select representative tables and generate group info
            filtered_tables = []
            
            for pattern_type, tables_in_group in pattern_groups.items():
                if len(tables_in_group) == 1:
                    # Single table - no grouping needed
                    table_name = tables_in_group[0]['table_name']
                    filtered_tables.append(table_name)
                    
                    full_table_name = f"{db_name}.{schema_name}.{table_name}"
                    grouped_analysis['group_info'][full_table_name] = None
                    grouped_analysis['table_column_inf'][full_table_name] = None
                    
                else:
                    # Multiple tables - select representative and create group info
                    representative = self._select_representative_table(tables_in_group, schema_key, table_column_info)
                    
                    filtered_tables.append(representative['table_name'])
                    
                    # Generate group info string
                    full_table_name = f"{db_name}.{schema_name}.{representative['table_name']}"
                    group_info_string = self._generate_group_info_string(tables_in_group, representative)
                    grouped_analysis['group_info'][full_table_name] = group_info_string
                    
                    # Generate table column info string
                    column_info_string = self._generate_table_column_info_string(tables_in_group, schema_key, table_column_info)
                    grouped_analysis['table_column_inf'][full_table_name] = column_info_string
            
            # Store filtered table list (only representatives)
            grouped_analysis['filtered_tables'][schema_key] = filtered_tables
            
        print(f"Table grouping complete.")
        return grouped_analysis
    
    def _group_tables_by_pattern(self, table_list: List[str], db_name: str, schema_name: str) -> Dict[str, List[Dict]]:
        """Group tables by detected patterns"""
        
        patterns = defaultdict(list)
        
        for table_name in table_list:
            pattern_type, base_name, suffix = self._classify_table_pattern(table_name)
            
            table_info = {
                'table_name': table_name,
                'base_name': base_name,
                'suffix': suffix,
                'pattern_type': pattern_type
            }
            
            # Group key combines pattern type and base name for better grouping
            if base_name:
                group_key = f"{pattern_type}_{base_name}"
            else:
                group_key = f"STANDALONE_{table_name}"
                
            patterns[group_key].append(table_info)
        
        return dict(patterns)
    
    def _classify_table_pattern(self, table_name: str) -> Tuple[str, str, str]:
        """Classify table pattern and return (pattern_type, base_name, identifier)"""
        
        # Year at end pattern (GSOD1929, GSOD2024, etc.)
        year_match = re.search(r'^(.+?)(\d{4})$', table_name)
        if year_match:
            base_name = year_match.group(1)
            year = year_match.group(2)
            if 1900 <= int(year) <= 2100:
                return 'YEARLY', base_name, year
        
        # Year in middle with suffix pattern (BLOCKGROUP_2010_5YR, CBSA_2007_1YR, etc.)
        year_middle_match = re.search(r'^(.+?)_(\d{4})_(.+)$', table_name)
        if year_middle_match:
            base_name = year_middle_match.group(1)
            year = year_middle_match.group(2)
            suffix = year_middle_match.group(3)
            if 1900 <= int(year) <= 2100:
                return 'YEARLY_SUFFIX', base_name, f"{year}_{suffix}"
        
        # Date streaming pattern (YYYYMMDD)
        date_match = re.search(r'^(.+)_(\d{8})$', table_name)
        if date_match:
            return 'DATE_STREAMING', date_match.group(1), date_match.group(2)
        
        # Date pattern with separators (YYYY_MM_DD)
        date_sep_match = re.search(r'^(.+)_(\d{4}_\d{2}_\d{2})$', table_name)
        if date_sep_match:
            return 'DATE_STREAMING', date_sep_match.group(1), date_sep_match.group(2)
        
        # Quarterly pattern (standard format: BASE_Q1_2024)
        quarterly_match = re.search(r'^(.+)_(Q[1-4]_\d{4})$', table_name)
        if quarterly_match:
            return 'QUARTERLY', quarterly_match.group(1), quarterly_match.group(2)
        
        # Quarterly pattern (underscore format: _1990_Q2)
        quarterly_underscore_match = re.search(r'^_(\d{4})_(Q[1-4])$', table_name)
        if quarterly_underscore_match:
            year = quarterly_underscore_match.group(1)
            quarter = quarterly_underscore_match.group(2)
            if 1900 <= int(year) <= 2100:
                return 'QUARTERLY', 'BLS_QCEW', f"{year}_{quarter}"
        
        # Monthly pattern
        monthly_match = re.search(r'^(.+)_((JAN|FEB|MAR|APR|MAY|JUN|JUL|AUG|SEP|OCT|NOV|DEC)_?\d{4})$', table_name)
        if monthly_match:
            return 'MONTHLY', monthly_match.group(1), monthly_match.group(2)
        
        # Numeric suffix without underscore (INDIV04, OPPEXP06, OTH00, etc.)
        numeric_suffix_match = re.search(r'^([A-Z]+[A-Z])(\d{2,4})$', table_name)
        if numeric_suffix_match:
            base_name = numeric_suffix_match.group(1)
            suffix = numeric_suffix_match.group(2)
            return 'NUMERIC_SUFFIX', base_name, suffix
        
        # Sequential numeric pattern with underscore
        seq_match = re.search(r'^(.+)_(\d{3,})$', table_name)
        if seq_match:
            return 'SEQUENTIAL', seq_match.group(1), seq_match.group(2)
        
        # Release/Version pattern with number in middle (REL17_DESCRIPTION, REL18_DESCRIPTION)
        release_match = re.search(r'^([A-Z]+)(\d{2,3})_(.+)$', table_name)
        if release_match:
            base_name = release_match.group(1)
            version = release_match.group(2)
            suffix = release_match.group(3)
            return 'RELEASE_VERSIONED', base_name, f"{version}_{suffix}"
        
        # Version pattern
        version_match = re.search(r'^(.+)_(V\d+|VER_\d+|VERSION_\d+)$', table_name)
        if version_match:
            return 'VERSIONED', version_match.group(1), version_match.group(2)
        
        # Archive pattern
        if re.search(r'_(ARCHIVE|BACKUP|BAK|HIST|HISTORICAL)$', table_name):
            base = re.sub(r'_(ARCHIVE|BACKUP|BAK|HIST|HISTORICAL)$', '', table_name)
            return 'ARCHIVE', base, 'ARCHIVE'
        
        # Default case - standalone table
        return 'STANDALONE', table_name, ''
    
    def _select_representative_table(self, tables_in_group: List[Dict], schema_key: str, table_column_info: Dict[str, Dict[str, List[str]]] = None) -> Dict:
        """Select the best representative table from a group based on column count (highest priority) and other factors"""
        
        table_scores = []
        
        for table_info in tables_in_group:
            score = 0
            table_name = table_info['table_name']
            
            # Get column count for this table
            column_count = 0
            if table_column_info and schema_key in table_column_info and table_name in table_column_info[schema_key]:
                column_count = len(table_column_info[schema_key][table_name])
            
            # Primary scoring: highest column count gets the highest score
            # Use a large multiplier to ensure column count dominates other factors
            score += column_count * 1000
            
            # Secondary scoring: other factors (much lower weight)
            # Prefer more recent tables for temporal patterns
            if table_info['pattern_type'] in ['YEARLY', 'DATE_STREAMING', 'QUARTERLY', 'MONTHLY']:
                if table_info['suffix']:
                    # For yearly patterns, prefer more recent years
                    if table_info['pattern_type'] == 'YEARLY' and table_info['suffix'].isdigit():
                        try:
                            year_val = int(table_info['suffix'])
                            if 2015 <= year_val <= 2022:  # Sweet spot for complete recent data
                                score += 25
                            elif 2010 <= year_val <= 2024:  # Still good recent data
                                score += 15
                            elif year_val >= 2000:  # Decent recent data
                                score += 10
                        except:
                            pass
                    
                    # For date patterns, prefer more recent dates
                    elif table_info['pattern_type'] == 'DATE_STREAMING' and re.match(r'\d{8}', table_info['suffix']):
                        try:
                            date_val = int(table_info['suffix'])
                            score += date_val / 100000  # Recent dates get higher scores
                        except:
                            pass
            
            # Prefer middle values for sequential patterns
            elif table_info['pattern_type'] == 'SEQUENTIAL':
                if table_info['suffix'] and table_info['suffix'].isdigit():
                    seq_val = int(table_info['suffix'])
                    if 50 <= seq_val <= 500:  # Middle range
                        score += 30
            
            table_scores.append({
                'table_info': table_info,
                'score': score,
                'column_count': column_count
            })
        
        # Select table with highest score (which will be dominated by column count)
        best_table = max(table_scores, key=lambda x: x['score'])
        return best_table['table_info']
    
    def _generate_group_info_string(self, tables_in_group: List[Dict], representative: Dict) -> str:
        """Generate a group info string showing all table names that this represents"""
        
        table_names = sorted([t['table_name'] for t in tables_in_group])
        
        if len(table_names) <= 10:
            other_tables = [name for name in table_names if name != representative['table_name']]
            return f"{representative['table_name']} represents a group of tables containing {', '.join(other_tables)}"
        else:
            other_tables = [name for name in table_names if name != representative['table_name']]
            first_few = other_tables[:3]
            last_few = other_tables[-3:]
            return f"{representative['table_name']} represents a group of {len(table_names)} tables containing {', '.join(first_few)}, ..., {', '.join(last_few)}"
    
    def _generate_table_column_info_string(self, tables_in_group: List[Dict], schema_key: str, table_column_info: Dict[str, Dict[str, List[str]]] = None) -> str:
        """Generate a table column info string showing all possible columns in the group"""
        
        if not table_column_info or schema_key not in table_column_info:
            return "Column information not available"
        
        all_columns = set()
        table_column_mappings = []
        
        for table_info in tables_in_group:
            table_name = table_info['table_name']
            if table_name in table_column_info[schema_key]:
                columns = table_column_info[schema_key][table_name]
                for column in columns:
                    all_columns.add(column)
                    table_column_mappings.append(f"{table_name}.{column}")
        
        if len(table_column_mappings) <= 20:
            return f"Group columns are {', '.join(sorted(table_column_mappings))}"
        else:
            sorted_mappings = sorted(table_column_mappings)
            first_few = sorted_mappings[:10]
            last_few = sorted_mappings[-10:]
            return f"Group columns are {', '.join(first_few)}, ..., {', '.join(last_few)} (total: {len(table_column_mappings)} columns)"


class LocalJSONDatabaseKeyFinder:
    """
    Enhanced key finder for local JSON database structures
    """
    
    def __init__(self, database_structure, db_name=None):
        self.database_structure = database_structure
        self.db_name = db_name or "unknown_db"
        self.tables = []
        self.table_columns = {}
        self.table_data = {}
        self.primary_keys = {}
        self.foreign_keys = defaultdict(list)
        
    def _extract_database_structure(self):
        """Extract tables and columns from database structure"""
        
        if "databases" in self.database_structure:
            # Multi-database structure
            for db_name, db_data in self.database_structure["databases"].items():
                if not self.db_name or self.db_name == "unknown_db":
                    self.db_name = db_name
                schemas = db_data.get("schemas", {})
                break  # Take first database
        elif "database" in self.database_structure:
            # Single database structure
            schemas = self.database_structure["database"]["schemas"]
            if not self.db_name or self.db_name == "unknown_db":
                self.db_name = self.database_structure["database"].get("name", "unknown_db")
        else:
            raise ValueError("Unsupported database structure format")
        
        # Extract tables from all schemas
        for schema_name, schema_data in schemas.items():
            if "tables" not in schema_data:
                continue
                
            for table_name, table_info in schema_data["tables"].items():
                full_table_name = f"{schema_name}.{table_name}"
                self.tables.append(full_table_name)
                
                # Extract column information
                columns = []
                column_names = table_info.get("column_names", [])
                column_types = table_info.get("column_types", [])
                
                for i, col_name in enumerate(column_names):
                    col_type = column_types[i] if i < len(column_types) else "UNKNOWN"
                    columns.append((col_name, col_type))
                
                self.table_columns[full_table_name] = columns
                self.table_data[full_table_name] = table_info
        
        return self.tables, self.table_columns

    def _get_sample_values(self, table_name, column_name):
        """Get sample values for a column from table data"""
        if table_name not in self.table_data:
            return []
        
        table_info = self.table_data[table_name]
        sample_rows = table_info.get("sample_rows", [])
        column_names = table_info.get("column_names", [])
        
        if column_name not in column_names:
            return []
        
        col_index = column_names.index(column_name)
        sample_values = []
        
        for row in sample_rows:
            if isinstance(row, list) and col_index < len(row):
                sample_values.append(row[col_index])
            elif isinstance(row, dict) and column_name in row:
                sample_values.append(row[column_name])
        
        return sample_values

    def _analyze_sample_values(self, sample_values):
        """Analyze sample values to determine characteristics"""
        if not sample_values:
            return {
                "total_count": 0,
                "null_count": 0,
                "distinct_count": 0,
                "uniqueness_ratio": 0,
                "null_ratio": 0
            }
        
        non_null_values = [v for v in sample_values if v is not None]
        unique_values = set(str(v) for v in non_null_values)
        
        total_count = len(sample_values)
        null_count = total_count - len(non_null_values)
        distinct_count = len(unique_values)
        
        return {
            "total_count": total_count,
            "null_count": null_count,
            "distinct_count": distinct_count,
            "uniqueness_ratio": distinct_count / max(1, len(non_null_values)) if non_null_values else 0,
            "null_ratio": null_count / max(1, total_count)
        }

    def find_potential_primary_keys(self):
        """Identify columns that are likely to be primary keys"""
        print(f"Finding potential primary keys for {len(self.tables)} tables...")
        
        for table_name in self.tables:
            pk_candidates = {}
            
            for col_name, data_type in self.table_columns[table_name]:
                sample_values = self._get_sample_values(table_name, col_name)
                stats = self._analyze_sample_values(sample_values)
                
                if stats["uniqueness_ratio"] < 0.9:
                    continue
                
                if stats["null_ratio"] > 0.1:
                    continue
                
                pk_score = 0
                
                # Score uniqueness
                if stats["uniqueness_ratio"] == 1.0:
                    pk_score += 30
                elif stats["uniqueness_ratio"] > 0.98:
                    pk_score += 20
                
                # No nulls is good for PKs
                if stats["null_count"] == 0:
                    pk_score += 20
                
                # Data type scoring
                if any(int_type in data_type.upper() for int_type in ['NUMBER', 'INTEGER', 'BIGINT', 'SMALLINT']):
                    pk_score += 15
                elif any(text_type in data_type.upper() for text_type in ['VARCHAR', 'CHAR', 'STRING', 'TEXT']):
                    pk_score += 5
                
                # Check for naming patterns
                table_base_name = table_name.split('.')[-1]
                name_patterns = [
                    (r'^id$', 15),
                    (r'^{}_id$'.format(table_base_name.lower()), 15),
                    (r'^{}_key$'.format(table_base_name.lower()), 15),
                    (r'^pk_', 15),
                    (r'^key$', 10),
                    (r'^code$', 8),
                    (r'^uuid$', 15),
                    (r'^guid$', 15),
                    (r'^serial$', 15),
                    (r'^seq', 10),
                    (r'id$', 5),
                    (r'uuid$', 10),
                    (r'code$', 5),
                    (r'num$', 5),
                    (r'no$', 5),
                    (r'^record', 8),
                    (r'^pid$', 15),
                    (r'^mid$', 15),
                    (r'^uid$', 15),
                    (r'^eid$', 15),
                    (r'^[a-z]+id$', 10),
                    (r'^[a-z]+_id$', 10),
                ]
                
                for pattern, score in name_patterns:
                    if re.search(pattern, col_name, re.IGNORECASE):
                        pk_score += score
                        break
                
                # Check for auto-increment indication in sample values
                try:
                    non_null_values = [v for v in sample_values if v is not None]
                    if non_null_values and all(isinstance(v, (int, float)) for v in non_null_values):
                        sorted_values = sorted([int(v) for v in non_null_values])
                        min_val = sorted_values[0]
                        
                        if min_val in [0, 1] and len(sorted_values) > 1:
                            expected_values = list(range(int(min_val), int(min_val) + len(sorted_values)))
                            if sorted_values == expected_values:
                                pk_score += 15
                except:
                    pass
                
                if pk_score >= 25:
                    pk_candidates[col_name] = {
                        'score': pk_score,
                        'data_type': data_type,
                        'uniqueness': stats["uniqueness_ratio"],
                        'null_ratio': stats["null_ratio"],
                        'sample_count': stats["total_count"]
                    }
            
            # Select the best primary key candidate(s)
            if pk_candidates:
                sorted_candidates = sorted(
                    pk_candidates.items(), 
                    key=lambda x: x[1]['score'], 
                    reverse=True
                )
                
                pk_columns = []
                threshold_score = sorted_candidates[0][1]['score'] * 0.8
                
                for col_name, info in sorted_candidates:
                    if info['score'] >= threshold_score:
                        pk_columns.append(col_name)
                
                self.primary_keys[table_name] = {
                    'columns': pk_columns,
                    'origin': 'potential'
                }
                
        return self.primary_keys

    def find_potential_foreign_keys(self):
        """Identify columns that are likely to be foreign keys"""
        print(f"Finding potential foreign keys for {len(self.tables)} tables...")
        
        for src_table in self.tables:
            processed_relationships = set()
            
            for src_col_name, src_data_type in self.table_columns[src_table]:
                for ref_table in self.tables:
                    if src_table == ref_table:
                        continue
                    
                    if ref_table not in self.primary_keys:
                        continue
                    
                    ref_col_list = self.primary_keys[ref_table].get('columns', [])
                    if not ref_col_list:
                        continue
                    
                    for ref_col in ref_col_list:
                        if (src_col_name, ref_table, ref_col) in processed_relationships:
                            continue
                        
                        # Get reference column data type
                        ref_col_type = None
                        for col_name, col_type in self.table_columns[ref_table]:
                            if col_name == ref_col:
                                ref_col_type = col_type
                                break
                        
                        # Define naming patterns for foreign keys
                        ref_table_base = ref_table.split('.')[-1]
                        fk_patterns = [
                            r'^{}_{}$'.format(ref_table_base.lower(), ref_col.lower()),
                            r'^{}{}$'.format(ref_table_base.lower(), ref_col.capitalize()),
                            r'^{}_id$'.format(ref_table_base.lower()),
                            r'^{}Id$'.format(ref_table_base.lower()),
                            r'^{}_key$'.format(ref_table_base.lower()),
                            r'^fk_{}_'.format(ref_table_base.lower()),
                            r'^{}$'.format(ref_col.lower())
                        ]
                        
                        name_pattern_match = False
                        for pattern in fk_patterns:
                            if re.match(pattern, src_col_name, re.IGNORECASE):
                                name_pattern_match = True
                                break
                        
                        if name_pattern_match:
                            confidence = "medium" if ref_col_type and src_data_type and ref_col_type.upper() == src_data_type.upper() else "low"
                            
                            self.foreign_keys[src_table].append({
                                'from': src_col_name,
                                'to_table': ref_table,
                                'to_column': ref_col,
                                'origin': 'potential',
                                'confidence': confidence
                            })
                            
                            processed_relationships.add((src_col_name, ref_table, ref_col))
                        
                        # Check data value matching
                        else:
                            try:
                                src_sample_values = self._get_sample_values(src_table, src_col_name)
                                ref_sample_values = self._get_sample_values(ref_table, ref_col)
                                
                                if not src_sample_values or not ref_sample_values:
                                    continue
                                
                                src_set = set(str(v) for v in src_sample_values if v is not None)
                                ref_set = set(str(v) for v in ref_sample_values if v is not None)
                                
                                if not src_set:
                                    continue
                                
                                invalid_refs = src_set - ref_set
                                
                                if len(invalid_refs) == 0:
                                    coverage_ratio = len(src_set) / max(1, len(ref_set))
                                    
                                    if (coverage_ratio > 0.01 and 
                                        (not ref_col_type or not src_data_type or 
                                         ref_col_type.upper() == src_data_type.upper())):
                                        
                                        self.foreign_keys[src_table].append({
                                            'from': src_col_name,
                                            'to_table': ref_table,
                                            'to_column': ref_col,
                                            'origin': 'potential',
                                            'confidence': 'high' if coverage_ratio > 0.3 else 'medium',
                                            'evidence': 'data_match',
                                            'coverage_ratio': round(coverage_ratio, 3)
                                        })
                                        
                                        processed_relationships.add((src_col_name, ref_table, ref_col))
                            except Exception as e:
                                continue
        
        return dict(self.foreign_keys)

    def analyze(self):
        """Run the full analysis to find potential primary and foreign keys"""
        print(f"Analyzing database structure for {self.db_name}")
        
        tables, table_columns = self._extract_database_structure()
        print(f"Found {len(tables)} tables: {', '.join([t.split('.')[-1] for t in tables])}")
        
        print("\nFinding potential primary keys...")
        pk_results = self.find_potential_primary_keys()
        
        print("\nFinding potential foreign keys...")
        fk_results = self.find_potential_foreign_keys()
        
        return {
            'tables': [t.split('.')[-1] for t in tables],
            'columns': {t.split('.')[-1]: cols for t, cols in table_columns.items()},
            'primary_keys': {t.split('.')[-1]: pk for t, pk in pk_results.items()},
            'foreign_keys': {t.split('.')[-1]: fk for t, fk in fk_results.items()}
        }


class LocalDatabaseAnalyzer:
    """
    Combined analyzer for local JSON database files with grouping and PK/FK detection
    """
    
    def __init__(self, base_path: str):
        self.base_path = base_path
        self.pattern_analyzer = LocalTablePatternAnalyzer()
        
    def load_database_structure(self, database_name: str) -> Dict[str, Any]:
        """Load database structure from local JSON files"""
        
        db_path = os.path.join(self.base_path, database_name)
        
        if not os.path.exists(db_path):
            available_dbs = [d for d in os.listdir(self.base_path) 
                           if os.path.isdir(os.path.join(self.base_path, d))]
            raise ValueError(f"Database '{database_name}' not found. Available: {available_dbs}")
        
        print(f"Loading database structure from: {db_path}")
        
        structure = {
            "databases": {
                database_name: self._load_database_directory(db_path, database_name)
            }
        }
        
        return structure
    
    def _load_database_directory(self, db_path: str, db_name: str) -> Dict[str, Any]:
        """Load a database directory structure"""
        
        db_structure = {
            "name": db_name,
            "schemas": {}
        }
        
        items = os.listdir(db_path)
        json_files = [f for f in items if f.endswith('.json')]
        subdirs = [d for d in items if os.path.isdir(os.path.join(db_path, d))]
        
        if json_files and not subdirs:
            # Direct table files in database directory
            db_structure["schemas"]["PUBLIC"] = self._load_schema_directory(db_path, "PUBLIC")
        elif subdirs:
            # Schema directories
            for schema_name in subdirs:
                schema_path = os.path.join(db_path, schema_name)
                schema_structure = self._load_schema_directory(schema_path, schema_name)
                if schema_structure["tables"]:
                    db_structure["schemas"][schema_name] = schema_structure
        
        return db_structure
    
    def _load_schema_directory(self, schema_path: str, schema_name: str) -> Dict[str, Any]:
        """Load a schema directory structure"""
        
        schema_structure = {
            "name": schema_name,
            "tables": {}
        }
        
        json_files = glob.glob(os.path.join(schema_path, "*.json"))
        
        for json_file in json_files:
            table_name = os.path.splitext(os.path.basename(json_file))[0]
            
            try:
                with open(json_file, 'r', encoding='utf-8') as f:
                    table_data = json.load(f)
                
                # Ensure required fields
                if "table_name" not in table_data:
                    table_data["table_name"] = table_name
                if "table_fullname" not in table_data:
                    table_data["table_fullname"] = f"{db_name}.{schema_name}.{table_name}"

                
                schema_structure["tables"][table_name] = table_data
                
            except Exception as e:
                print(f"Error loading {json_file}: {e}")
        
        return schema_structure
    
    def _extract_table_column_info(self, structure: Dict[str, Any]) -> Dict[str, Dict[str, List[str]]]:
        """Extract column information for all tables to help with representative selection"""
        
        table_column_info = {}
        
        for db_name, db_data in structure["databases"].items():
            for schema_name, schema_data in db_data["schemas"].items():
                schema_key = f"{db_name}.{schema_name}"
                table_column_info[schema_key] = {}
                
                for table_name, table_data in schema_data["tables"].items():
                    column_names = table_data.get("column_names", [])
                    table_column_info[schema_key][table_name] = column_names
        
        return table_column_info
    
    def apply_table_grouping(self, structure: Dict[str, Any]) -> Dict[str, Any]:
        """Apply table grouping to reduce the number of tables"""
        
        print("\n=== APPLYING TABLE GROUPING ===")
        
        # Extract table column information first
        table_column_info = self._extract_table_column_info(structure)
        
        # Convert structure to the format expected by pattern analyzer
        tables_dict = {}
        
        for db_name, db_data in structure["databases"].items():
            for schema_name, schema_data in db_data["schemas"].items():
                schema_key = f"{db_name}.{schema_name}"
                table_list = list(schema_data["tables"].keys())
                tables_dict[schema_key] = table_list
        
        total_tables = sum(len(table_list) for table_list in tables_dict.values())
        print(f"Total tables before grouping: {total_tables}")
        
        # Apply grouping with column information
        grouped_analysis = self.pattern_analyzer.analyze_and_group_tables(tables_dict, table_column_info)
        
        total_representative_tables = sum(len(filtered_list) for filtered_list in grouped_analysis['filtered_tables'].values())
        reduction_percentage = ((total_tables - total_representative_tables) / total_tables * 100) if total_tables > 0 else 0
        
        print(f"Tables after grouping: {total_representative_tables}")
        print(f"Reduction: {reduction_percentage:.1f}%")
        
        return grouped_analysis
    
    def create_grouped_structure(self, original_structure: Dict[str, Any], 
                                grouped_analysis: Dict[str, Any]) -> Dict[str, Any]:
        """Create structure with only representative tables"""
        
        print("\n=== CREATING GROUPED STRUCTURE ===")
        
        grouped_structure = {
            "databases": {}
        }
        
        for db_name, db_data in original_structure["databases"].items():
            grouped_structure["databases"][db_name] = {
                "name": db_name,
                "schemas": {}
            }
            
            for schema_name, schema_data in db_data["schemas"].items():
                schema_key = f"{db_name}.{schema_name}"
                
                if schema_key not in grouped_analysis['filtered_tables']:
                    continue
                
                representative_tables = grouped_analysis['filtered_tables'][schema_key]
                
                grouped_structure["databases"][db_name]["schemas"][schema_name] = {
                    "name": schema_name,
                    "tables": {}
                }
                
                for table_name in representative_tables:
                    if table_name not in schema_data["tables"]:
                        continue
                    
                    original_table_info = schema_data["tables"][table_name].copy()
                    
                    # Add grouping information
                    full_table_name = f"{db_name}.{schema_name}.{table_name}"
                    group_info = grouped_analysis['group_info'].get(full_table_name)
                    table_column_inf = grouped_analysis['table_column_inf'].get(full_table_name)
                    
                    if group_info:
                        original_table_info["table_info"] = group_info
                        print(f"    🔗 {table_name}: {group_info}")
                    else:
                        print(f"    📋 {table_name}: Standalone table")
                    
                    if table_column_inf:
                        original_table_info["table_column_inf"] = table_column_inf
                        print(f"    📊 {table_name}: {table_column_inf}")
                    
                    grouped_structure["databases"][db_name]["schemas"][schema_name]["tables"][table_name] = original_table_info
        
        return grouped_structure
    
    def apply_pkfk_detection(self, structure: Dict[str, Any]) -> Dict[str, Any]:
        """Apply PK/FK detection to the structure"""
        
        print("\n=== APPLYING PK/FK DETECTION ===")
        
        # Create key finder for the structure
        key_finder = LocalJSONDatabaseKeyFinder(structure)
        
        # Run PK/FK analysis
        key_analysis = key_finder.analyze()
        
        print(f"PK/FK detection completed:")
        print(f"  Tables with primary keys: {len(key_analysis['primary_keys'])}")
        print(f"  Tables with foreign keys: {len(key_analysis['foreign_keys'])}")
        print(f"  Total FK relationships: {sum(len(fks) for fks in key_analysis['foreign_keys'].values())}")
        
        return key_analysis
    
    def generate_simple_output(self, structure: Dict[str, Any], key_analysis: Dict[str, Any]) -> Dict[str, Any]:
        """Generate simple output format matching the sample"""
        
        print("\n=== GENERATING SIMPLE OUTPUT ===")
        
        output = {
            "tables": {},
            "relationships": []
        }
        
        # Process each table
        for db_name, db_data in structure["databases"].items():
            for schema_name, schema_data in db_data["schemas"].items():
                for table_name, table_data in schema_data["tables"].items():
                    
                    # Get column information
                    column_names = table_data.get("column_names", [])
                    column_types = table_data.get("column_types", [])
                    sample_rows = table_data.get("sample_rows", [])
                    
                    # Create columns list
                    columns = []
                    for i, col_name in enumerate(column_names):
                        col_type = column_types[i] if i < len(column_types) else "UNKNOWN"
                        
                        column_info = {
                            "name": col_name,
                            "type": col_type,
                            "is_primary_key": False,
                            "is_foreign_key": False
                        }
                        
                        # Check if this column is a primary key
                        if table_name in key_analysis['primary_keys']:
                            pk_info = key_analysis['primary_keys'][table_name]
                            if col_name in pk_info['columns']:
                                column_info["is_primary_key"] = True
                                column_info["pk_origin"] = pk_info['origin']
                        
                        # Check if this column is a foreign key
                        if table_name in key_analysis['foreign_keys']:
                            for fk in key_analysis['foreign_keys'][table_name]:
                                if fk['from'] == col_name:
                                    column_info["is_foreign_key"] = True
                                    column_info["fk_origin"] = fk['origin']
                                    to_full = fk['to_table']
                                    if to_full.count('.') == 0:
                                        to_full = f"{db_name}.{schema_name}.{to_full}"
                                    elif to_full.count('.') == 1:
                                        to_full = f"{db_name}.{to_full}"
                                    # else already fully qualified
                                    column_info["references_table"] = to_full
                                    column_info["references_column"] = fk['to_column']
                                    break
                        
                        columns.append(column_info)
                    
                    # Create samples dictionary
                    samples = {}
                    for i, col_name in enumerate(column_names):
                        col_samples = []
                        for row in sample_rows:
                            if isinstance(row, list) and i < len(row):
                                col_samples.append(row[i])
                            elif isinstance(row, dict) and col_name in row:
                                col_samples.append(row[col_name])
                        samples[col_name] = col_samples
                    
                    # Store the full table name for this table (we'll need it for FK lookup)
                    full_table_name = f"{db_name}.{schema_name}.{table_name}"
                    
                    # Create foreign keys list for the table
                    foreign_keys = []
                    if table_name in key_analysis['foreign_keys']:
                        for fk in key_analysis['foreign_keys'][table_name]:
                            # Find the actual full table name for the reference
                            ref_table_full = None
                            for table_key in [f"{db_name}.{schema_name}.{table_name}" for db_name, db_data in structure["databases"].items() for schema_name, schema_data in db_data["schemas"].items() for table_name in schema_data["tables"].keys()]:
                                if table_key.endswith(f".{fk['to_table']}"):
                                    ref_table_full = table_key
                                    break
                            
                            # Fallback if not found
                            if not ref_table_full:
                                ref_table_full = fk['to_table']
                            
                            fk_entry = {
                                "column": fk['from'],
                                "references": {
                                    "table": ref_table_full,  # Now uses the full table name
                                    "column": fk['to_column']
                                },
                                "fk_origin": fk['origin'],
                                "confidence": fk.get('confidence', 'medium')
                            }
                            foreign_keys.append(fk_entry)
                    
                    # Build table entry
                    table_entry = {
                        "name": full_table_name,
                        "columns": columns,
                        "samples": samples,
                        "foreign_keys": foreign_keys
                    }
                    
                    # Add table_info if it exists (for grouped tables)
                    if "table_info" in table_data:
                        table_entry["table_info"] = table_data["table_info"]
                    
                    # Add table_column_inf if it exists (for grouped tables)
                    if "table_column_inf" in table_data:
                        table_entry["table_column_inf"] = table_data["table_column_inf"]
                    
                    output["tables"][full_table_name] = table_entry

        
        # Process relationships
        for src_table, fks in key_analysis['foreign_keys'].items():
            for fk in fks:
                # Use table names as they appear in the tables section keys
                # Find the actual full table name from the tables dictionary
                from_full = None
                to_full = None
                
                # Find source table full name
                for table_key in output["tables"].keys():
                    if table_key.endswith(f".{src_table}"):
                        from_full = table_key
                        break
                
                # Find destination table full name  
                for table_key in output["tables"].keys():
                    if table_key.endswith(f".{fk['to_table']}"):
                        to_full = table_key
                        break
                
                # Fallback if not found (shouldn't happen but safety check)
                if not from_full:
                    from_full = src_table
                if not to_full:
                    to_full = fk['to_table']

                relationship = {
                    "from_table": from_full,
                    "from_column": fk['from'],
                    "to_table": to_full,
                    "to_column": fk['to_column'],
                    "type": "simple"
                }
                output["relationships"].append(relationship)
        return output



    
    def save_simple_output(self, output: Dict[str, Any], database_name: str, 
                          output_dir: str = "./analysis_results") -> str:
        """Save simple output to JSON file"""
        
        os.makedirs(output_dir, exist_ok=True)
        
        output_file = os.path.join(output_dir, f"{database_name}_db_summary.json")
        
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(output, f, indent=2, ensure_ascii=False)
            print(f"Simple output saved: {output_file}")
            return output_file
        except Exception as e:
            print(f"Error saving simple output: {e}")
            return None
    
    def run_simple_analysis(self, database_names: Union[str, List[str]], 
                           output_dir: str = "./analysis_results") -> Dict[str, Any]:
        """Run complete analysis and generate simple output format"""
        
        print("STARTING SIMPLE DATABASE ANALYSIS")
        print("=" * 60)
        
        if database_names == "all":
            # Get all available databases
            available_dbs = [d for d in os.listdir(self.base_path) 
                           if os.path.isdir(os.path.join(self.base_path, d))]
            database_names = available_dbs
            print(f"Processing ALL databases: {database_names}")
        elif isinstance(database_names, str):
            database_names = [database_names]
            print(f"Processing specific database: {database_names}")
        else:
            print(f"Processing specific databases: {database_names}")
        
        results = {}
        
        for db_name in database_names:
            print(f"\n{'='*20} PROCESSING {db_name} {'='*20}")
            
            try:
                # Step 1: Load database structure
                print(f"Loading database structure for {db_name}...")
                structure = self.load_database_structure(db_name)
                
                # Step 2: Apply table grouping
                grouped_analysis = self.apply_table_grouping(structure)
                
                # Step 3: Create grouped structure
                grouped_structure = self.create_grouped_structure(structure, grouped_analysis)
                
                # Step 4: Apply PK/FK detection
                key_analysis = self.apply_pkfk_detection(grouped_structure)
                
                # Step 5: Generate simple output
                simple_output = self.generate_simple_output(grouped_structure, key_analysis)
                
                # Step 6: Save simple output
                print(f"\nSaving results for {db_name}...")
                saved_file = self.save_simple_output(simple_output, db_name, output_dir)
                
                results[db_name] = {
                    "output": simple_output,
                    "saved_file": saved_file,
                    "status": "success"
                }
                
                # Print summary for this database
                self._print_simple_summary(db_name, simple_output)
                
            except Exception as e:
                print(f"Error processing database {db_name}: {e}")
                results[db_name] = {
                    "status": "error",
                    "error": str(e)
                }
        
        # Print final summary
        self._print_final_simple_summary(results, output_dir)
        
        return results
    
    def _print_simple_summary(self, db_name: str, output: Dict[str, Any]):
        """Print summary for a single database"""
        
        print(f"\nSUMMARY FOR {db_name}")
        print("-" * 40)
        
        tables = output.get("tables", {})
        relationships = output.get("relationships", [])
        
        print(f"Tables: {len(tables)}")
        print(f"Relationships: {len(relationships)}")
        
        # Count tables with grouping info
        grouped_tables = sum(1 for table in tables.values() if "table_info" in table)
        print(f"Grouped Tables: {grouped_tables}")
        print(f"Standalone Tables: {len(tables) - grouped_tables}")
        
        # Count primary and foreign keys
        tables_with_pk = 0
        tables_with_fk = 0
        
        for table in tables.values():
            has_pk = any(col["is_primary_key"] for col in table["columns"])
            has_fk = any(col["is_foreign_key"] for col in table["columns"])
            
            if has_pk:
                tables_with_pk += 1
            if has_fk:
                tables_with_fk += 1
        
        print(f"Tables with PK: {tables_with_pk}")
        print(f"Tables with FK: {tables_with_fk}")
    
    def _print_final_simple_summary(self, results: Dict[str, Any], output_dir: str):
        """Print final summary of all processed databases"""
        
        print(f"\n{'='*60}")
        print("SIMPLE DATABASE ANALYSIS COMPLETED!")
        print(f"{'='*60}")
        
        successful_dbs = [db for db, result in results.items() if result.get("status") == "success"]
        failed_dbs = [db for db, result in results.items() if result.get("status") == "error"]
        
        print(f"Successfully processed: {len(successful_dbs)} databases")
        if successful_dbs:
            print(f"   {', '.join(successful_dbs)}")
        
        if failed_dbs:
            print(f"Failed to process: {len(failed_dbs)} databases")
            print(f"   {', '.join(failed_dbs)}")
        
        print(f"\nResults saved to: {output_dir}")
        
        # Show sample files created
        if successful_dbs:
            sample_db = successful_dbs[0]
            sample_file = results[sample_db].get("saved_file")
            if sample_file:
                print(f"\nSample file created: {os.path.basename(sample_file)}")


def extract_local_db_summary(base_path: str,
                            database_name: str = None,
                            sample_limit: int = 10,
                            include_samples: bool = True,
                            include_column_names: bool = True,
                            include_data_types: bool = True,
                            detect_primary_keys: bool = True,
                            detect_foreign_keys: bool = True,
                            include_key_confidence: bool = True,
                            include_row_count: bool = False,
                            include_column_count: bool = True,
                            skip_empty_tables: bool = False,
                            include_table_relationships: bool = True,
                            apply_table_grouping: bool = True,
                            max_string_display_length: int = 100):
    """
    Extract comprehensive local database summary including schema, statistics, and sample data.
    
    Args:
        base_path: Base path to database directories
        database_name: Specific database to analyze (None for all accessible)
        sample_limit: Number of sample rows per table
        include_samples: Include sample data
        include_column_names: Include column names
        include_data_types: Include data types
        detect_primary_keys: Detect primary keys
        detect_foreign_keys: Detect foreign keys
        include_key_confidence: Include key detection confidence
        include_row_count: Include row counts
        include_column_count: Include column counts
        skip_empty_tables: Skip empty tables
        include_table_relationships: Include relationship summary
        apply_table_grouping: Apply table grouping to reduce table count
        max_string_display_length: String truncation length
        
    Returns:
        Dictionary containing database summary in simple format
    """
    
    try:
        # Create analyzer
        analyzer = LocalDatabaseAnalyzer(base_path)
        
        # Determine which databases to process
        if database_name:
            databases_to_process = [database_name]
        else:
            # Get all available databases
            available_dbs = [d for d in os.listdir(base_path) 
                           if os.path.isdir(os.path.join(base_path, d))]
            databases_to_process = available_dbs
        
        all_results = {}
        
        for db_name in databases_to_process:
            print(f"Processing database: {db_name}")
            
            # Step 1: Load database structure
            structure = analyzer.load_database_structure(db_name)
            
            # Step 2: Apply table grouping if requested
            if apply_table_grouping:
                grouped_analysis = analyzer.apply_table_grouping(structure)
                grouped_structure = analyzer.create_grouped_structure(structure, grouped_analysis)
            else:
                # Create a dummy grouped_analysis but still normalize the structure
                grouped_analysis = {'filtered_tables': {}, 'group_info': {}, 'table_column_inf': {}}
                # Populate filtered_tables with all tables (no filtering)
                for db_name, db_data in structure["databases"].items():
                    for schema_name, schema_data in db_data["schemas"].items():
                        schema_key = f"{db_name}.{schema_name}"
                        grouped_analysis['filtered_tables'][schema_key] = list(schema_data["tables"].keys())
                # Use create_grouped_structure to normalize even without grouping
                grouped_structure = analyzer.create_grouped_structure(structure, grouped_analysis)
            # Step 3: Apply PK/FK detection
            key_analysis = analyzer.apply_pkfk_detection(grouped_structure)
            
            # Step 4: Generate enhanced output with additional statistics
            db_summary = _generate_enhanced_output(
                grouped_structure, 
                key_analysis, 
                grouped_analysis,
                sample_limit=sample_limit,
                include_samples=include_samples,
                include_column_names=include_column_names,
                include_data_types=include_data_types,
                detect_primary_keys=detect_primary_keys,
                detect_foreign_keys=detect_foreign_keys,
                include_key_confidence=include_key_confidence,
                include_row_count=include_row_count,
                include_column_count=include_column_count,
                skip_empty_tables=skip_empty_tables,
                include_table_relationships=include_table_relationships,
                max_string_display_length=max_string_display_length
            )
            
            all_results[db_name] = db_summary
        
        # If only one database, return it directly, otherwise return all
        if len(all_results) == 1:
            return list(all_results.values())[0]
        else:
            return all_results
            
    except Exception as e:
        print(f"Error during analysis: {e}")
        import traceback
        traceback.print_exc()
        return None


def _generate_enhanced_output(structure: Dict[str, Any], 
                            key_analysis: Dict[str, Any],
                            grouped_analysis: Dict[str, Any],
                            sample_limit: int = 10,
                            include_samples: bool = True,
                            include_column_names: bool = True,
                            include_data_types: bool = True,
                            detect_primary_keys: bool = True,
                            detect_foreign_keys: bool = True,
                            include_key_confidence: bool = True,
                            include_row_count: bool = False,
                            include_column_count: bool = True,
                            skip_empty_tables: bool = False,
                            include_table_relationships: bool = True,
                            max_string_display_length: int = 100) -> Dict[str, Any]:
    """Generate enhanced output with additional statistics similar to Snowflake function"""
    
    output = {
        "tables": {},
        "relationships": []
    }
    
    # Process each table
    for db_name, db_data in structure["databases"].items():
        for schema_name, schema_data in db_data["schemas"].items():
            for table_name, table_data in schema_data["tables"].items():
                
                # Define full_table_name at the beginning of the loop
                full_table_name = f"{db_name}.{schema_name}.{table_name}"
                
                # Skip empty tables if requested
                if skip_empty_tables:
                    sample_rows = table_data.get("sample_rows", [])
                    if not sample_rows:
                        continue
                
                # Get column information
                column_names = table_data.get("column_names", [])
                column_types = table_data.get("column_types", [])
                sample_rows = table_data.get("sample_rows", [])
                
                # Create columns list
                columns = []
                for i, col_name in enumerate(column_names):
                    if not include_column_names and not include_data_types:
                        continue
                        
                    col_type = column_types[i] if i < len(column_types) else "UNKNOWN"
                    
                    column_info = {}
                    
                    if include_column_names:
                        column_info["name"] = col_name
                    
                    if include_data_types:
                        column_info["type"] = col_type
                    
                    # Initialize key flags
                    if detect_primary_keys:
                        column_info["is_primary_key"] = False
                    
                    if detect_foreign_keys:
                        column_info["is_foreign_key"] = False
                    
                    # Check if this column is a primary key
                    if detect_primary_keys and table_name in key_analysis['primary_keys']:
                        pk_info = key_analysis['primary_keys'][table_name]
                        if col_name in pk_info['columns']:
                            column_info["is_primary_key"] = True
                            if include_key_confidence:
                                column_info["pk_origin"] = pk_info['origin']
                    
                    # Check if this column is a foreign key
                    if detect_foreign_keys and table_name in key_analysis['foreign_keys']:
                        for fk in key_analysis['foreign_keys'][table_name]:
                            if fk['from'] == col_name:
                                column_info["is_foreign_key"] = True
                                if include_key_confidence:
                                    column_info["fk_origin"] = fk['origin']

                                    # Find the actual full table name for the reference
                                    to_full = None
                                    for table_key in [f"{db_name}.{schema_name}.{table_name}" for db_name, db_data in structure["databases"].items() for schema_name, schema_data in db_data["schemas"].items() for table_name in schema_data["tables"].keys()]:
                                        if table_key.endswith(f".{fk['to_table']}"):
                                            to_full = table_key
                                            break
                                    
                                    # Fallback if not found
                                    if not to_full:
                                        to_full = fk['to_table']
                                    
                                    column_info["references_table"] = to_full
                                    column_info["references_column"] = fk['to_column']
                                break

                    columns.append(column_info)

                
                # Create samples dictionary
                samples = {}
                if include_samples:
                    for i, col_name in enumerate(column_names):
                        col_samples = []
                        sample_count = 0
                        
                        for row in sample_rows:
                            if sample_count >= sample_limit:
                                break
                                
                            value = None
                            if isinstance(row, list) and i < len(row):
                                value = row[i]
                            elif isinstance(row, dict) and col_name in row:
                                value = row[col_name]
                            
                            # Handle string truncation
                            if isinstance(value, str) and len(value) > max_string_display_length:
                                value = value[:max_string_display_length-3] + "..."
                            
                            col_samples.append(value)
                            sample_count += 1
                        
                        samples[col_name] = col_samples

                # Create foreign keys list for the table
                foreign_keys = []
                if detect_foreign_keys and table_name in key_analysis['foreign_keys']:
                    for fk in key_analysis['foreign_keys'][table_name]:
                        # Find the actual full table name for the reference
                        ref_table_full = None
                        for table_key in [f"{db_name}.{schema_name}.{table_name}" for db_name, db_data in structure["databases"].items() for schema_name, schema_data in db_data["schemas"].items() for table_name in schema_data["tables"].keys()]:
                            if table_key.endswith(f".{fk['to_table']}"):
                                ref_table_full = table_key
                                break
                        
                        # Fallback if not found
                        if not ref_table_full:
                            ref_table_full = fk['to_table']
                        
                        fk_entry = {
                            "column": fk['from'],
                            "references": {
                                "table": ref_table_full,  # Now uses the full table name
                                "column": fk['to_column']
                            }
                        }
                        
                        if include_key_confidence:
                            fk_entry["fk_origin"] = fk['origin']
                            fk_entry["confidence"] = fk.get('confidence', 'medium')
                        
                        foreign_keys.append(fk_entry)
                
                # Build table entry (full_table_name is already defined above)
                table_entry = {
                    "name": full_table_name
                }
                
                if columns:
                    table_entry["columns"] = columns
                
                if include_samples and samples:
                    table_entry["samples"] = samples
                
                if foreign_keys:
                    table_entry["foreign_keys"] = foreign_keys
                
                # Add column count if requested
                if include_column_count:
                    table_entry["column_count"] = len(column_names)
                
                # Add row count if requested (estimate from samples)
                if include_row_count:
                    table_entry["row_count"] = len(sample_rows)  # This is just sample count, not actual row count
                
                # Add table_info if it exists (for grouped tables)
                group_info = grouped_analysis.get('group_info', {}).get(full_table_name)
                if group_info:
                    table_entry["table_info"] = group_info
                
                # Add table_column_inf if it exists (for grouped tables)
                table_column_inf = grouped_analysis.get('table_column_inf', {}).get(full_table_name)
                if table_column_inf:
                    table_entry["table_column_inf"] = table_column_inf
                
                # Store the table entry
                output["tables"][full_table_name] = table_entry

    
    # Process relationships
    if include_table_relationships:
        for src_table, fks in key_analysis['foreign_keys'].items():
            for fk in fks:
                # Use table names as they appear in the tables section keys
                # Find the actual full table name from the tables dictionary
                from_full = None
                to_full = None
                
                # Find source table full name
                for table_key in output["tables"].keys():
                    if table_key.endswith(f".{src_table}"):
                        from_full = table_key
                        break
                
                # Find destination table full name  
                for table_key in output["tables"].keys():
                    if table_key.endswith(f".{fk['to_table']}"):
                        to_full = table_key
                        break
                
                # Fallback if not found (shouldn't happen but safety check)
                if not from_full:
                    from_full = src_table
                if not to_full:
                    to_full = fk['to_table']

                relationship = {
                    "from_table": from_full,
                    "from_column": fk['from'],
                    "to_table": to_full,
                    "to_column": fk['to_column'],
                    "type": "simple"
                }
                output["relationships"].append(relationship)

    return output



def save_local_db_summary(db_summary, output_path, indent=2):
    """
    Save local database summary to JSON file
    """
    def json_serialize(obj):
        if hasattr(obj, 'isoformat'):  # Handle datetime objects
            return obj.isoformat()
        return str(obj)
    
    # Create output directory if it doesn't exist
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(db_summary, f, default=json_serialize, indent=indent, ensure_ascii=False)
    
    print(f"Local database summary saved to {output_path}")


def main():
    """Main function with command line interface"""
    
    # Simple configuration
    BASE_PATH = os.environ.get("SNOWFLAKE_LOCAL_DB_ROOT", "./snowflake_dbs")
    DATABASES_TO_PROCESS = ['NOAA_GSOD'] #'all' #or ['AUSTIN'] for specific database
    
    parser = argparse.ArgumentParser(description='Simple Local Database Analysis with Grouping and PK/FK Detection')
    
    parser.add_argument('--base-path', help='Base path to database directories', default=BASE_PATH)
    parser.add_argument('--databases', nargs='+', help='Specific databases to process', default=DATABASES_TO_PROCESS)
    parser.add_argument('--output-dir', help='Output directory for results', default='./results_snow_d_v2')
    parser.add_argument('--use-extract-function', action='store_true', help='Use extract_local_db_summary function instead')
    
    args = parser.parse_args()
    
    try:
        if args.use_extract_function:
            # Use the new extract function (similar to Snowflake version)
            if len(args.databases) == 1 and args.databases[0] != 'all':
                db_summary = extract_local_db_summary(
                    base_path=args.base_path,
                    database_name=args.databases[0],
                    sample_limit=10,
                    include_samples=True,
                    detect_primary_keys=True,
                    detect_foreign_keys=True,
                    apply_table_grouping=True
                )
                
                # Save the result
                output_file = os.path.join(args.output_dir, f"{args.databases[0]}_extract_summary.json")
                save_local_db_summary(db_summary, output_file)
                
                return db_summary
            else:
                print("Extract function works best with single database. Use --databases DBNAME")
                return None
        else:
            # Use the original analyzer approach
            analyzer = LocalDatabaseAnalyzer(args.base_path)
            
            # Run simple analysis
            results = analyzer.run_simple_analysis(
                database_names=args.databases,
                output_dir=args.output_dir
            )
            
            return results
        
    except Exception as e:
        print(f"Error during analysis: {e}")
        import traceback
        traceback.print_exc()
        return None


if __name__ == "__main__":
    main()