import logging
import math
from decimal import Decimal
from typing import Dict, Any, List, Optional
import reversion

from sqlalchemy import MetaData, Table, select, func, text, inspect, case
from sqlalchemy.sql.sqltypes import (
    Integer, Float, Numeric, BigInteger, SmallInteger, DECIMAL,
    String, Text, Enum, DateTime, Date, TIMESTAMP
)

logger = logging.getLogger(__name__)


def safe_float(val):
    if isinstance(val, Decimal):
        return float(val)
    return val


def get_column_stats(conn, table_inspector, table_name: str, column_name: str, cardinality_limit: int = 20) -> Dict[str, Any]:
    logger.debug(f"Analyzing column: {table_name}.{column_name}")
    stats = {}

    try:
        if column_name not in table_inspector.columns:
            logger.error(f"Column '{column_name}' not found in table '{table_name}'")
            return {}   
        col = table_inspector.c[column_name]
        dialect_name = conn.dialect.name.lower()

        try:
            basic_query = select(
                func.count().label("row_count"),
                func.sum(case((col.is_(None), 1), else_=0)).label("null_count"),
                func.count(func.distinct(col)).label("cardinality")
            ).select_from(table_inspector)
            row_count, null_count, cardinality = conn.execute(basic_query).fetchone()
            stats.update({
                "row_count": row_count,
                "null_count": null_count,
                "cardinality": cardinality,
                "null_percentage": round((null_count / row_count) * 100, 2) if row_count else 0
            })
        except Exception as e:
            logger.warning(f"Basic stats failed for {table_name}.{column_name}: {e}")

        try:
            inspector = inspect(conn)
            indexed_columns = {c for idx in inspector.get_indexes(table_name) for c in idx["column_names"]}
            stats["is_indexed"] = column_name in indexed_columns
        except Exception as e:
            logger.warning(f"Index check failed for {table_name}: {e}")
            stats["is_indexed"] = False

        is_date_type = isinstance(col.type, (DateTime, Date, TIMESTAMP))

        if isinstance(col.type, (Integer, Float, Numeric, BigInteger, SmallInteger, DECIMAL)) and not is_date_type:
            try:
                basic_numeric_query = select(
                    func.avg(col).label("mean"),
                    func.min(col).label("min"),
                    func.max(col).label("max")
                ).where(col.isnot(None)).select_from(table_inspector)

                result = conn.execute(basic_numeric_query).fetchone()
                mean, min_val, max_val = result

                stats.update({
                    "mean": safe_float(mean),
                    "min": safe_float(min_val),
                    "max": safe_float(max_val),
                    "range": (safe_float(max_val) - safe_float(min_val)) if (min_val is not None and max_val is not None) else None
                })
            except Exception as e:
                logger.warning(f"Basic numeric stats failed for {table_name}.{column_name}: {e}")

            try:
                variance_query = select(
                    (func.avg(col * col) - func.avg(col) * func.avg(col)).label("variance")
                ).where(col.isnot(None)).select_from(table_inspector)

                variance = conn.execute(variance_query).scalar()
                variance = safe_float(variance)

                std_dev = math.sqrt(variance) if (variance is not None and variance >= 0) else None
                stats["std_dev"] = std_dev
            except Exception as e:
                logger.warning(f"Variance/StdDev stats failed for {table_name}.{column_name}: {e}")

        elif isinstance(col.type, (String, Text, Enum)):
            try:
                if stats.get("cardinality", 0) <= cardinality_limit:
                    unique_query = select(col, func.count().label("cnt")).group_by(col).order_by(func.count().desc()).limit(cardinality_limit)
                    rows = conn.execute(unique_query).fetchall()
                    stats["unique_values"] = [{"value": str(row[0]), "count": row[1]} for row in rows]
                else:
                    top_query = select(col, func.count().label("cnt")).group_by(col).order_by(func.count().desc()).limit(5)
                    rows = conn.execute(top_query).fetchall()
                    stats["top_values"] = [{"value": str(row[0]), "count": row[1]} for row in rows]
            except Exception as e:
                logger.warning(f"String frequency analysis failed for {table_name}.{column_name}: {e}")

            try:
                length_func = func.length
                length_query = select(
                    func.min(length_func(col)).label("min_length"),
                    func.max(length_func(col)).label("max_length")
                ).select_from(table_inspector)
                min_length, max_length = conn.execute(length_query).fetchone()
                stats["min_length"] = min_length
                stats["max_length"] = max_length
            except Exception as e:
                logger.warning(f"String length stats failed for {table_name}.{column_name}: {e}")

        elif isinstance(col.type, (DateTime, Date, TIMESTAMP)):
            try:
                dt_query = select(
                    func.min(col).label("min_date"),
                    func.max(col).label("max_date")
                ).where(col.isnot(None)).select_from(table_inspector)
                min_date, max_date = conn.execute(dt_query).fetchone()
                stats["min_date"] = str(min_date) if min_date else None
                stats["max_date"] = str(max_date) if max_date else None
                stats["date_range_days"] = (max_date - min_date).days if (min_date and max_date) else None
            except Exception as e:
                logger.warning(f"Date range stats failed for {table_name}.{column_name}: {e}")

    except Exception as top_level_err:
        logger.error(f"Top-level failure while processing {table_name}.{column_name}: {top_level_err}")
        try:
            conn.rollback()
        except Exception:
            pass
        return {}

    return stats


