import logging
from sqlalchemy import text
from sqlalchemy.engine import Engine


def _dialect(engine: Engine) -> str:
    return engine.dialect.name.lower()


def _q(engine: Engine) -> str:
    return "`" if _dialect(engine) == "mysql" else '"'


def get_tables(engine: Engine) -> list[str]:
    with engine.connect() as conn:
        d = _dialect(engine)
        if d == "sqlite":
            rows = conn.execute(text("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"))
        elif d == "mysql":
            rows = conn.execute(text(
                "SELECT table_name FROM information_schema.tables "
                "WHERE table_schema = DATABASE() AND table_type = 'BASE TABLE' ORDER BY table_name"
            ))
        else:
            rows = conn.execute(text(
                "SELECT table_name FROM information_schema.tables "
                "WHERE table_schema NOT IN ('information_schema','pg_catalog') "
                "AND table_type = 'BASE TABLE' ORDER BY table_name"
            ))
        return [row[0] for row in rows]


def get_primary_keys(engine: Engine, table: str) -> set[str]:
    d = _dialect(engine)
    try:
        with engine.connect() as conn:
            if d == "mysql":
                rows = conn.execute(text(
                    "SELECT kcu.COLUMN_NAME FROM information_schema.KEY_COLUMN_USAGE kcu "
                    "JOIN information_schema.TABLE_CONSTRAINTS tc "
                    "  ON kcu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME "
                    "  AND kcu.TABLE_NAME = tc.TABLE_NAME "
                    "  AND kcu.TABLE_SCHEMA = tc.TABLE_SCHEMA "
                    "WHERE tc.CONSTRAINT_TYPE = 'PRIMARY KEY' "
                    "AND kcu.TABLE_NAME = :t AND kcu.TABLE_SCHEMA = DATABASE()"
                ), {"t": table})
            else:
                rows = conn.execute(text(
                    "SELECT kcu.column_name FROM information_schema.table_constraints tc "
                    "JOIN information_schema.key_column_usage kcu "
                    "  ON tc.constraint_name = kcu.constraint_name "
                    "  AND tc.table_name = kcu.table_name "
                    "WHERE tc.constraint_type = 'PRIMARY KEY' AND tc.table_name = :t"
                ), {"t": table})
            return {row[0] for row in rows}
    except Exception:
        return set()


def get_foreign_keys(engine: Engine, table: str) -> list[dict]:
    d = _dialect(engine)
    if d == "sqlite":
        return []
    try:
        with engine.connect() as conn:
            if d == "mysql":
                rows = conn.execute(text(
                    "SELECT kcu.COLUMN_NAME, kcu.REFERENCED_TABLE_NAME, kcu.REFERENCED_COLUMN_NAME "
                    "FROM information_schema.KEY_COLUMN_USAGE kcu "
                    "WHERE kcu.TABLE_NAME = :t AND kcu.TABLE_SCHEMA = DATABASE() "
                    "AND kcu.REFERENCED_TABLE_NAME IS NOT NULL"
                ), {"t": table})
            else:
                rows = conn.execute(text(
                    "SELECT kcu.column_name, ccu.table_name, ccu.column_name "
                    "FROM information_schema.table_constraints tc "
                    "JOIN information_schema.key_column_usage kcu "
                    "  ON tc.constraint_name = kcu.constraint_name "
                    "  AND tc.table_name = kcu.table_name "
                    "JOIN information_schema.constraint_column_usage ccu "
                    "  ON ccu.constraint_name = tc.constraint_name "
                    "WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_name = :t"
                ), {"t": table})
            return [{"column": r[0], "references": f"{r[1]}.{r[2]}"} for r in rows]
    except Exception as e:
        logging.warning("FK introspection failed: %s", e)
        return []


