"""
MySQL 连接封装：按配置创建连接，超时与只读检查由调用方配合 SafetyConfig 使用。
"""
from __future__ import annotations

import contextlib
from typing import Generator

import pymysql
from pymysql.cursors import DictCursor

from .config import ConnectionConfig, SafetyConfig


def get_connection_config(
    profile: str | None,
    default: ConnectionConfig,
    connection_map: dict[str, ConnectionConfig],
) -> ConnectionConfig:
    """
    解析 profile 得到连接配置。
    profile 为空时返回 default；否则从 connection_map 取，不存在则抛出 ValueError。
    """
    if not (profile or "").strip():
        return default
    key = (profile or "").strip().lower()
    if key not in connection_map:
        raise ValueError(f"未知的 connection profile: {profile}，可用: {', '.join(sorted(connection_map))}")
    return connection_map[key]


@contextlib.contextmanager
def connect(
    config: ConnectionConfig,
    safety: SafetyConfig,
) -> Generator[pymysql.Connection, None, None]:
    """
    创建 MySQL 连接，设置字符集与超时。调用方负责在只读模式下不执行写操作。
    """
    conn = pymysql.connect(
        host=config.host,
        port=config.port,
        user=config.user,
        password=config.password,
        database=config.database,
        charset=config.charset,
        cursorclass=DictCursor,
        connect_timeout=10,
        read_timeout=safety.query_timeout_seconds,
        write_timeout=safety.query_timeout_seconds,
        defer_connect=False,
    )
    try:
        yield conn
    finally:
        conn.close()


def list_databases(
    conn: pymysql.Connection,
    allowed: list[str] | None = None,
) -> list[str]:
    """列出当前连接可见的数据库名。allowed=None 表示所有库；allowed=[] 表示不允许任何库；allowed 非空则只返回这些库。"""
    with conn.cursor() as cur:
        cur.execute("SHOW DATABASES")
        all_dbs = [row["Database"] for row in cur.fetchall()]
    if allowed is None:
        return all_dbs
    if not allowed:
        return []  # 配置了但空数组：不允许任何库
    all_set = set(all_dbs)
    return [d for d in allowed if d in all_set]


def list_tables(conn: pymysql.Connection, database: str) -> list[dict]:
    """列出指定库下的表，含可选 TABLE_ROWS（近似）"""
    with conn.cursor() as cur:
        cur.execute("USE `%s`" % database.replace("`", "``"))
        cur.execute(
            """
            SELECT TABLE_NAME AS name, TABLE_ROWS AS rows_approx, TABLE_COMMENT AS comment
            FROM information_schema.TABLES
            WHERE TABLE_SCHEMA = %s AND TABLE_TYPE = 'BASE TABLE'
            ORDER BY TABLE_NAME
            """,
            (database,),
        )
        return [dict(row) for row in cur.fetchall()]


def describe_table(conn: pymysql.Connection, database: str, table: str) -> dict:
    """
    返回表结构：列名、类型、主键、索引、可选行数/大小。
    """
    with conn.cursor() as cur:
        cur.execute("USE `%s`" % database.replace("`", "``"))
        # 列信息
        cur.execute(
            """
            SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_KEY, COLUMN_DEFAULT, EXTRA
            FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
            ORDER BY ORDINAL_POSITION
            """,
            (database, table),
        )
        columns = [dict(row) for row in cur.fetchall()]
        # 表行数/大小（近似）
        cur.execute(
            """
            SELECT TABLE_ROWS, ROUND((DATA_LENGTH + INDEX_LENGTH) / 1024) AS size_kb
            FROM information_schema.TABLES
            WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
            """,
            (database, table),
        )
        row = cur.fetchone()
        rows_approx = row["TABLE_ROWS"] if row and row.get("TABLE_ROWS") is not None else None
        size_kb = row["size_kb"] if row and row.get("size_kb") is not None else None
    return {
        "database": database,
        "table": table,
        "columns": columns,
        "rows_approx": rows_approx,
        "size_kb": size_kb,
    }


def execute_select(
    conn: pymysql.Connection,
    sql: str,
    max_rows: int,
    limit: int | None = None,
    offset: int | None = None,
) -> tuple[list[dict], bool]:
    """
    执行 SELECT，返回 (rows, truncated)。
    若 SQL 中无 LIMIT，则追加 LIMIT max_rows（或调用方传入的 limit/offset），避免一次拉取过多。
    """
    sql_clean = sql.rstrip(";").strip()
    has_limit = "LIMIT" in sql_clean.upper().split("--")[0]
    with conn.cursor() as cur:
        if not has_limit:
            l, o = (limit if limit is not None else max_rows), (offset or 0)
            sql_to_run = f"{sql_clean} LIMIT {int(l)} OFFSET {int(o)}"
        else:
            sql_to_run = sql_clean
        cur.execute(sql_to_run)
        rows = cur.fetchall()
        truncated = len(rows) >= (limit if limit is not None else max_rows) and not has_limit
        return [dict(r) for r in rows], truncated