def get_sample_rows(conn, table_inspector, n: int = 10) -> List[List[Any]]:
    try:
        latest_column = None
        for col in table_inspector.columns:
            if col.primary_key:
                latest_column = col
                break
            if isinstance(col.type, (DateTime, Date, TIMESTAMP)):
                latest_column = col

        if latest_column is not None:
            query = select(table_inspector).order_by(latest_column.desc()).limit(n)
        else:
            query = select(table_inspector).limit(n)

        result = conn.execute(query).fetchall()
        return [list(row) for row in result]

    except Exception as e:
        logger.error(f"Error fetching sample rows for table {table_inspector.name}: {e}")
        try:
            conn.rollback()
        except Exception:
            pass
        return []


def get_table_info(datasource, table_name: str, sample_rows_count: int = 10) -> Dict[str, Any]:
    from terno_dbi.connectors import ConnectorFactory
    from terno_dbi.core import models
    connector = ConnectorFactory.create_connector(
        datasource.type,
        datasource.connection_str,
        credentials=datasource.connection_json
    )

    result = {
        "table_name": table_name,
        "datasource_name": datasource.display_name,
        "datasource_type": datasource.type,
        "columns": [],
        "sample_rows": [],
        "column_names": [],
        "relationships": []
    }

    try:
        with connector.get_connection() as conn:
            metadata = MetaData()
            table_inspector = Table(table_name, metadata, autoload_with=conn.engine)
            for col in table_inspector.columns:
                col_info = {
                    "name": col.name,
                    "type": str(col.type),
                    "nullable": col.nullable,
                    "primary_key": col.primary_key,
                    "stats": get_column_stats(conn, table_inspector, table_name, col.name)
                }
                result["columns"].append(col_info)
                result["column_names"].append(col.name)
            sample = get_sample_rows(conn, table_inspector, sample_rows_count)
            result["sample_rows"] = [[str(v) if v is not None else None for v in row] for row in sample]
            try:
                inspector = inspect(conn)
                fks = inspector.get_foreign_keys(table_name)
                result["relationships"] = [
                    {
                        "column": fk["constrained_columns"],
                        "references_table": fk["referred_table"],
                        "references_column": fk["referred_columns"]
                    }
                    for fk in fks
                ]
            except Exception as e:
                logger.warning(f"Could not get foreign keys for {table_name}: {e}")

        try:
            table_model = models.Table.objects.get(name=table_name, data_source=datasource)
            result["existing_description"] = table_model.description
            result["existing_public_name"] = table_model.public_name
            result["table_id"] = table_model.id

            for col_info in result["columns"]:
                try:
                    col_model = models.TableColumn.objects.get(table=table_model, name=col_info["name"])
                    col_info["existing_description"] = col_model.description
                    col_info["existing_public_name"] = col_model.public_name
                    col_info["column_id"] = col_model.id
                except models.TableColumn.DoesNotExist:
                    pass
        except models.Table.DoesNotExist:
            pass

    except Exception as e:
        logger.exception(f"Error getting table info for {table_name}: {e}")
        result["error"] = str(e)

    return result