def get_columns(engine: Engine, table: str) -> list[dict]:
    d = _dialect(engine)
    with engine.connect() as conn:
        if d == "sqlite":
            rows = conn.execute(text(f'PRAGMA table_info("{table}")'))
            return [{"name": r[1], "type": r[2], "nullable": not r[3], "primary_key": bool(r[5])} for r in rows]
        elif d == "mysql":
            pk_cols = get_primary_keys(engine, table)
            rows = conn.execute(text(
                "SELECT column_name, data_type, is_nullable FROM information_schema.columns "
                "WHERE table_name = :t AND table_schema = DATABASE() ORDER BY ordinal_position"
            ), {"t": table})
        else:
            pk_cols = get_primary_keys(engine, table)
            rows = conn.execute(text(
                "SELECT column_name, data_type, is_nullable FROM information_schema.columns "
                "WHERE table_name = :t ORDER BY ordinal_position"
            ), {"t": table})
        return [{"name": r[0], "type": r[1], "nullable": r[2].upper() == "YES", "primary_key": r[0] in pk_cols} for r in rows]


def get_row_count(engine: Engine, table: str) -> int:
    q = _q(engine)
    with engine.connect() as conn:
        return conn.execute(text(f"SELECT COUNT(*) FROM {q}{table}{q}")).scalar() or 0


def get_column_stats(engine: Engine, table: str, column: str, row_count: int) -> dict:
    q = _q(engine)
    try:
        with engine.connect() as conn:
            row = conn.execute(text(
                f"SELECT COUNT(*) - COUNT({q}{column}{q}) AS null_count, "
                f"COUNT(DISTINCT {q}{column}{q}) AS distinct_count FROM {q}{table}{q}"
            )).fetchone()
            null_count = row[0] or 0
            distinct_count = row[1] or 0
        null_rate = null_count / row_count if row_count > 0 else 0.0
        return {"null_count": null_count, "null_rate": null_rate, "distinct_count": distinct_count}
    except Exception:
        return {"null_count": None, "null_rate": None, "distinct_count": None}


def get_sample_values(engine: Engine, table: str, column: str, n: int = 5) -> list:
    q = _q(engine)
    with engine.connect() as conn:
        result = conn.execute(text(
            f"SELECT {q}{column}{q} FROM {q}{table}{q} "
            f"WHERE {q}{column}{q} IS NOT NULL "
            f"GROUP BY {q}{column}{q} ORDER BY COUNT(*) DESC LIMIT {n}"
        ))
        return [row[0] for row in result]


def describe_table(engine: Engine, table: str) -> dict:
    columns = get_columns(engine, table)
    row_count = get_row_count(engine, table)
    columns_info = []
    for col in columns:
        stats = get_column_stats(engine, table, col["name"], row_count)
        samples = get_sample_values(engine, table, col["name"])
        columns_info.append({**col, **stats, "samples": samples})
    return {"table": table, "row_count": row_count, "columns": columns_info, "foreign_keys": get_foreign_keys(engine, table)}


def explore_schema(engine: Engine, user_context: str | None = None) -> str:
    tables = get_tables(engine)
    if not tables:
        return "No tables found in the database."
    lines = []
    if user_context:
        lines += ["--- User-provided context ---", user_context, "-----------------------------", ""]
    lines.append(f"Database schema ({len(tables)} tables)\n{'='*50}")
    for table in tables:
        info = describe_table(engine, table)
        lines.append(f"\nTable: {table}  ({info['row_count']:,} rows)")
        lines.append("-" * 40)
        for col in info["columns"]:
            pk_flag = " [PK]" if col["primary_key"] else ""
            null_flag = "" if col["nullable"] else " NOT NULL"
            samples_str = ", ".join(str(s) for s in col["samples"]) if col["samples"] else "—"
            lines.append(f"  {col['name']}{pk_flag}  {col['type']}{null_flag}")
            stat_parts = []
            if col.get("null_rate") is not None and col["null_rate"] > 0:
                stat_parts.append(f"{col['null_rate']:.0%} null")
            if col.get("distinct_count") is not None:
                stat_parts.append(f"{col['distinct_count']:,} distinct")
            stat_str = " | ".join(stat_parts)
            lines.append(f"    {stat_str + ' | ' if stat_str else ''}samples: {samples_str}")
        if info["foreign_keys"]:
            lines.append("  Foreign keys:")
            for fk in info["foreign_keys"]:
                lines.append(f"    {fk['column']} → {fk['references']}")
    return "\n".join(lines)
