"""Schema introspection and per-column profiling for a user's database.

Identifiers (table/column names) are quoted via the engine's dialect preparer,
which handles reserved words, mixed case, and embedded quotes correctly across
dialects. Values used in SQL come from SQLAlchemy inspection of the DB itself,
not user input.
"""

from typing import Optional

import pandas as pd
from sqlalchemy import Float, Integer, Numeric, inspect
from sqlalchemy.engine import Engine

from src.middlewares.logging import get_logger

logger = get_logger("db_extractor")

TOP_VALUES_THRESHOLD = 0.05  # show top values if distinct_ratio <= 5%


def _qi(engine: Engine, name: str) -> str:
    """Dialect-correct identifier quoting (schema.table also handled if dotted)."""
    preparer = engine.dialect.identifier_preparer
    if "." in name:
        schema, _, table = name.partition(".")
        return f"{preparer.quote(schema)}.{preparer.quote(table)}"
    return preparer.quote(name)


def get_schema(
    engine: Engine, exclude_tables: Optional[frozenset[str]] = None
) -> dict[str, list[dict]]:
    """Returns {table_name: [{name, type, is_numeric, is_primary_key, foreign_key}, ...]}."""
    exclude = exclude_tables or frozenset()
    inspector = inspect(engine)
    schema = {}
    for table_name in inspector.get_table_names():
        if table_name in exclude:
            continue

        pk = inspector.get_pk_constraint(table_name)
        pk_cols = set(pk["constrained_columns"]) if pk else set()

        fk_map = {}
        for fk in inspector.get_foreign_keys(table_name):
            for col, ref_col in zip(fk["constrained_columns"], fk["referred_columns"]):
                fk_map[col] = f"{fk['referred_table']}.{ref_col}"

        cols = inspector.get_columns(table_name)
        schema[table_name] = [
            {
                "name": c["name"],
                "type": str(c["type"]),
                "is_numeric": isinstance(c["type"], (Integer, Numeric, Float)),
                "is_primary_key": c["name"] in pk_cols,
                "foreign_key": fk_map.get(c["name"]),
            }
            for c in cols
        ]
    logger.info("extracted schema", table_count=len(schema))
    return schema


def get_row_count(engine: Engine, table_name: str) -> int:
    return pd.read_sql(f"SELECT COUNT(*) FROM {_qi(engine, table_name)}", engine).iloc[0, 0]


def profile_column(
    engine: Engine,
    table_name: str,
    col_name: str,
    is_numeric: bool,
    row_count: int,
) -> dict:
    """Returns null_count, distinct_count, min/max, top values, and sample values."""
    if row_count == 0:
        return {
            "null_count": 0,
            "distinct_count": 0,
            "distinct_ratio": 0.0,
            "sample_values": [],
        }

    qt = _qi(engine, table_name)
    qc = _qi(engine, col_name)

    # Combined stats query: null_count, distinct_count, and min/max (if numeric).
    # One round-trip instead of two.
    select_cols = [
        f"COUNT(*) - COUNT({qc}) AS nulls",
        f"COUNT(DISTINCT {qc}) AS distincts",
    ]
    if is_numeric:
        select_cols.append(f"MIN({qc}) AS min_val")
        select_cols.append(f"MAX({qc}) AS max_val")
        select_cols.append(f"AVG({qc}) AS mean_val")
        # PERCENTILE_CONT is supported by Postgres and SQL Server; MySQL would need
        # a dialect-specific fallback when that connector is added.
        select_cols.append(
            f"PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY {qc}) AS median_val"
        )
    stats = pd.read_sql(f"SELECT {', '.join(select_cols)} FROM {qt}", engine)

    null_count = int(stats.iloc[0]["nulls"])
    distinct_count = int(stats.iloc[0]["distincts"])
    distinct_ratio = distinct_count / row_count if row_count > 0 else 0

    profile = {
        "null_count": null_count,
        "distinct_count": distinct_count,
        "distinct_ratio": round(distinct_ratio, 4),
    }

    if is_numeric:
        profile["min"] = stats.iloc[0]["min_val"]
        profile["max"] = stats.iloc[0]["max_val"]
        profile["mean"] = stats.iloc[0]["mean_val"]
        profile["median"] = stats.iloc[0]["median_val"]

    if 0 < distinct_ratio <= TOP_VALUES_THRESHOLD:
        top = pd.read_sql(
            f"SELECT {qc}, COUNT(*) AS cnt FROM {qt} "
            f"GROUP BY {qc} ORDER BY cnt DESC LIMIT 10",
            engine,
        )
        profile["top_values"] = list(zip(top[col_name].tolist(), top["cnt"].tolist()))

    sample = pd.read_sql(f"SELECT {qc} FROM {qt} LIMIT 5", engine)
    profile["sample_values"] = sample[col_name].tolist()

    return profile


def profile_table(engine: Engine, table_name: str, columns: list[dict]) -> list[dict]:
    """Profile every column in a table. Returns [{col, profile, text}, ...].

    Per-column errors are logged and skipped so one bad column doesn't abort
    the whole table.
    """
    row_count = get_row_count(engine, table_name)
    if row_count == 0:
        logger.info("skipping empty table", table=table_name)
        return []

    results = []
    for col in columns:
        try:
            profile = profile_column(
                engine, table_name, col["name"], col.get("is_numeric", False), row_count
            )
            text = build_text(table_name, row_count, col, profile)
            results.append({"col": col, "profile": profile, "text": text})
        except Exception as e:
            logger.error(
                "column profiling failed",
                table=table_name,
                column=col["name"],
                error=str(e),
            )
            continue
    return results


def build_text(table_name: str, row_count: int, col: dict, profile: dict) -> str:
    col_name = col["name"]
    col_type = col["type"]

    key_label = ""
    if col.get("is_primary_key"):
        key_label = " [PRIMARY KEY]"
    elif col.get("foreign_key"):
        key_label = f" [FK -> {col['foreign_key']}]"

    text = f"Table: {table_name} ({row_count} rows)\n"
    text += f"Column: {col_name} ({col_type}){key_label}\n"
    text += f"Null count: {profile['null_count']}\n"
    text += f"Distinct count: {profile['distinct_count']} ({profile['distinct_ratio']:.1%})\n"
    if "min" in profile:
        text += f"Min: {profile['min']}, Max: {profile['max']}\n"
        text += f"Mean: {profile['mean']}, Median: {profile['median']}\n"
    if "top_values" in profile:
        top_str = ", ".join(f"{v} ({c})" for v, c in profile["top_values"])
        text += f"Top values: {top_str}\n"
    text += f"Sample values: {profile['sample_values']}"
    return text