def get_datasource_tables_info(datasource_id: int, table_names: Optional[List[str]] = None) -> Dict[str, Any]:
    from terno_dbi.core import models

    try:
        datasource = models.DataSource.objects.get(id=datasource_id, enabled=True)
    except models.DataSource.DoesNotExist:
        return {"error": f"Datasource {datasource_id} not found or not enabled"}

    if table_names:
        tables = models.Table.objects.filter(data_source=datasource, name__in=table_names)
    else:
        tables = models.Table.objects.filter(data_source=datasource)

    result = {
        "datasource_id": datasource_id,
        "datasource_name": datasource.display_name,
        "datasource_type": datasource.type,
        "tables_count": tables.count(),
        "tables": []
    }

    for table in tables:
        table_info = get_table_info(datasource, table.name)
        result["tables"].append(table_info)

    return result


SYSTEM_SCHEMAS = {'INFORMATION_SCHEMA', 'information_schema', 'pg_catalog', 'pg_toast'}


def build_row_counts_lookup(raw_counts: Dict[str, int]) -> Dict[str, int]:
    """
    Build a flexible lookup map from connector row-count results.

    Connectors may return keys as unqualified ("orders") or qualified
    ("public.orders"), and in varying cases.  This function normalises
    every key to lowercase *and* stores an extra entry for the
    unqualified (last-dot-segment) form so both look-ups succeed.
    """
    lookup: Dict[str, int] = {}
    for key, count in raw_counts.items():
        lower_key = key.lower()
        lookup[lower_key] = count
        base = lower_key.rsplit('.', 1)[-1]
        if base != lower_key:
            # Only overwrite if not already present (first wins)
            lookup.setdefault(base, count)
    return lookup


def resolve_row_count(
    table_name: str, counts_map: Dict[str, int]
) -> int | None:
    """
    Try to find a match for *table_name* inside *counts_map*.

    Attempts (in order):
    1. Exact lowercased match  ("public.orders" -> "public.orders")
    2. Unqualified match       ("public.orders" -> "orders")
    """
    lower = table_name.lower()
    if lower in counts_map:
        return counts_map[lower]
    base = lower.rsplit('.', 1)[-1]
    if base in counts_map:
        return counts_map[base]
    return None


