src/indexing/extract_db_summary_bigquery.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
BigQuery Database Analyzer
- Lists tables, schema, sample rows, and metadata
- Groups similarly-patterned tables to reduce redundancy
- Heuristically detects potential primary/foreign keys
- Produces a simplified JSON summary and optional enriched output

Requirements:
  pip install google-cloud-bigquery google-api-core
"""

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

try:
    from google.oauth2 import service_account
    from google.cloud import bigquery
except ImportError:
    print("Please install required packages: pip install google-cloud-bigquery google-api-core")
    raise


# ===============================
# Table Pattern Analyzer
# ===============================

class BigQueryTablePatternAnalyzer:
    """Analyzes table name patterns and selects representative tables."""
    def __init__(self, client: Optional["bigquery.Client"] = None):
        self.client = client
        self.table_groups: Dict[str, Any] = {}
        self.representative_tables: Dict[str, Any] = {}
        self.group_metadata: Dict[str, Any] = {}

    def analyze_and_group_tables(
        self,
        tables_dict: Dict[str, List[str]],
        table_column_info: Optional[Dict[str, Dict[str, List[str]]]] = None
    ) -> Dict[str, Any]:
        print("Analyzing table patterns and grouping similar tables...")

        grouped_analysis = {
            'filtered_tables': {},
            'group_info': {},
            'table_column_inf': {}
        }

        for dataset_key, table_list in tables_dict.items():
            if not table_list:
                continue

            print(f"Processing dataset: {dataset_key}")
            project_id, dataset_id = dataset_key.split('.', 1)

            pattern_groups = self._group_tables_by_pattern(table_list, project_id, dataset_id)

            filtered_tables: List[str] = []
            for _, tables_in_group in pattern_groups.items():
                if len(tables_in_group) == 1:
                    table_name = tables_in_group[0]['table_name']
                    filtered_tables.append(table_name)
                    full_table_name = f"{project_id}.{dataset_id}.{table_name}"
                    grouped_analysis['group_info'][full_table_name] = None
                    grouped_analysis['table_column_inf'][full_table_name] = None
                else:
                    representative = self._select_representative_table(
                        tables_in_group, dataset_key, table_column_info
                    )
                    filtered_tables.append(representative['table_name'])

                    full_table_name = f"{project_id}.{dataset_id}.{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

                    column_info_string = self._generate_table_column_info_string(
                        tables_in_group, dataset_key, table_column_info
                    )
                    grouped_analysis['table_column_inf'][full_table_name] = column_info_string

            grouped_analysis['filtered_tables'][dataset_key] = filtered_tables

        print("Table grouping complete.")
        return grouped_analysis

    def create_grouped_structure(
        self,
        original_structure: Dict[str, Any],
        grouped_analysis: Dict[str, Any]
    ) -> Dict[str, Any]:
        print("\n=== CREATING GROUPED STRUCTURE ===")
        grouped_structure = {"datasets": {}}

        for dataset_key, dataset_data in original_structure.get("datasets", {}).items():
            if dataset_key not in grouped_analysis['filtered_tables']:
                continue

            representative_tables = grouped_analysis['filtered_tables'][dataset_key]
            grouped_structure["datasets"][dataset_key] = {
                "project_id": dataset_data["project_id"],
                "dataset_id": dataset_data["dataset_id"],
                "tables": {}
            }

            for table_name in representative_tables:
                if table_name not in dataset_data["tables"]:
                    continue

                original_table_info = dict(dataset_data["tables"][table_name])
                project_id = dataset_data["project_id"]
                dataset_id = dataset_data["dataset_id"]
                full_table_name = f"{project_id}.{dataset_id}.{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

                grouped_structure["datasets"][dataset_key]["tables"][table_name] = original_table_info

        return grouped_structure

    def apply_pkfk_detection(self, structure: Dict[str, Any]) -> Dict[str, Any]:
        print("\n=== APPLYING PK/FK DETECTION ===")
        if self.client is None:
            raise RuntimeError("BigQuery client is not set in BigQueryTablePatternAnalyzer.")

        project_datasets: List[Tuple[str, str]] = []
        for dataset_key in structure.get("datasets", {}).keys():
            project_id, dataset_id = dataset_key.split('.', 1)
            project_datasets.append((project_id, dataset_id))

        key_finder = BigQueryKeyFinder(self.client, project_datasets)

        key_finder.tables = []
        key_finder.table_columns = {}
        key_finder.table_sample_data = {}

        for _, dataset_data in structure.get("datasets", {}).items():
            for table_name, table_data in dataset_data.get("tables", {}).items():
                project_id = dataset_data["project_id"]
                dataset_id = dataset_data["dataset_id"]
                full_table_name = f"{project_id}.{dataset_id}.{table_name}"

                key_finder.tables.append(full_table_name)

                col_names = table_data.get("column_names", [])
                col_types = table_data.get("column_types", [])
                key_finder.table_columns[full_table_name] = list(zip(col_names, col_types))

                sample_rows = table_data.get("sample_rows", [])
                sample_data: Dict[str, List[Any]] = {c: [] for c in col_names}
                for row in sample_rows:
                    if isinstance(row, list):
                        for i, c in enumerate(col_names):
                            if i < len(row):
                                sample_data[c].append(row[i])
                    elif isinstance(row, dict):
                        for c in col_names:
                            if c in row:
                                sample_data[c].append(row[c])
                key_finder.table_sample_data[full_table_name] = sample_data

        print("\nFinding potential primary keys...")
        pk_results = key_finder.find_potential_primary_keys()

        print("\nFinding potential foreign keys...")
        fk_results = key_finder.find_potential_foreign_keys()

        print("PK/FK detection completed:")
        print(f"  Tables with primary keys: {len(pk_results)}")
        print(f"  Tables with foreign keys: {len(fk_results)}")
        print(f"  Total FK relationships: {sum(len(v) for v in fk_results.values())}")

        # Return full table names instead of simple names
        return {
            'tables': [t for t in key_finder.tables],
            'columns': {t: cols for t, cols in key_finder.table_columns.items()},
            'primary_keys': pk_results,
            'foreign_keys': fk_results
        }

    # ---- helpers ----

    def _group_tables_by_pattern(
        self,
        table_list: List[str],
        project_id: str,
        dataset_id: str
    ) -> Dict[str, List[Dict[str, Any]]]:
        patterns: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
        for table_name in table_list:
            pattern_type, base_name, suffix = self._classify_table_pattern(table_name)
            info = {
                'table_name': table_name,
                'base_name': base_name,
                'suffix': suffix,
                'pattern_type': pattern_type
            }
            group_key = f"{pattern_type}_{base_name}" if base_name else f"STANDALONE_{table_name}"
            patterns[group_key].append(info)
        return dict(patterns)

    def _classify_table_pattern(self, table_name: str) -> Tuple[str, str, str]:
        m = re.search(r'^(.+?)(\d{4})$', table_name)
        if m:
            base, year = m.group(1), m.group(2)
            try:
                y = int(year)
                if 1900 <= y <= 2100:
                    return 'YEARLY', base, year
            except Exception:
                pass

        m = re.search(r'^(.+?)_(\d{4})_(.+)$', table_name)
        if m:
            base, year, suffix = m.group(1), m.group(2), m.group(3)
            try:
                y = int(year)
                if 1900 <= y <= 2100:
                    return 'YEARLY_SUFFIX', base, f"{year}_{suffix}"
            except Exception:
                pass

        m = re.search(r'^(.+)_(\d{8})$', table_name)
        if m:
            return 'DATE_STREAMING', m.group(1), m.group(2)

        m = re.search(r'^(.+)_(\d{4}_\d{2}_\d{2})$', table_name)
        if m:
            return 'DATE_STREAMING', m.group(1), m.group(2)

        m = re.search(r'^(.+)_(Q[1-4]_\d{4})$', table_name)
        if m:
            return 'QUARTERLY', m.group(1), m.group(2)

        m = re.search(r'^(.+)_((JAN|FEB|MAR|APR|MAY|JUN|JUL|AUG|SEP|OCT|NOV|DEC)_?\d{4})$', table_name)
        if m:
            return 'MONTHLY', m.group(1), m.group(2)

        m = re.search(r'^([A-Z]+[A-Z])(\d{2,4})$', table_name)
        if m:
            return 'NUMERIC_SUFFIX', m.group(1), m.group(2)

        m = re.search(r'^(.+)_(\d{3,})$', table_name)
        if m:
            return 'SEQUENTIAL', m.group(1), m.group(2)

        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'

        return 'STANDALONE', table_name, ''

    def _select_representative_table(
        self,
        tables_in_group: List[Dict[str, Any]],
        dataset_key: str,
        table_column_info: Optional[Dict[str, Dict[str, List[str]]]] = None
    ) -> Dict[str, Any]:
        scores = []
        for t in tables_in_group:
            score = 0
            table_name = t['table_name']

            column_count = 0
            if table_column_info and dataset_key in table_column_info:
                column_count = len(table_column_info[dataset_key].get(table_name, []))
            score += column_count * 1000

            if t['pattern_type'] in ['YEARLY', 'DATE_STREAMING', 'QUARTERLY', 'MONTHLY'] and t['suffix']:
                year_digits = re.search(r'(\d{4})', t['suffix'])
                if year_digits:
                    try:
                        y = int(year_digits.group(1))
                        if y >= 2020:
                            score += 30
                        elif y >= 2015:
                            score += 20
                        elif y >= 2010:
                            score += 10
                    except Exception:
                        pass

            scores.append({'table_info': t, 'score': score, 'column_count': column_count})

        return max(scores, key=lambda x: x['score'])['table_info']

    def _generate_group_info_string(self, tables_in_group: List[Dict[str, Any]], representative: Dict[str, Any]) -> str:
        names = sorted([t['table_name'] for t in tables_in_group])
        others = [n for n in names if n != representative['table_name']]
        if len(names) <= 10:
            return f"{representative['table_name']} represents a group of tables containing {', '.join(others)}"
        first3, last3 = others[:3], others[-3:]
        return f"{representative['table_name']} represents a group of {len(names)} tables containing {', '.join(first3)}, ..., {', '.join(last3)}"

    def _generate_table_column_info_string(
        self,
        tables_in_group: List[Dict[str, Any]],
        dataset_key: str,
        table_column_info: Optional[Dict[str, Dict[str, List[str]]]] = None
    ) -> str:
        if not table_column_info or dataset_key not in table_column_info:
            return "Column information not available"

        table_column_mappings: List[str] = []
        for t in tables_in_group:
            table_name = t['table_name']
            cols = table_column_info[dataset_key].get(table_name, [])
            for c in cols:
                table_column_mappings.append(f"{table_name}.{c}")

        if len(table_column_mappings) <= 20:
            return f"Group columns are {', '.join(sorted(table_column_mappings))}"
        sorted_map = sorted(table_column_mappings)
        return f"Group columns are {', '.join(sorted_map[:10])}, ..., {', '.join(sorted_map[-10:])} (total: {len(table_column_mappings)} columns)"


# ===============================
# Key Finder
# ===============================

class BigQueryKeyFinder:
    """Heuristic PK/FK detector using table samples and naming patterns."""
    def __init__(self, client: "bigquery.Client", project_datasets: List[Tuple[str, str]]):
        self.client = client
        self.project_datasets = project_datasets
        self.tables: List[str] = []
        self.table_columns: Dict[str, List[Tuple[str, str]]] = {}
        self.table_sample_data: Dict[str, Dict[str, List[Any]]] = {}
        self.primary_keys: Dict[str, Dict[str, Any]] = {}
        self.foreign_keys: Dict[str, List[Dict[str, Any]]] = defaultdict(list)

    def _extract_database_structure(self):
        for project_id, dataset_id in self.project_datasets:
            try:
                dataset_ref = self.client.dataset(dataset_id, project=project_id)
                tables = list(self.client.list_tables(dataset_ref))
                for table in tables:
                    full = f"{project_id}.{dataset_id}.{table.table_id}"
                    self.tables.append(full)

                    table_ref = dataset_ref.table(table.table_id)
                    table_obj = self.client.get_table(table_ref)

                    cols = [(f.name, f.field_type) for f in table_obj.schema]
                    self.table_columns[full] = cols
                    self._get_sample_data(full, table_obj)

            except Exception as e:
                print(f"Error accessing {project_id}.{dataset_id}: {e}")
                continue

        return self.tables, self.table_columns

    def _get_sample_data(self, full_table_name: str, table_obj: "bigquery.table.Table", sample_size: int = 100):
        try:
            q = f"SELECT * FROM `{full_table_name}` LIMIT {sample_size}"
            rows = self.client.query(q).result()
            sample: Dict[str, List[Any]] = {f.name: [] for f in table_obj.schema}
            for r in rows:
                for f in table_obj.schema:
                    sample[f.name].append(r.get(f.name))
            self.table_sample_data[full_table_name] = sample
        except Exception as e:
            print(f"Error getting sample data for {full_table_name}: {e}")
            self.table_sample_data[full_table_name] = {}

    def _get_sample_values(self, table_name: str, column_name: str) -> List[Any]:
        return self.table_sample_data.get(table_name, {}).get(column_name, [])

    @staticmethod
    def _analyze_sample_values(sample_values: List[Any]) -> Dict[str, Union[int, float]]:
        if not sample_values:
            return {"total_count": 0, "null_count": 0, "distinct_count": 0, "uniqueness_ratio": 0.0, "null_ratio": 0.0}
        non_null = [v for v in sample_values if v is not None]
        uniq = set(str(v) for v in non_null)
        total = len(sample_values)
        nulls = total - len(non_null)
        return {
            "total_count": total,
            "null_count": nulls,
            "distinct_count": len(uniq),
            "uniqueness_ratio": (len(uniq) / max(1, len(non_null))) if non_null else 0.0,
            "null_ratio": nulls / max(1, total)
        }
    
    def find_potential_primary_keys(self) -> Dict[str, Dict[str, Any]]:
        print(f"Finding potential primary keys for {len(self.tables)} tables...")
        for table in self.tables:
            pk_candidates: Dict[str, Dict[str, Any]] = {}
            for col_name, data_type in self.table_columns.get(table, []):
                stats = self._analyze_sample_values(self._get_sample_values(table, col_name))
                if stats["uniqueness_ratio"] < 0.9:
                    continue
                if stats["null_ratio"] > 0.1:
                    continue

                score = 0
                if stats["uniqueness_ratio"] == 1.0:
                    score += 30
                elif stats["uniqueness_ratio"] > 0.98:
                    score += 20
                if stats["null_count"] == 0:
                    score += 20
                typ = (data_type or "").upper()
                if typ in ['INTEGER', 'INT64', 'NUMERIC', 'BIGNUMERIC']:
                    score += 15
                elif typ in ['STRING', 'BYTES']:
                    score += 5
                
                # ✅ FIX: Ensure col_name is a string before regex operations
                col_name_str = str(col_name) if col_name is not None else ""
                base = table.split('.')[-1].lower()
                
                patterns = [
                    (r'^id$', 15),
                    (rf'^{base}_id$', 15),
                    (rf'^{base}_key$', 15),
                    (r'^pk_', 15),
                    (r'^key$', 10),
                    (r'^code$', 8),
                    (r'^uuid$', 15),
                    (r'^guid$', 15),
                    (r'id$', 5),
                ]
                for pat, pts in patterns:
                    if re.search(pat, col_name_str, re.IGNORECASE):
                        score += pts
                        break

                if score >= 25:
                    pk_candidates[col_name] = {
                        'score': score,
                        'data_type': data_type,
                        'uniqueness': stats["uniqueness_ratio"],
                        'null_ratio': stats["null_ratio"],
                        'sample_count': stats["total_count"]
                    }

            if pk_candidates:
                sorted_cands = sorted(pk_candidates.items(), key=lambda x: x[1]['score'], reverse=True)
                top_score = sorted_cands[0][1]['score']
                threshold = 0.8 * top_score
                pk_cols = [c for c, info in sorted_cands if info['score'] >= threshold]
                self.primary_keys[table] = {'columns': pk_cols, 'origin': 'potential'}

        return self.primary_keys

    def find_potential_foreign_keys(self) -> Dict[str, List[Dict[str, Any]]]:
        print(f"Finding potential foreign keys for {len(self.tables)} tables...")
        for src in self.tables:
            seen: set = set()
            for src_col, src_type in self.table_columns.get(src, []):
                # ✅ FIX: Ensure src_col is a string
                if src_col is None:
                    continue
                src_col_str = str(src_col)
                
                for ref in self.tables:
                    if ref == src:
                        continue
                    if ref not in self.primary_keys:
                        continue
                    ref_pk_cols = self.primary_keys[ref].get('columns', [])
                    if not ref_pk_cols:
                        continue

                    ref_base = ref.split('.')[-1].lower()
                    for ref_col in ref_pk_cols:
                        # ✅ FIX: Ensure ref_col is a string
                        if ref_col is None:
                            continue
                        ref_col_str = str(ref_col)
                        
                        if (src_col_str, ref, ref_col_str) in seen:
                            continue

                        ref_col_type = None
                        for c, t in self.table_columns.get(ref, []):
                            if c == ref_col:
                                ref_col_type = t
                                break

                        patterns = [
                            rf'^{ref_base}_{ref_col_str}$',
                            rf'^{ref_base}{ref_col_str.capitalize()}$',
                            rf'^{ref_base}_id$',
                            rf'^{ref_col_str}$'
                        ]
                        name_ok = any(re.match(p, src_col_str, re.IGNORECASE) for p in patterns)
                        if name_ok:
                            confidence = "medium" if (ref_col_type and src_type and ref_col_type.upper() == (src_type or '').upper()) else "low"
                            self.foreign_keys[src].append({
                                'from': src_col_str,
                                'to_table': ref,
                                'to_column': ref_col_str,
                                'origin': 'potential',
                                'confidence': confidence
                            })
                            seen.add((src_col_str, ref, ref_col_str))

        return dict(self.foreign_keys)

    def analyze(self) -> Dict[str, Any]:
        print("Analyzing BigQuery database structure")
        self._extract_database_structure()
        self.find_potential_primary_keys()
        self.find_potential_foreign_keys()
        return {
            'tables': [t for t in self.tables],
            'columns': {t: cols for t, cols in self.table_columns.items() },
            'primary_keys': self.primary_keys,
            'foreign_keys': self.foreign_keys
        }


# ===============================
# Database Analyzer
# ===============================

class BigQueryDatabaseAnalyzer:
    """Orchestrates loading datasets, grouping tables, and PK/FK detection."""
    def __init__(self, credentials_path: Optional[str] = None, project_id: Optional[str] = None):
        self.credentials_path = credentials_path
        self.project_id = project_id
        self.client = self._create_client()
        self.pattern_analyzer = BigQueryTablePatternAnalyzer(self.client)

    def _create_client(self) -> "bigquery.Client":
        try:
            if self.credentials_path:
                if not os.path.exists(self.credentials_path):
                    raise FileNotFoundError(
                        f"BigQuery credentials file not found at {self.credentials_path}. "
                        f"Please check the path."
                    )
                print(f"✅ Using service account credentials from {self.credentials_path}")
                creds = service_account.Credentials.from_service_account_file(self.credentials_path)
                return bigquery.Client(credentials=creds, project=self.project_id)

            # If no credentials_path provided at all, fallback to ADC
            print("⚠️ No credentials_path provided, falling back to Application Default Credentials (ADC)")
            return bigquery.Client(project=self.project_id)

        except Exception as e:
            print(f"❌ Error creating BigQuery client: {e}")
            raise



    def list_available_datasets(self, project_ids: Optional[List[str]] = None) -> Dict[str, List[str]]:
        if project_ids is None:
            project_ids = [self.project_id] if self.project_id else ['bigquery-public-data']
        out: Dict[str, List[str]] = {}
        for pid in project_ids:
            try:
                ds = list(self.client.list_datasets(project=pid))
                out[pid] = [d.dataset_id for d in ds]
                print(f"Found {len(out[pid])} datasets in {pid}")
            except Exception as e:
                print(f"Error listing datasets for {pid}: {e}")
                out[pid] = []
        return out

    def expand_datasets(self, project_datasets: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
        """
        Expand (project, dataset_pattern) into concrete (project, dataset_id) pairs.
        Supports:
          - glob patterns: 'austin_*'
          - plain prefix convenience: 'austin' -> matches 'austin' and 'austin_*'
        """
        expanded: List[Tuple[str, str]] = []
        for project_id, pattern in project_datasets:
            try:
                all_ds = [d.dataset_id for d in self.client.list_datasets(project=project_id)]
            except Exception as e:
                print(f"Error listing datasets for {project_id}: {e}")
                continue

            if "*" not in pattern and "?" not in pattern and "[" not in pattern:
                candidates = [ds for ds in all_ds if ds == pattern or ds.startswith(pattern + "_")]
            else:
                candidates = [ds for ds in all_ds if fnmatch.fnmatch(ds, pattern)]

            if not candidates:
                print(f"  No datasets matched '{project_id}.{pattern}'")
            for ds in candidates:
                expanded.append((project_id, ds))
        return expanded

    def load_dataset_structure(self, project_datasets: List[Tuple[str, str]], sample_size: int = 10) -> Dict[str, Any]:
        print("Loading dataset structures from BigQuery...")
        structure = {"datasets": {}}
        for project_id, dataset_id in project_datasets:
            key = f"{project_id}.{dataset_id}"
            print(f"Loading dataset: {key}")
            ds = self._load_dataset_tables(project_id, dataset_id, sample_size=sample_size)
            if ds["tables"]:
                structure["datasets"][key] = ds
                print(f"  Loaded {len(ds['tables'])} tables")
        return structure

    def _load_dataset_tables(self, project_id: str, dataset_id: str, sample_size: int = 10) -> Dict[str, Any]:
        """
        Hybrid loader: combines local schema (DDL.csv) with live BigQuery table info.
        - Reads schema definitions from spider2-lite/resource/databases/bigquery/{dataset_id}/DDL.csv
        - Enriches with live samples & metadata from BigQuery
        """
        out = {"project_id": project_id, "dataset_id": dataset_id, "tables": {}}

        # --- 1. Try to load schema from local DDL.csv (FIXED: handle nested structures) ---
        local_db_root = os.path.join(os.environ.get("BQ_LOCAL_DB_ROOT", "./bigquery_dbs"), dataset_id)
        local_schemas: Dict[str, Dict[str, Any]] = {}
        ddl_paths = []

        # Check for nested structure (e.g., TCGA_bioclin_v0/isb-cgc.GDC_metadata/DDL.csv)
        if os.path.exists(local_db_root):
            try:
                items = os.listdir(local_db_root)
                for item in items:
                    item_path = os.path.join(local_db_root, item)
                    # Look for subdirectories matching project.dataset pattern
                    if os.path.isdir(item_path) and '.' in item:
                        nested_ddl = os.path.join(item_path, "DDL.csv")
                        if os.path.exists(nested_ddl):
                            ddl_paths.append(nested_ddl)
                            print(f"  📁 Found nested DDL: {item}/DDL.csv")
            except Exception as e:
                print(f"⚠️ Error scanning for nested DDL files: {e}")

        # Also check for direct DDL.csv at root level
        ddl_path = os.path.join(local_db_root, "DDL.csv")
        if os.path.exists(ddl_path):
            ddl_paths.insert(0, ddl_path)  # Prioritize root-level DDL

        # Load from all found DDL.csv files
        for ddl_file in ddl_paths:
            try:
                import csv
                with open(ddl_file, newline="", encoding="utf-8") as f:
                    reader = csv.DictReader(f)
                    for row in reader:
                        tname = row.get("table_name")
                        cname = row.get("column_name")
                        ctype = row.get("data_type")
                        if not tname:
                            continue
                        if tname not in local_schemas:
                            local_schemas[tname] = {
                                "table_name": tname,
                                "table_fullname": f"{project_id}.{dataset_id}.{tname}",
                                "column_names": [],
                                "column_types": [],
                                "sample_rows": [],
                                "row_count": None,
                                "size_bytes": None,
                                "created": None,
                                "modified": None
                            }
                        local_schemas[tname]["column_names"].append(cname)
                        local_schemas[tname]["column_types"].append(ctype)
                dir_name = os.path.basename(os.path.dirname(ddl_file))
                print(f"  ✅ Loaded {len([t for t in local_schemas.keys()])} tables from {dir_name}/DDL.csv")
            except Exception as e:
                print(f"⚠️ Error reading DDL.csv from {ddl_file}: {e}")

        # --- 2. Try to load tables from BigQuery API ---
        try:
            dataset_ref = self.client.dataset(dataset_id, project=project_id)
            tables = list(self.client.list_tables(dataset_ref))

            for table in tables:
                tname = table.table_id
                table_ref = dataset_ref.table(tname)

                try:
                    table_obj = self.client.get_table(table_ref)

                    # Prefer local schema if available
                    if tname in local_schemas:
                        col_names = local_schemas[tname]["column_names"]
                        col_types = local_schemas[tname]["column_types"]
                    else:
                        col_names = [f.name for f in table_obj.schema]
                        col_types = [f.field_type for f in table_obj.schema]

                    # Fetch online sample rows
                    samples = self._get_sample_rows(f"{project_id}.{dataset_id}.{tname}", sample_size)

                    out["tables"][tname] = {
                        "table_name": tname,
                        "table_fullname": f"{project_id}.{dataset_id}.{tname}",
                        "column_names": col_names,
                        "column_types": col_types,
                        "sample_rows": samples,
                        "row_count": table_obj.num_rows,
                        "size_bytes": table_obj.num_bytes,
                        "created": table_obj.created.isoformat() if table_obj.created else None,
                        "modified": table_obj.modified.isoformat() if table_obj.modified else None
                    }

                except Exception as e:
                    print(f"⚠️ Error loading table {tname} from BigQuery: {e}")
                    if tname in local_schemas:
                        out["tables"][tname] = local_schemas[tname]

        except Exception as e:
            print(f"❌ Error accessing dataset {project_id}.{dataset_id}: {e}")
            # If BigQuery fails entirely, fallback to local schema only
            if local_schemas:
                out["tables"].update(local_schemas)

        return out


    def _get_sample_rows(self, table_name: str, sample_size: int = 10) -> list:
        project_id = "western-trilogy-473322-n2"  # your project where jobUser permission exists
        query = f"SELECT * FROM `{table_name}` LIMIT {sample_size}"
        try:
            # Run the query using your own project as the job execution project
            rows = self.client.query(query, project=project_id).result()
            return [list(row.values()) for row in rows]
        except Exception as e:
            # Return empty list so PK/FK detection doesn't fail
            print(f"⚠️ Skipping samples for {table_name}: {str(e)[:100]}")
            return []  




    def _extract_table_column_info(self, structure: Dict[str, Any]) -> Dict[str, Dict[str, List[str]]]:
        out: Dict[str, Dict[str, List[str]]] = {}
        for dataset_key, dataset_data in structure.get("datasets", {}).items():
            out[dataset_key] = {}
            for table_name, t in dataset_data.get("tables", {}).items():
                out[dataset_key][table_name] = t.get("column_names", [])
        return out

    def apply_table_grouping(self, structure: Dict[str, Any]) -> Dict[str, Any]:
        print("\n=== APPLYING TABLE GROUPING ===")
        table_column_info = self._extract_table_column_info(structure)
        tables_dict = {dk: list(dd["tables"].keys()) for dk, dd in structure.get("datasets", {}).items()}
        total = sum(len(v) for v in tables_dict.values())
        print(f"Total tables before grouping: {total}")
        grouped = self.pattern_analyzer.analyze_and_group_tables(tables_dict, table_column_info)
        reps = sum(len(v) for v in grouped['filtered_tables'].values())
        red = ((total - reps) / total * 100.0) if total else 0.0
        print(f"Tables after grouping: {reps}")
        print(f"Reduction: {red:.1f}%")
        return grouped

    def run_simple_analysis(self, project_datasets: List[Tuple[str, str]], output_dir: str = "./bigquery_analysis_results", sample_size: int = 10) -> Dict[str, Any]:
        print("STARTING BIGQUERY DATABASE ANALYSIS")
        print("=" * 60)

        # EXPAND nested/pattern datasets here
        project_datasets = self.expand_datasets(project_datasets)

        dataset_names = "_".join([f"{p}_{d}" for p, d in project_datasets])
        print(f"Processing datasets: {project_datasets}")

        try:
            print("Loading dataset structures...")
            structure = self.load_dataset_structure(project_datasets, sample_size=sample_size)
            if not structure["datasets"]:
                print("No datasets loaded successfully")
                return {"status": "error", "error": "No datasets loaded"}

            grouped_analysis = self.apply_table_grouping(structure)
            grouped_structure = self.pattern_analyzer.create_grouped_structure(structure, grouped_analysis)
            key_analysis = self.pattern_analyzer.apply_pkfk_detection(grouped_structure)
            simple_output = self._generate_simple_output(grouped_structure, key_analysis)

            print("\nSaving results...")
            saved = self._save_simple_output(simple_output, dataset_names, output_dir)

            result = {"output": simple_output, "saved_file": saved, "status": "success"}
            self._print_simple_summary(dataset_names, simple_output)
            return result

        except Exception as e:
            print(f"Error processing datasets: {e}")
            import traceback
            traceback.print_exc()
            return {"status": "error", "error": str(e)}

    # ----- helpers for simple output -----

    def _generate_simple_output(
        self,
        structure: Dict[str, Any],
        key_analysis: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Generate simplified JSON output (Snowflake-style structure):
        {
          "tables": {
            "project.dataset.table": {
              "name": "project.dataset.table",
              "columns": [...],
              "samples": {...},
              "foreign_keys": [...],
              "row_count": ...,
              "size_bytes": ...,
              "created": "...",
              "modified": "..."
            }
          },
          "relationships": [...]
        }
        """
        print("\n=== GENERATING SIMPLE OUTPUT ===")
        output: Dict[str, Any] = {"tables": {}, "relationships": []}

        for _, dataset_data in structure.get("datasets", {}).items():
            for table_name, t in dataset_data.get("tables", {}).items():
                project_id = dataset_data["project_id"]
                dataset_id = dataset_data["dataset_id"]
                full_table_name = f"{project_id}.{dataset_id}.{table_name}"

                col_names = t.get("column_names", [])
                col_types = t.get("column_types", [])
                sample_rows = t.get("sample_rows", [])

                # build columns metadata
                cols: List[Dict[str, Any]] = []
                for i, c in enumerate(col_names):
                    ctype = col_types[i] if i < len(col_types) else "UNKNOWN"
                    col_info = {
                        "name": c,
                        "type": ctype,
                        "is_primary_key": False,
                        "is_foreign_key": False
                    }

                    pk_info = key_analysis.get("primary_keys", {}).get(full_table_name)
                    if pk_info and c in pk_info.get("columns", []):
                        col_info["is_primary_key"] = True
                        col_info["pk_origin"] = pk_info.get("origin")

                    fk_list = key_analysis.get("foreign_keys", {}).get(full_table_name, [])
                    for fk in fk_list:
                        if fk.get("from") == c:
                            col_info["is_foreign_key"] = True
                            col_info["fk_origin"] = fk.get("origin")
                            col_info["references_table"] = fk.get("to_table")
                            col_info["references_column"] = fk.get("to_column")
                            break

                    cols.append(col_info)

                # sample rows per column
                samples: Dict[str, List[Any]] = {c: [] for c in col_names}
                for row in sample_rows:
                    if isinstance(row, list):
                        for i, c in enumerate(col_names):
                            if i < len(row):
                                samples[c].append(row[i])
                    elif isinstance(row, dict):
                        for c in col_names:
                            if c in row:
                                samples[c].append(row[c])

                # foreign keys info
                fks: List[Dict[str, Any]] = []
                for fk in key_analysis.get("foreign_keys", {}).get(full_table_name, []):
                    fks.append({
                        "column": fk['from'],
                        "references": {
                            "table": fk['to_table'],
                            "column": fk['to_column']
                        },
                        "fk_origin": fk.get('origin'),
                        "confidence": fk.get('confidence', 'medium')
                    })

                # final entry for this table
                entry = {
                    "name": full_table_name,          # ✅ keep full name inside
                    "columns": cols,
                    "samples": samples,
                    "foreign_keys": fks,
                    "row_count": t.get("row_count"),
                    "size_bytes": t.get("size_bytes"),
                    "created": t.get("created"),
                    "modified": t.get("modified")
                }

                if "table_info" in t:
                    entry["table_info"] = t["table_info"]
                if "table_column_inf" in t:
                    entry["table_column_inf"] = t["table_column_inf"]

                # ✅ ensure consistent key = full table name
                output["tables"][full_table_name] = entry

        # relationships block
        for full_table_name, fk_list in key_analysis.get("foreign_keys", {}).items():
            for fk in fk_list:
                output["relationships"].append({
                    "from_table": full_table_name,
                    "from_column": fk['from'],
                    "to_table": fk['to_table'],
                    "to_column": fk['to_column'],
                    "type": "simple"
                })

        return output


    def _save_simple_output(self, output: Dict[str, Any], dataset_names: str, output_dir: str) -> Optional[str]:
        os.makedirs(output_dir, exist_ok=True)
        path = os.path.join(output_dir, f"{dataset_names}_bigquery_summary.json")
        try:
            with open(path, "w", encoding="utf-8") as f:
                json.dump(output, f, indent=2, ensure_ascii=False, default=str)
            print(f"Simple output saved: {path}")
            return path
        except Exception as e:
            print(f"Error saving simple output: {e}")
            return None

    def _print_simple_summary(self, dataset_names: str, output: Dict[str, Any]):
        print(f"\nSUMMARY FOR {dataset_names}")
        print("-" * 40)
        tables = output.get("tables", {})
        rels = output.get("relationships", [])
        print(f"Tables: {len(tables)}")
        print(f"Relationships: {len(rels)}")

        grouped = sum(1 for t in tables.values() if "table_info" in t)
        print(f"Grouped Tables: {grouped}")
        print(f"Standalone Tables: {len(tables) - grouped}")

        tables_with_pk = 0
        tables_with_fk = 0
        for t in tables.values():
            cols = t.get("columns", [])
            if any(c.get("is_primary_key") for c in cols):
                tables_with_pk += 1
            if any(c.get("is_foreign_key") for c in cols):
                tables_with_fk += 1
        print(f"Tables with PK: {tables_with_pk}")
        print(f"Tables with FK: {tables_with_fk}")

        total_size = sum((t.get("size_bytes") or 0) for t in tables.values())
        total_rows = sum((t.get("row_count") or 0) for t in tables.values())
        print(f"Total estimated rows: {total_rows:,}")
        if total_size > 0:
            print(f"Total estimated size: {total_size / (1024**3):.2f} GB")


