#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""MySQL 异步处理器"""
import logging
import time
from typing import Any, Dict, List, Optional

from core.database_manager.handlers.common import (
    format_connection_error,
    format_size,
    log_database_connect_failure,
    serialize_row,
)
from core.database_manager.handlers.pools import get_aiomysql_pool
from core.database_manager.handlers.transaction_mixin import HandlerTransactionMixin
from core.database_manager.sql_utils import quote_identifier, split_sql_statements

logger = logging.getLogger(__name__)

class MySQLHandler(HandlerTransactionMixin):
    """MySQL异步处理器"""

    def __init__(self, host: str, port: int, user: str, password: str, database: str):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.conn = None
        self._pool = None
        self._last_connect_error: Optional[str] = None
        self._in_transaction = False

    async def _release_connection(self) -> None:
        if self.conn and self._pool:
            try:
                self._pool.release(self.conn)
            except Exception as exc:
                logger.warning("Error releasing MySQL connection: %s", exc)
        self.conn = None
        self._pool = None

    async def connect(self, database: str = None) -> bool:
        db = database or self.database
        try:
            await self._release_connection()

            pool = await get_aiomysql_pool(
                self.host, self.port, self.user, self.password, db,
            )
            self._pool = pool
            self.conn = await pool.acquire()
            self._last_connect_error = None
            return True
        except Exception as e:
            self._last_connect_error = format_connection_error(e)
            log_database_connect_failure(
                db_type="mysql",
                host=self.host,
                port=self.port,
                user=self.user,
                database=db,
                detail=self._last_connect_error,
            )
            self.conn = None
            self._pool = None
            return False

    async def close(self):
        await self._release_connection()

    async def _execute_query(self, query: str, params: tuple = None) -> List[Dict[str, Any]]:
        import aiomysql
        async with self.conn.cursor(aiomysql.DictCursor) as cursor:
            await cursor.execute(query, params or ())
            rows = await cursor.fetchall()
            return [dict(row) for row in rows]

    async def _execute_command(self, command: str, params: tuple = None) -> int:
        async with self.conn.cursor() as cursor:
            await cursor.execute(command, params or ())
            return cursor.rowcount

    async def get_databases(self) -> List[Dict[str, Any]]:
        try:
            if not await self.connect():
                return []
            query = """
                SELECT SCHEMA_NAME as name, DEFAULT_CHARACTER_SET_NAME as encoding, DEFAULT_COLLATION_NAME as collation,
                       (SELECT COUNT(*) FROM information_schema.TABLES WHERE TABLE_SCHEMA = SCHEMA_NAME AND TABLE_TYPE = 'BASE TABLE') as tables_count
                FROM information_schema.SCHEMATA
                WHERE SCHEMA_NAME NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys') ORDER BY SCHEMA_NAME
            """
            databases = await self._execute_query(query)
            for db in databases:
                size_result = await self._execute_query("SELECT ROUND(SUM(data_length + index_length), 2) as size_bytes FROM information_schema.TABLES WHERE table_schema = %s", (db['name'],))
                size_bytes = size_result[0]['size_bytes'] if size_result and size_result[0]['size_bytes'] else 0
                db['size_bytes'] = int(size_bytes) if size_bytes else 0
                db['size'] = format_size(db['size_bytes'])
            return databases
        except Exception as e:
            logger.error(f"Failed to get databases: {e}")
            return []
        finally:
            await self.close()

    async def create_database(self, name: str, charset: str = "utf8mb4", collation: str = "utf8mb4_unicode_ci", **kwargs) -> bool:
        try:
            if not await self.connect():
                return False
            qname = quote_identifier(name, "mysql")
            await self._execute_command(
                f"CREATE DATABASE {qname} CHARACTER SET {charset} COLLATE {collation}"
            )
            return True
        except Exception as e:
            logger.error(f"Failed to create database {name}: {e}")
            raise
        finally:
            await self.close()

    async def drop_database(self, name: str) -> bool:
        try:
            if not await self.connect():
                return False
            qname = quote_identifier(name, "mysql")
            await self._execute_command(f"DROP DATABASE {qname}")
            return True
        except Exception as e:
            logger.error(f"Failed to drop database {name}: {e}")
            raise
        finally:
            await self.close()

    async def get_schemas(self, database: str = None) -> List[Dict[str, Any]]:
        databases = await self.get_databases()
        return [{"name": db["name"], "owner": None, "tables_count": db.get("tables_count", 0)} for db in databases]

    async def get_tables(self, database: str = None, schema_name: str = None) -> List[Dict[str, Any]]:
        try:
            db_name = database or schema_name or self.database
            if not await self.connect(db_name):
                return []
            query = """
                SELECT TABLE_SCHEMA as schema_name, TABLE_NAME as table_name, TABLE_TYPE as table_type,
                       TABLE_ROWS as row_count, DATA_LENGTH as data_length, INDEX_LENGTH as index_length,
                       (DATA_LENGTH + INDEX_LENGTH) as total_size_bytes, TABLE_COMMENT as description
                FROM information_schema.TABLES WHERE TABLE_SCHEMA = %s AND TABLE_TYPE = 'BASE TABLE' ORDER BY TABLE_NAME
            """
            tables = await self._execute_query(query, (db_name,))
            for table in tables:
                table['total_size'] = format_size(table.get('total_size_bytes', 0) or 0)
            return tables
        except Exception as e:
            logger.error(f"Failed to get tables: {e}")
            return []
        finally:
            await self.close()

    async def get_table_columns(self, table_name: str, schema_name: str = None) -> List[Dict[str, Any]]:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                return []
            query = """
                SELECT COLUMN_NAME as column_name, DATA_TYPE as data_type, IS_NULLABLE = 'YES' as is_nullable,
                       COLUMN_DEFAULT as column_default, CHARACTER_MAXIMUM_LENGTH as character_maximum_length,
                       NUMERIC_PRECISION as numeric_precision, NUMERIC_SCALE as numeric_scale,
                       ORDINAL_POSITION as ordinal_position, COLUMN_KEY = 'PRI' as is_primary_key,
                       COLUMN_KEY = 'UNI' as is_unique, COLUMN_COMMENT as description
                FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s ORDER BY ORDINAL_POSITION
            """
            rows = await self._execute_query(query, (db_name, table_name))
            return rows
        except Exception as e:
            logger.error(f"Failed to get table columns: {e}")
            return []
        finally:
            await self.close()

    async def get_table_indexes(self, table_name: str, schema_name: str = None) -> List[Dict[str, Any]]:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                return []
            query = """
                SELECT INDEX_NAME as index_name, INDEX_TYPE as index_type,
                       GROUP_CONCAT(COLUMN_NAME ORDER BY SEQ_IN_INDEX) as columns,
                       NON_UNIQUE = 0 as is_unique, INDEX_NAME = 'PRIMARY' as is_primary
                FROM information_schema.STATISTICS WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
                GROUP BY INDEX_NAME, INDEX_TYPE, NON_UNIQUE ORDER BY INDEX_NAME
            """
            indexes = await self._execute_query(query, (db_name, table_name))
            for idx in indexes:
                idx['definition'] = f"INDEX {idx['index_name']} ({idx['columns']})"
            return indexes
        except Exception as e:
            logger.error(f"Failed to get table indexes: {e}")
            return []
        finally:
            await self.close()

    async def get_table_constraints(self, table_name: str, schema_name: str = None) -> List[Dict[str, Any]]:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                return []
            query = """
                SELECT tc.CONSTRAINT_NAME as constraint_name, tc.CONSTRAINT_TYPE as constraint_type,
                       GROUP_CONCAT(kcu.COLUMN_NAME) as columns, kcu.REFERENCED_TABLE_NAME as referenced_table, NULL as referenced_columns
                FROM information_schema.TABLE_CONSTRAINTS tc
                JOIN information_schema.KEY_COLUMN_USAGE kcu ON tc.CONSTRAINT_NAME = kcu.CONSTRAINT_NAME AND tc.TABLE_SCHEMA = kcu.TABLE_SCHEMA AND tc.TABLE_NAME = kcu.TABLE_NAME
                WHERE tc.TABLE_SCHEMA = %s AND tc.TABLE_NAME = %s
                GROUP BY tc.CONSTRAINT_NAME, tc.CONSTRAINT_TYPE, kcu.REFERENCED_TABLE_NAME ORDER BY tc.CONSTRAINT_TYPE, tc.CONSTRAINT_NAME
            """
            constraints = await self._execute_query(query, (db_name, table_name))
            for const in constraints:
                const['definition'] = f"{const['constraint_type']} ({const['columns']})"
            return constraints
        except Exception as e:
            logger.error(f"Failed to get table constraints: {e}")
            return []
        finally:
            await self.close()

    async def get_table_structure(self, table_name: str, database: str = None, schema_name: str = None) -> Dict[str, Any]:
        db_name = database or schema_name or self.database
        tables = await self.get_tables(database=db_name, schema_name=db_name)
        table_info = next((t for t in tables if t['table_name'] == table_name), None)
        if not table_info:
            raise ValueError(f"Table {db_name}.{table_name} not found")
        columns = await self.get_table_columns(table_name, db_name)
        indexes = await self.get_table_indexes(table_name, db_name)
        constraints = await self.get_table_constraints(table_name, db_name)
        return {"table_info": table_info, "columns": columns, "indexes": indexes, "constraints": constraints}

    async def get_table_ddl(self, table_name: str, schema_name: str = None) -> str:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                return "-- 无法连接数据库"
            result = await self._execute_query(f"SHOW CREATE TABLE `{table_name}`")
            if result:
                return result[0].get('Create Table', f"-- 无法获取表 {table_name} 的DDL")
            return f"-- 无法获取表 {table_name} 的DDL"
        except Exception as e:
            return f"-- 获取DDL失败: {str(e)}"
        finally:
            await self.close()

    async def get_views(self, database: str = None, schema_name: str = None) -> List[Dict[str, Any]]:
        try:
            db_name = database or schema_name or self.database
            if not await self.connect(db_name):
                return []
            query = """
                SELECT TABLE_NAME as view_name, TABLE_SCHEMA as schema_name, VIEW_DEFINITION as view_definition,
                       IS_UPDATABLE as is_updatable, CHECK_OPTION as check_option, 'VIEW' as view_type
                FROM information_schema.VIEWS WHERE TABLE_SCHEMA = %s ORDER BY TABLE_NAME
            """
            result = await self._execute_query(query, (db_name,))
            for row in result:
                row['is_updatable'] = row.get('is_updatable') == 'YES'
            return result
        except Exception as e:
            logger.error(f"Failed to get views: {e}")
            return []
        finally:
            await self.close()

    async def get_view_structure(self, view_name: str, schema_name: str = None) -> Dict[str, Any]:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                raise ValueError("Failed to connect")
            view_query = "SELECT TABLE_NAME as view_name, TABLE_SCHEMA as schema_name, VIEW_DEFINITION as view_definition, IS_UPDATABLE as is_updatable, CHECK_OPTION as check_option, 'VIEW' as view_type FROM information_schema.VIEWS WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s"
            view_info = await self._execute_query(view_query, (db_name, view_name))
            if not view_info:
                raise ValueError(f"View {db_name}.{view_name} not found")
            view_info = view_info[0]
            view_info['is_updatable'] = view_info.get('is_updatable') == 'YES'
            columns_query = "SELECT COLUMN_NAME as column_name, DATA_TYPE as data_type, IS_NULLABLE = 'YES' as is_nullable, ORDINAL_POSITION as ordinal_position, COLUMN_COMMENT as description FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s ORDER BY ORDINAL_POSITION"
            columns = await self._execute_query(columns_query, (db_name, view_name))
            definition_sql = await self._get_view_definition_internal(view_name)
            dependencies = await self._get_view_dependencies_internal(view_name, db_name)
            return {"view_info": view_info, "columns": columns, "dependencies": dependencies, "definition_sql": definition_sql}
        except Exception as e:
            logger.error(f"Failed to get view structure: {e}")
            raise
        finally:
            await self.close()

    async def _get_view_definition_internal(self, view_name: str) -> str:
        try:
            result = await self._execute_query(f"SHOW CREATE VIEW `{view_name}`")
            if result:
                return result[0].get('Create View', f"-- 无法获取视图 {view_name} 的定义")
        except Exception as e:
            logger.error(f"Failed to get view definition: {e}")
        return f"-- 无法获取视图 {view_name} 的定义"

    async def _get_view_dependencies_internal(self, view_name: str, schema_name: str) -> List[str]:
        try:
            query = "SELECT DISTINCT REFERENCED_TABLE_NAME as table_name FROM information_schema.VIEW_TABLE_USAGE WHERE VIEW_SCHEMA = %s AND VIEW_NAME = %s AND REFERENCED_TABLE_NAME IS NOT NULL ORDER BY REFERENCED_TABLE_NAME"
            result = await self._execute_query(query, (schema_name, view_name))
            return [row['table_name'] for row in result]
        except Exception as e:
            logger.error(f"Failed to get view dependencies: {e}")
            return []

    async def get_view_definition(self, view_name: str, schema_name: str = None) -> str:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                return "-- 无法连接数据库"
            result = await self._get_view_definition_internal(view_name)
            return result
        except Exception as e:
            return f"-- 获取视图定义失败: {str(e)}"
        finally:
            await self.close()

    async def get_view_dependencies(self, view_name: str, schema_name: str = None) -> List[str]:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                return []
            result = await self._get_view_dependencies_internal(view_name, db_name)
            return result
        except Exception as e:
            return []
        finally:
            await self.close()

    async def query_data(self, table_name: str, schema_name: str = None, page: int = 1, page_size: int = 20, where: str = None, order_by: str = None, database: str = None) -> Dict[str, Any]:
        try:
            db_name = database or schema_name or self.database
            if not await self.connect(db_name):
                return {"columns": [], "rows": [], "total": 0, "page": page, "page_size": page_size}
            count_q = f"SELECT COUNT(*) as total FROM `{table_name}`" + (f" WHERE {where}" if where else "")
            count_result = await self._execute_query(count_q)
            total = count_result[0]['total'] if count_result else 0
            offset = (page - 1) * page_size
            data_q = f"SELECT * FROM `{table_name}`" + (f" WHERE {where}" if where else "") + (f" ORDER BY {order_by}" if order_by else "") + f" LIMIT {page_size} OFFSET {offset}"
            rows = await self._execute_query(data_q)
            rows_list = [serialize_row(row) for row in rows]
            columns = list(rows_list[0].keys()) if rows_list else []
            return {"columns": columns, "rows": rows_list, "total": total, "page": page, "page_size": page_size}
        except Exception as e:
            logger.error(f"Failed to query data: {e}")
            return {"columns": [], "rows": [], "total": 0, "page": page, "page_size": page_size}
        finally:
            await self.close()

    async def _run_sql_query(self, sql: str) -> list:
        return await self._execute_query(sql)

    async def _run_sql_command(self, sql: str) -> int:
        return await self._execute_command(sql)

    async def insert_data(self, table_name: str, data: Dict[str, Any], schema_name: str = None) -> Dict[str, Any]:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                return {"success": False, "message": "数据库连接失败", "affected_rows": 0}
            columns = list(data.keys())
            values = list(data.values())
            placeholders = ', '.join(['%s'] * len(values))
            query = f"INSERT INTO `{table_name}` ({', '.join(columns)}) VALUES ({placeholders})"
            affected_rows = await self._execute_command(query, tuple(values))
            return {"success": True, "message": "插入成功", "affected_rows": affected_rows}
        except Exception as e:
            return {"success": False, "message": str(e), "affected_rows": 0}
        finally:
            await self.close()

    async def update_data(self, table_name: str, data: Dict[str, Any], where: str, schema_name: str = None) -> Dict[str, Any]:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                return {"success": False, "message": "数据库连接失败", "affected_rows": 0}
            set_clause = ', '.join([f'{k} = %s' for k in data.keys()])
            query = f"UPDATE `{table_name}` SET {set_clause} WHERE {where}"
            affected_rows = await self._execute_command(query, tuple(data.values()))
            return {"success": True, "message": f"更新成功，影响 {affected_rows} 行", "affected_rows": affected_rows}
        except Exception as e:
            return {"success": False, "message": str(e), "affected_rows": 0}
        finally:
            await self.close()

    async def delete_data(self, table_name: str, where: str, schema_name: str = None) -> Dict[str, Any]:
        try:
            db_name = schema_name or self.database
            if not await self.connect(db_name):
                return {"success": False, "message": "数据库连接失败", "affected_rows": 0}
            query = f"DELETE FROM `{table_name}` WHERE {where}"
            affected_rows = await self._execute_command(query)
            return {"success": True, "message": f"删除成功，影响 {affected_rows} 行", "affected_rows": affected_rows}
        except Exception as e:
            return {"success": False, "message": str(e), "affected_rows": 0}
        finally:
            await self.close()

    async def execute_ddl(self, sql: str, database: str = None, schema_name: str = None) -> Dict[str, Any]:
        try:
            db_name = database or schema_name or self.database
            if not await self.connect(db_name):
                return {"success": False, "message": "数据库连接失败", "affected_rows": 0}
            statements = split_sql_statements(sql)
            if not statements:
                return {"success": False, "message": "无有效SQL语句", "affected_rows": 0}
            await self.conn.begin()
            try:
                for index, statement in enumerate(statements, start=1):
                    try:
                        await self._execute_command(statement)
                    except Exception as exc:
                        await self.conn.rollback()
                        return {
                            "success": False,
                            "message": f"DDL执行失败 (第{index}条): {exc}",
                            "affected_rows": 0,
                        }
                await self.conn.commit()
            except Exception:
                await self.conn.rollback()
                raise
            return {"success": True, "message": "DDL执行成功", "affected_rows": 0}
        except Exception as e:
            return {"success": False, "message": f"DDL执行失败: {str(e)}", "affected_rows": 0}
        finally:
            await self.close()