def _sync_from_information_schema(connector, datasource, result, overwrite=False):
    """
    Fallback sync using INFORMATION_SCHEMA when SQLAlchemy/SQLShield reflection fails.

    This is useful for databases where only views exist (no tables), or when
    views are stored in a way that SQLAlchemy's reflection doesn't pick them up.

    Works for Snowflake, BigQuery, PostgreSQL, MySQL, etc.

    Args:
        connector: Database connector instance
        datasource: DataSource model instance
        result: Result dict to update (mutated in place)
        overwrite: If True, update existing table/column metadata

    Returns:
        Number of tables/views discovered
    """
    from terno_dbi.core import models
    from sqlalchemy import text

    tables_discovered = 0

    target_schema = None
    conn_str = datasource.connection_str

    if 'snowflake://' in conn_str:
        parts = conn_str.split('/')
        if len(parts) >= 5:
            target_schema = parts[-1].split('?')[0]
    logger.info(f"Syncing from INFORMATION_SCHEMA (target_schema={target_schema})")

    try:
        with connector.get_connection() as conn:
            if target_schema and target_schema.upper() not in SYSTEM_SCHEMAS:
                query = text("""
                    SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, DATA_TYPE, ORDINAL_POSITION
                    FROM INFORMATION_SCHEMA.COLUMNS
                    WHERE UPPER(TABLE_SCHEMA) = UPPER(:schema)
                    ORDER BY TABLE_SCHEMA, TABLE_NAME, ORDINAL_POSITION
                """)
                rows = conn.execute(query, {"schema": target_schema}).fetchall()
            else:
                query = text("""
                    SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, DATA_TYPE, ORDINAL_POSITION
                    FROM INFORMATION_SCHEMA.COLUMNS
                    WHERE UPPER(TABLE_SCHEMA) NOT IN ('INFORMATION_SCHEMA', 'PG_CATALOG', 'PG_TOAST')
                    ORDER BY TABLE_SCHEMA, TABLE_NAME, ORDINAL_POSITION
                """)
                rows = conn.execute(query).fetchall()     
            logger.info(f"INFORMATION_SCHEMA returned {len(rows)} column rows")

            tables_dict = {}
            found_table_names = set()

            for row in rows:
                schema_name, table_name, column_name, data_type, _ = row
                key = (schema_name, table_name)
                if key not in tables_dict:
                    tables_dict[key] = []
                tables_dict[key].append((column_name, data_type))

            logger.info(f"Found {len(tables_dict)} tables/views in INFORMATION_SCHEMA")

            for (schema_name, table_name), columns in tables_dict.items():
                try:
                    if target_schema and schema_name.upper() == target_schema.upper():
                        full_table_name = table_name
                    if target_schema and schema_name.upper() == target_schema.upper():
                        full_table_name = table_name
                    else:
                        full_table_name = f"{schema_name}.{table_name}"

                    found_table_names.add(full_table_name)

                    existing_table = models.Table.objects.filter(
                        data_source=datasource,
                        name=full_table_name
                    ).first()

                    if existing_table and not overwrite:
                        result["tables_skipped"] += 1
                        result["tables"].append({
                            "name": full_table_name,
                            "status": "skipped",
                            "id": existing_table.id
                        })
                        continue

                    if existing_table:
                        table_model = existing_table
                        result["tables_updated"] += 1
                    else:
                        table_model = models.Table.objects.create(
                            data_source=datasource,
                            name=full_table_name,
                            public_name=full_table_name,
                        )
                        result["tables_created"] += 1
                        tables_discovered += 1

                        tables_discovered += 1

                    columns_count = 0
                    found_column_names = set()

                    for col_name, data_type in columns:
                        existing_col = models.TableColumn.objects.filter(
                            table=table_model,
                            name=col_name
                        ).first()

                        if existing_col:
                            if overwrite:
                                existing_col.data_type = str(data_type) if data_type else 'UNKNOWN'
                                existing_col.save()
                        else:
                            models.TableColumn.objects.create(
                                table=table_model,
                                name=col_name,
                                public_name=col_name,
                                data_type=str(data_type) if data_type else 'UNKNOWN',
                            )
                            result["columns_created"] += 1
                        found_column_names.add(col_name)
                        columns_count += 1

                    models.TableColumn.objects.filter(
                        table=table_model
                    ).exclude(name__in=found_column_names).delete()

                    result["tables"].append({
                        "name": full_table_name,
                        "status": "created" if not existing_table else "updated",
                        "id": table_model.id,
                        "columns": columns_count,
                        "source": "information_schema"
                    })

                except Exception as e:
                    logger.exception(f"Error syncing table {schema_name}.{table_name}: {e}")
                    result["tables"].append({
                        "name": f"{schema_name}.{table_name}",
                        "status": "error",
                        "status": "error",
                        "error": str(e)
                    })

            models.Table.objects.filter(
                data_source=datasource
            ).exclude(name__in=found_table_names).delete()

    except Exception as e:
        logger.exception(f"Error querying INFORMATION_SCHEMA: {e}")

    return tables_discovered