def save_bigquery_db_summary(db_summary: Dict[str, Any], output_path: str, indent: int = 2) -> None:
    """Save JSON to disk (handles datetimes via isoformat)."""
    def json_serialize(obj):
        if hasattr(obj, 'isoformat'):
            return obj.isoformat()
        return str(obj)
    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"BigQuery database summary saved to {output_path}")


# ===============================
# CLI
# ===============================

def main():
    # ===============================
    # CLI defaults
    # ===============================

    DEFAULT_CREDENTIALS = os.environ.get("BQ_SERVICE_ACCOUNT_JSON", None)
    DEFAULT_PROJECT = os.environ.get("BQ_PROJECT_ID", None)

    parser = argparse.ArgumentParser(description='BigQuery Database Analysis with Grouping and PK/FK Detection')
    parser.add_argument('--credentials', default=DEFAULT_CREDENTIALS, help='Path to service account JSON file')
    parser.add_argument('--project', default=DEFAULT_PROJECT, help='Default project ID')
    parser.add_argument('--datasets', nargs='+',
                        default=['bigquery-public-data.usa_names', 'bigquery-public-data.census_bureau_usa'],
                        help='Datasets to process (format: project.dataset | supports glob/prefix like austin_*)')
    parser.add_argument('--output-dir', default='./bigquery_results', help='Output directory')
    parser.add_argument('--sample-size', type=int, default=10, help='Sample rows per table')
    args = parser.parse_args()

    # parse incoming dataset strings
    project_datasets: List[Tuple[str, str]] = []
    for ds in args.datasets:
        if '.' in ds:
            p, d = ds.split('.', 1)
            project_datasets.append((p, d))
        else:
            project_datasets.append((args.project or 'bigquery-public-data', ds))

    analyzer = BigQueryDatabaseAnalyzer(args.credentials, args.project)

    # Use the analyzer approach with expansion
    expanded = analyzer.expand_datasets(project_datasets)
    analyzer.run_simple_analysis(project_datasets=expanded,
                                 output_dir=args.output_dir,
                                 sample_size=args.sample_size)

