import os
import sys
import asyncio
from pathlib import Path
import aiomysql
import datetime
import inspect

sys.path.insert(0, str(Path(__file__).parent.parent))

from config import settings
from database.models import Base
from sqlalchemy import inspect as sqlalchemy_inspect
from utils.logger import session_logger

async def get_mysql_connection():
    """建立MySQL数据库连接"""
    try:
        conn = await aiomysql.connect(
            host=settings.db_host,
            port=settings.db_port,
            user=settings.db_user,
            password=settings.db_password,
            db=settings.db_name,
            autocommit=False
        )
        return conn
    except Exception as e:
        print(f"数据库连接失败: {e}")
        return None

async def get_all_tables(conn):
    """获取数据库中所有表"""
    async with conn.cursor() as cursor:
        await cursor.execute("SHOW TABLES")
        tables = await cursor.fetchall()
        return [table[0] for table in tables]

async def get_model_tables():
    """获取模型定义的所有表"""
    return [table.__tablename__ for table in Base.__subclasses__()]

async def get_table_columns(conn, table_name):
    """获取表的所有列"""
    async with conn.cursor() as cursor:
        await cursor.execute(f"SHOW COLUMNS FROM `{table_name}`")
        columns = await cursor.fetchall()
        return [column[0] for column in columns]

async def get_model_columns(model_class):
    """获取模型定义的所有列"""
    inspector = sqlalchemy_inspect(model_class)
    return [column.key for column in inspector.mapper.column_attrs]

async def identify_unused_tables(conn):
    """识别未使用的表"""
    db_tables = await get_all_tables(conn)
    model_tables = await get_model_tables()
    
    unused_tables = [table for table in db_tables if table not in model_tables]
    return unused_tables

async def identify_unused_columns(conn):
    """识别未使用的列"""
    results = {}
    model_dict = {model.__tablename__: model for model in Base.__subclasses__()}
    
    for table_name, model_class in model_dict.items():
        try:
            db_columns = await get_table_columns(conn, table_name)
            model_columns = await get_model_columns(model_class)
            
            unused_columns = [col for col in db_columns if col not in model_columns]
            if unused_columns:
                results[table_name] = unused_columns
        except Exception as e:
            print(f"检查表 {table_name} 列时出错: {e}")
    
    return results

async def get_table_foreign_keys(conn, table_name):
    """获取表的外键约束"""
    async with conn.cursor() as cursor:
        query = """
        SELECT
            CONSTRAINT_NAME,
            TABLE_NAME,
            COLUMN_NAME,
            REFERENCED_TABLE_NAME,
            REFERENCED_COLUMN_NAME
        FROM
            INFORMATION_SCHEMA.KEY_COLUMN_USAGE
        WHERE
            REFERENCED_TABLE_NAME IS NOT NULL
            AND TABLE_SCHEMA = %s
            AND TABLE_NAME = %s
        """
        await cursor.execute(query, (settings.db_name, table_name))
        foreign_keys = await cursor.fetchall()
        return foreign_keys

async def get_referenced_tables(conn, table_name):
    """获取引用指定表的所有表"""
    async with conn.cursor() as cursor:
        query = """
        SELECT
            TABLE_NAME,
            CONSTRAINT_NAME
        FROM
            INFORMATION_SCHEMA.KEY_COLUMN_USAGE
        WHERE
            REFERENCED_TABLE_NAME = %s
            AND TABLE_SCHEMA = %s
        """
        await cursor.execute(query, (table_name, settings.db_name))
        references = await cursor.fetchall()
        
        if references:
            print(f"表 {table_name} 被以下表引用:")
            for ref in references:
                print(f"  - {ref[0]} (约束: {ref[1]})")
                
        return references

async def disable_foreign_key_checks(conn):
    """临时禁用外键约束检查"""
    async with conn.cursor() as cursor:
        await cursor.execute("SET FOREIGN_KEY_CHECKS = 0")
        print("已临时禁用外键约束检查")

async def enable_foreign_key_checks(conn):
    """重新启用外键约束检查"""
    async with conn.cursor() as cursor:
        await cursor.execute("SET FOREIGN_KEY_CHECKS = 1")
        print("已重新启用外键约束检查")

async def sort_tables_for_deletion(conn, tables):
    """对表进行排序，以确保先删除引用者，再删除被引用者"""
    dependency_graph = {}
    for table in tables:
        references = await get_referenced_tables(conn, table)
        dependency_graph[table] = [ref[0] for ref in references if ref[0] in tables]
    
    result = []
    visited = set()
    temp_mark = set()
    
    def visit(node):
        if node in temp_mark:
            return
        if node not in visited:
            temp_mark.add(node)
            for dependent in dependency_graph.get(node, []):
                visit(dependent)
            temp_mark.remove(node)
            visited.add(node)
            result.append(node)
    
    for table in tables:
        if table not in visited:
            visit(table)
    
    return result[::-1]

async def drop_tables(conn, tables, dry_run=True):
    """删除未使用的表"""
    if not tables:
        return True
    
    if not dry_run:
        try:
            print("禁用外键约束检查以确保安全删除所有表...")
            await disable_foreign_key_checks(conn)
            
            for table in tables:
                try:
                    async with conn.cursor() as cursor:
                        await cursor.execute(f"DROP TABLE IF EXISTS `{table}`")
                    print(f"已删除表: {table}")
                except Exception as e:
                    print(f"删除表 {table} 时出错: {e}")
                    await conn.rollback()
                    await enable_foreign_key_checks(conn)
                    return False
            
            await conn.commit()
            await enable_foreign_key_checks(conn)
            return True
        except Exception as e:
            print(f"删除表时出错: {e}")
            try:
                await enable_foreign_key_checks(conn)
                await conn.rollback()
            except:
                pass
            return False
    else:
        for table in tables:
            print(f"将删除表: {table}")
        return True

