import pymysql
from pymysql.cursors import DictCursor
from typing import Optional, Dict, List, Any

class DBSchemaExporter:
    """
    MySQL数据库表结构导出工具，支持提取表和字段注释
    
    使用示例:
    >>> exporter = DBSchemaExporter("localhost", "user", "password", "dbname")
    >>> schema = exporter.export(include_comments=True)
    >>> print(schema)
    >>> exporter.save_to_file("schema.txt")
    """
    
    def __init__(
        self, 
        host: str, 
        user: str, 
        password: str, 
        database: str, 
        port: int = 3306, 
        charset: str = 'utf8mb4'
    ):
        """初始化数据库连接配置"""
        self.config = {
            'host': host,
            'user': user,
            'password': password,
            'database': database,
            'port': port,
            'charset': charset,
            'cursorclass': DictCursor
        }
        self.schema = None
    
    def export(self, include_comments: bool = True, 
               include_tables: Optional[List[str]] = None, 
               exclude_tables: Optional[List[str]] = None) -> Optional[str]:
        """
        导出数据库表结构，支持包含注释
        
        Args:
            include_comments: 是否包含表和字段注释
            include_tables: 只包含指定的表
            exclude_tables: 排除指定的表
            
        Returns:
            包含所有表结构的字符串，失败时返回None
        """
        try:
            # 获取表和字段注释信息
            table_comments = {}
            column_comments = {}
            
            if include_comments:
                table_comments, column_comments = self._fetch_comments()
            
            with pymysql.connect(**self.config) as connection:
                with connection.cursor() as cursor:
                    # 获取所有表名
                    cursor.execute("SHOW TABLES")
                    all_tables = [row[f"Tables_in_{self.config['database']}"] for row in cursor.fetchall()]
                    
                    # 过滤表
                    if include_tables:
                        tables = [t for t in all_tables if t in include_tables]
                    else:
                        tables = all_tables
                    
                    if exclude_tables:
                        tables = [t for t in tables if t not in exclude_tables]
                    
                    # 获取每个表的结构
                    table_schemas = []
                    for table in tables:
                        # 获取表的创建语句
                        cursor.execute(f"SHOW CREATE TABLE `{table}`")
                        create_table = cursor.fetchone()["Create Table"]
                        
                        # 提取表结构部分
                        table_structure = create_table.split(f"CREATE TABLE `{table}`", 1)[1].strip()
                        
                        # 添加表注释（如果有）
                        if include_comments and table in table_comments:
                            comment_line = f"-- 表注释: {table_comments[table]}"
                            table_schemas.append(f"{comment_line}\n{table} (\n{table_structure}\n)\n")
                        else:
                            table_schemas.append(f"{table} (\n{table_structure}\n)\n")
            
            self.schema = "\n\n".join(table_schemas)
            return self.schema
            
        except Exception as e:
            print(f"导出失败: {e}")
            return None
    
    def _fetch_comments(self) -> tuple:
        """获取所有表和字段的注释信息"""
        table_comments = {}
        column_comments = {}
        
        with pymysql.connect(**self.config) as connection:
            with connection.cursor() as cursor:
                # 获取表注释
                cursor.execute("""
                    SELECT TABLE_NAME, TABLE_COMMENT 
                    FROM INFORMATION_SCHEMA.TABLES 
                    WHERE TABLE_SCHEMA = %s
                """, (self.config['database'],))
                
                for row in cursor.fetchall():
                    table_comments[row['TABLE_NAME']] = row['TABLE_COMMENT']
                
                # 获取字段注释
                cursor.execute("""
                    SELECT TABLE_NAME, COLUMN_NAME, COLUMN_COMMENT 
                    FROM INFORMATION_SCHEMA.COLUMNS 
                    WHERE TABLE_SCHEMA = %s
                """, (self.config['database'],))
                
                for row in cursor.fetchall():
                    table = row['TABLE_NAME']
                    column = row['COLUMN_NAME']
                    comment = row['COLUMN_COMMENT']
                    
                    if table not in column_comments:
                        column_comments[table] = {}
                    column_comments[table][column] = comment
        
        return table_comments, column_comments
    
    def save_to_file(self, file_path: str, overwrite: bool = False) -> bool:
        """
        将导出的表结构保存到文件
        
        Args:
            file_path: 文件路径
            overwrite: 是否覆盖已存在的文件
            
        Returns:
            保存成功返回True，失败返回False
        """
        if self.schema is None:
            print("没有导出的表结构，请先调用export()方法")
            return False
        
        try:
            # 检查文件是否存在
            import os
            if os.path.exists(file_path) and not overwrite:
                print(f"文件已存在: {file_path}，设置overwrite=True以覆盖")
                return False
                
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(self.schema)
            return True
            
        except Exception as e:
            print(f"保存失败: {e}")
            return False
    
    def get_table_details(self) -> Optional[Dict[str, List[Dict[str, Any]]]]:
        """
        获取每个表的详细列信息，包括注释
        
        Returns:
            包含表和列信息的字典，格式为:
            {
                'table1': [{'Field': 'id', 'Type': 'int', 'Comment': '主键'}, ...],
                'table2': [...]
            }
        """
        try:
            table_details = {}
            with pymysql.connect(**self.config) as connection:
                with connection.cursor() as cursor:
                    # 获取所有表名
                    cursor.execute("SHOW TABLES")
                    tables = [row[f"Tables_in_{self.config['database']}"] for row in cursor.fetchall()]
                    
                    for table in tables:
                        cursor.execute(f"SHOW FULL COLUMNS FROM `{table}`")
                        columns = cursor.fetchall()
                        table_details[table] = columns
            
            return table_details
            
        except Exception as e:
            print(f"获取表详细信息失败: {e}")
            return None


# 使用示例
if __name__ == "__main__":
    # 配置数据库连接
    exporter = DBSchemaExporter(
        host="localhost",
        user="your_username",
        password="your_password",
        database="your_database"
    )
    
    # 示例1: 导出包含注释的完整表结构
    schema_with_comments = exporter.export(include_comments=True)
    if schema_with_comments:
        print("带注释的完整表结构:")
        print(schema_with_comments)
        exporter.save_to_file("full_schema_with_comments.txt", overwrite=True)
    
    # 示例2: 导出特定表的结构（不带注释）
    specific_schema = exporter.export(
        include_comments=False,
        include_tables=["users", "orders"]
    )
    if specific_schema:
        print("\n特定表的结构（不带注释）:")
        print(specific_schema)
        exporter.save_to_file("specific_schema.txt", overwrite=True)
    
    # 示例3: 获取详细的列信息（包括注释）
    table_details = exporter.get_table_details()
    if table_details:
        print("\n表列详细信息:")
        for table, columns in table_details.items():
            print(f"\n{table}:")
            for column in columns[:3]:  # 只显示每个表的前3列
                comment = column.get('Comment', '')
                print(f"  - {column['Field']} ({column['Type']}){' - ' + comment if comment else ''}")