def extract_bigquery_db_summary(
    credentials_path: str,
    project_id: str,
    project_datasets: List[Tuple[str, str]],
    sample_limit: int = 10,
    detect_primary_keys: bool = True,
    detect_foreign_keys: bool = True,
    apply_table_grouping: bool = True
) -> Dict[str, Any]:
    """
    Wrapper to match db_summary_gen.py expectations.
    Uses BigQueryDatabaseAnalyzer to build a summary.
    """
    analyzer = BigQueryDatabaseAnalyzer(credentials_path, project_id)

    # Expand dataset patterns (wildcards, prefixes)
    expanded = analyzer.expand_datasets(project_datasets)

    # Load structure
    structure = analyzer.load_dataset_structure(expanded, sample_size=sample_limit)
    if not structure.get("datasets"):
        print("No datasets loaded successfully")
        return {}

    grouped_analysis = analyzer.apply_table_grouping(structure) if apply_table_grouping else None
    grouped_structure = (
        analyzer.pattern_analyzer.create_grouped_structure(structure, grouped_analysis)
        if grouped_analysis
        else structure
    )

    key_analysis = analyzer.pattern_analyzer.apply_pkfk_detection(grouped_structure) \
        if (detect_primary_keys or detect_foreign_keys) else {}

    return analyzer._generate_simple_output(grouped_structure, key_analysis)


if __name__ == "__main__":
    main()