import pymysql


class DatabaseInfo:
    def __init__(self, host, user, port, password, database):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.connection = None
    
    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()

    def connect(self):
        try:
            self.connection = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database
            )
        except pymysql.Error as e:
            print(f"数据库连接失败: {e}")

    def disconnect(self):
        if self.connection:
            self.connection.close()

    def get_all_tables(self):
        if not self.connection:
            self.connect()
        try:
            with self.connection.cursor() as cursor:
                cursor.execute("SHOW TABLE STATUS")
                tables = cursor.fetchall()
                table_info_list = []
                for table in tables:
                    table_info = {
                        "tableName": table[0],
                        "tableComment": table[17]  # 表注释信息在第 19 列（索引为 18）
                    }
                    table_info_list.append(table_info)
            return table_info_list
        except pymysql.Error as e:
            print(f"获取表名失败: {e}")
            return []
    
    def get_table_info(self, table_name):
        if not self.connection:
            self.connect()
        try:
            with self.connection.cursor() as cursor:
                cursor.execute("SHOW TABLE STATUS LIKE %s", (table_name,))
                table = cursor.fetchone()
                table_info = {
                    "tableName": table[0],
                    "tableComment": table[17]  # 表注释信息在第 19 列（索引为 18）
                }
            return table_info
        except pymysql.Error as e:
            print(f"获取表名失败: {e}")
            return {}
    
    def parse_type_and_validation(self, column_type):
        # 简单的类型解析，实际可能需要更复杂的逻辑
        if "int" in column_type.lower():
            return dict(type="int", validation="int")
        elif "varchar" in column_type.lower():
            number = column_type.split("(")[1].split(")")[0].strip()
            return dict(type="str", validations=[
                dict(type="str", min=1, max=int(number))
            ])
        # elif "datetime" in column_type.lower():
        #     return dict(type="str", validations=[
        #         dict(type="str", format="%Y-%m-%d %H:%M:%S")
        #     ])
        else:
            return dict(type=column_type)

    def get_table_structure(self, table_name):
        if not self.connection:
            self.connect()
        try:
            table_info = self.get_table_info(table_name)
            with self.connection.cursor() as cursor:
                # 使用 SHOW FULL COLUMNS FROM 获取包含注释的列信息
                cursor.execute(f"SHOW FULL COLUMNS FROM {table_name}")
                columns = cursor.fetchall()
                structure = []
                for column in columns:
                    column_info = {
                        "name": column[0],
                        "type": column[1],
                        "required": column[2] is not None,
                        "primary_key": column[4] == 'PRI',
                        # "default": column[4],
                        "description": column[8]  # 注释信息在第 9 列（索引为 8）
                    }
                    column_info.update(self.parse_type_and_validation(column_info["type"]))  # 添加解析后的类型和验证信息
                    structure.append(column_info)
            table_info['fields'] = structure
            return table_info
        except pymysql.Error as e:
            print(f"获取表 {table_name} 结构失败: {e}")
            return {'fields': []}


if __name__ == "__main__":
    # 请根据实际情况修改数据库连接信息
    db_info = DatabaseInfo(
        host="lt.330770.xyz",
        port=3307,
        user="root",
        password="rap_sky",
        database="rpa"
    )
    tables = db_info.get_all_tables()
    print("所有表名:")
    for table in tables:
        print(table)
        tableName = table.get('TableName')
        structure = db_info.get_table_structure(tableName)
        print(f"表 {tableName} 的结构信息:")
        for column in structure:
            print(column)
    db_info.disconnect()