@reversion.create_revision()
def sync_metadata(datasource_id: int, overwrite: bool = False) -> Dict[str, Any]:
    from terno_dbi.core import models
    from terno_dbi.connectors import ConnectorFactory

    try:
        datasource = models.DataSource.objects.get(id=datasource_id, enabled=True)
    except models.DataSource.DoesNotExist:
        return {"error": f"Datasource {datasource_id} not found or not enabled"}

    connector = ConnectorFactory.create_connector(
        datasource.type,
        datasource.connection_str,
        credentials=datasource.connection_json
    )

    result = {
        "datasource_id": datasource_id,
        "datasource_name": datasource.display_name,
        "tables_created": 0,
        "tables_updated": 0,
        "tables_skipped": 0,
        "columns_created": 0,
        "columns_deleted": 0,
        "foreign_keys_created": 0,
        "tables": []
    }

    reversion.set_comment("Automatic Metadata Sync")

    try:
        if not datasource.dialect_name or not datasource.dialect_version:
            try:
                dialect_name, dialect_version = connector.get_dialect_info()
                datasource.dialect_name = dialect_name
                datasource.dialect_version = dialect_version
                datasource.save(update_fields=['dialect_name', 'dialect_version'])
            except Exception as e:
                logger.warning(f"Could not get dialect info: {e}")

        mdb = connector.get_metadata()

        logger.info(f"Found {len(mdb.tables)} tables in database via SQLShield")

        if len(mdb.tables) == 0:
            logger.info("SQLShield found 0 tables")

            # Skip INFORMATION_SCHEMA fallback for SQLite and Oracle as it doesn't exist
            if datasource.type not in ('sqlite', 'oracle'):
                logger.info("Falling back to INFORMATION_SCHEMA")
                tables_discovered = _sync_from_information_schema(
                    connector, datasource, result, overwrite
                )
                result["tables_synced"] = result["tables_created"] + result["tables_updated"]
                result["sync_method"] = "information_schema"
                logger.info(f"INFORMATION_SCHEMA fallback discovered {tables_discovered} tables/views")

                # Fetch row counts for fallback tables
                try:
                    from terno_dbi.core import models
                    found_tables = [t for t in result["tables"] if t.get("status") in ("created", "updated", "skipped")]
                    if found_tables:
                        table_names_list = [t["name"] for t in found_tables]
                        logger.info(f"Fetching row counts for {len(table_names_list)} fallback tables")

                        raw_counts = connector.get_table_row_counts(tables=table_names_list)
                        counts_map = build_row_counts_lookup(raw_counts)

                        updates_count = 0
                        for t_info in found_tables:
                            t_name = t_info["name"]
                            count = resolve_row_count(t_name, counts_map)
                            if count is not None:
                                models.Table.objects.filter(id=t_info["id"]).update(
                                    estimated_row_count=count
                                )
                                updates_count += 1
                        logger.info(f"Updated row counts for {updates_count} fallback tables")

                except Exception as e:
                    logger.warning(f"Fallback row stats warning: {e}")

                return result

        result["sync_method"] = "sqlshield"

        row_counts_map = {}
        try:
            table_names_list = list(mdb.tables.keys())
            raw_counts = connector.get_table_row_counts(tables=table_names_list)
            row_counts_map = build_row_counts_lookup(raw_counts)
            logger.info(f"Fetched row counts for {len(row_counts_map)} tables")
        except Exception as e:
            logger.warning(f"Failed to fetch row counts (skipping stats): {e}")

        found_table_names = set()

        for tbl_name, tbl in mdb.tables.items():
            try:
                actual_table_name = tbl.name
                existing_table = models.Table.objects.filter(
                    data_source=datasource, 
                    name=actual_table_name
                ).first()

                if existing_table:
                    table_model = existing_table
                    result["tables_updated"] += 1
                else:
                    table_model = models.Table.objects.create(
                        data_source=datasource,
                        name=actual_table_name,
                        public_name=actual_table_name,
                    )
                    result["tables_created"] += 1

                count = resolve_row_count(actual_table_name, row_counts_map)
                if count is not None:
                    table_model.estimated_row_count = count
                    table_model.save(update_fields=['estimated_row_count'])

                found_table_names.add(actual_table_name)

                columns_count = 0
                found_column_names = set()
                for col_name, col in tbl.columns.items():
                    existing_col = models.TableColumn.objects.filter(
                        table=table_model,
                        name=col_name
                    ).first()

                    if existing_col:
                        if overwrite:
                            existing_col.data_type = str(col.type)
                            existing_col.save()
                    else:
                        models.TableColumn.objects.create(
                            table=table_model,
                            name=col_name,
                            public_name=col_name,
                            data_type=str(col.type),
                        )
                        result["columns_created"] += 1
                    found_column_names.add(col_name)
                    columns_count += 1

                missing_cols = models.TableColumn.objects.filter(
                    table=table_model
                ).exclude(name__in=found_column_names)

                if missing_cols.exists():
                    count = missing_cols.count()
                    missing_cols.delete()
                    result["columns_deleted"] += count
                    logger.info(f"Deleted {count} stale columns from {actual_table_name}")

                result["tables"].append({
                    "name": actual_table_name,
                    "status": "created" if not existing_table else "updated",
                    "id": table_model.id,
                    "columns": columns_count
                })

            except Exception as e:
                logger.exception(f"Error syncing table {tbl_name}: {e}")
                result["tables"].append({
                    "name": tbl_name,
                    "status": "error",
                    "error": str(e)
                })

        for tbl_name, tbl in mdb.tables.items():
            try:
                foreign_keys = tbl.Foreign_Keys
                if not foreign_keys:
                    continue

                table = models.Table.objects.filter(
                    name=tbl.name, data_source=datasource
                ).first()

                if not table:
                    continue

                for fk in foreign_keys:
                    try:
                        constrained_columns = models.TableColumn.objects.filter(
                            name=fk.constrained_columns[0].name,
                            table__data_source=datasource
                        ).first()

                        referred_table = models.Table.objects.filter(
                            name=fk.referred_table.name, 
                            data_source=datasource
                        ).first()

                        referred_columns = models.TableColumn.objects.filter(
                            name=fk.referred_columns[0].name,
                            table__data_source=datasource
                        ).first()

                        if not all([constrained_columns, referred_table, referred_columns]):
                            continue

                        existing_fk = models.ForeignKey.objects.filter(
                            constrained_table=table,
                            constrained_columns=constrained_columns,
                            referred_table=referred_table,
                            referred_columns=referred_columns
                        ).first()

                        if not existing_fk:
                            models.ForeignKey.objects.create(
                                constrained_table=table,
                                constrained_columns=constrained_columns,
                                referred_table=referred_table,
                                referred_columns=referred_columns
                            )
                            result["foreign_keys_created"] += 1
                    except Exception as fk_error:
                        logger.warning(f"Error creating FK for {tbl_name}: {fk_error}")

            except Exception as e:
                logger.warning(f"Error processing FKs for table {tbl_name}: {e}")

        missing_tables = models.Table.objects.filter(data_source=datasource).exclude(name__in=found_table_names)
        if missing_tables.exists():
            count = missing_tables.count()

            missing_tables.delete()
            result['tables_deleted'] = count
            logger.info(f"Deleted {count} stale tables from data source {datasource.display_name}")

        result["tables_synced"] = result["tables_created"] + result["tables_updated"]
        return result

    except Exception as e:
        logger.exception(f"Error syncing metadata for datasource {datasource_id}: {e}")
        return {"error": str(e)}