async def drop_columns(conn, column_dict, dry_run=True):
    """删除未使用的列"""
    for table, columns in column_dict.items():
        try:
            for column in columns:
                if dry_run:
                    print(f"将从表 {table} 删除列: {column}")
                else:
                    async with conn.cursor() as cursor:
                        await cursor.execute(f"ALTER TABLE `{table}` DROP COLUMN `{column}`")
                    print(f"已从表 {table} 删除列: {column}")
        except Exception as e:
            print(f"从表 {table} 删除列时出错: {e}")
            if not dry_run:
                await conn.rollback()
                return False
    
    if not dry_run:
        await conn.commit()
    return True

async def cleanup_database(dry_run=True):
    """清理数据库中未使用的表和列"""
    conn = await get_mysql_connection()
    if not conn:
        return False
    
    try:
        unused_tables = await identify_unused_tables(conn)
        unused_columns = await identify_unused_columns(conn)
        
        print("\n=== 数据库清理报告 ===")
        
        if unused_tables:
            print(f"\n发现 {len(unused_tables)} 个未使用的表:")
            for table in unused_tables:
                print(f"  - {table}")
            
            if not dry_run:
                success = await drop_tables(conn, unused_tables, dry_run=False)
                if not success:
                    print("删除表失败，已回滚操作")
                    return False
        else:
            print("\n未发现任何未使用的表")
        
        if unused_columns:
            total_columns = sum(len(cols) for cols in unused_columns.values())
            print(f"\n发现 {total_columns} 个未使用的列:")
            for table, columns in unused_columns.items():
                for column in columns:
                    print(f"  - {table}.{column}")
            
            if not dry_run:
                success = await drop_columns(conn, unused_columns, dry_run=False)
                if not success:
                    print("删除列失败，已回滚操作")
                    return False
        else:
            print("\n未发现任何未使用的列")
        
        if dry_run:
            print("\n这是一次模拟运行，未执行任何实际更改。要执行更改，请使用 --execute 选项。")
        else:
            print("\n数据库清理成功完成！")
        
        return True
    except Exception as e:
        print(f"清理数据库时发生错误: {e}")
        return False
    finally:
        conn.close()

async def remove_orphaned_records(dry_run=True):
    """删除孤立记录（有外键但对应的父记录不存在）"""
    conn = await get_mysql_connection()
    if not conn:
        return False
    
    try:
        async with conn.cursor() as cursor:
            await cursor.execute("""
                SELECT COUNT(*) FROM conversation_messages
                WHERE conversation_id NOT IN (SELECT id FROM conversation_histories)
            """)
            count = (await cursor.fetchone())[0]
            
            if count > 0:
                print(f"\n发现 {count} 条孤立的会话消息记录")
                if not dry_run:
                    await cursor.execute("""
                        DELETE FROM conversation_messages
                        WHERE conversation_id NOT IN (SELECT id FROM conversation_histories)
                    """)
                    print(f"已删除 {count} 条孤立的会话消息记录")
            else:
                print("\n未发现孤立的会话消息记录")
        
        async with conn.cursor() as cursor:
            await cursor.execute("""
                SELECT COUNT(*) FROM role_languages
                WHERE role_id NOT IN (SELECT id FROM roles)
            """)
            count = (await cursor.fetchone())[0]
            
            if count > 0:
                print(f"\n发现 {count} 条孤立的角色语言记录")
                if not dry_run:
                    await cursor.execute("""
                        DELETE FROM role_languages
                        WHERE role_id NOT IN (SELECT id FROM roles)
                    """)
                    print(f"已删除 {count} 条孤立的角色语言记录")
            else:
                print("\n未发现孤立的角色语言记录")
        
        if not dry_run:
            await conn.commit()
        
        return True
    except Exception as e:
        print(f"清理孤立记录时发生错误: {e}")
        if not dry_run:
            await conn.rollback()
        return False
    finally:
        conn.close()

async def drop_specific_tables(conn, tables_to_drop, dry_run=True):
    """删除指定的表，无论它们是否被使用"""
    if dry_run:
        print(f"\n将删除以下指定表:")
        for table in tables_to_drop:
            print(f"  - {table}")
        return True
    
    try:
        await disable_foreign_key_checks(conn)
        
        for table in tables_to_drop:
            try:
                async with conn.cursor() as cursor:
                    await cursor.execute(f"DROP TABLE IF EXISTS `{table}`")
                print(f"已删除表: {table}")
            except Exception as e:
                print(f"删除表 {table} 时出错: {e}")
        
        await conn.commit()
        print(f"指定表已成功删除")
        return True
    except Exception as e:
        await conn.rollback()
        print(f"删除指定表时出错: {e}")
        return False
    finally:
        await enable_foreign_key_checks(conn)

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="数据库清理工具")
    parser.add_argument('--execute', action='store_true', help='执行实际更改（不使用此选项则只进行模拟运行）')
    parser.add_argument('--drop-tables', type=str, help='指定要删除的表，用逗号分隔')
    args = parser.parse_args()
    
    dry_run = not args.execute
    
    async def main():
        if args.drop_tables:
            conn = await get_mysql_connection()
            if not conn:
                return False
            
            try:
                tables_to_drop = [t.strip() for t in args.drop_tables.split(',')]
                await drop_specific_tables(conn, tables_to_drop, dry_run)
            finally:
                conn.close()
        else:
            result = await cleanup_database(dry_run)
            if result:
                await remove_orphaned_records(dry_run)
    
    asyncio.run(main())