import pymysql
import sys
import os
import re

# 默认数据库配置
DB_CONFIG = {
    "host": "rm-bp1ux6tuk49er80t9xo.mysql.rds.aliyuncs.com",
    "port": 3306,
    "user": "ylt",
    "password": "Ycharge666",
    "database": "yltcharge",
    "charset": "utf8mb4"
}

# MySQL 到 Doris 类型映射
TYPE_MAPPING = {
    "int": "INT",
    "tinyint": "TINYINT",
    "smallint": "SMALLINT",
    "mediumint": "INT",
    "bigint": "BIGINT",
    "float": "FLOAT",
    "double": "DOUBLE",
    "decimal": "DECIMAL",
    "char": "CHAR",
    "varchar": "VARCHAR",
    "text": "STRING",
    "mediumtext": "STRING",
    "longtext": "STRING",
    "date": "DATE",
    "datetime": "DATETIME",
    "timestamp": "DATETIME",
    "bit": "BOOLEAN"
}

def get_table_schema(table_name):
    connection = pymysql.connect(**DB_CONFIG)
    try:
        with connection.cursor(pymysql.cursors.DictCursor) as cursor:
            # 获取列信息
            sql = """
            SELECT COLUMN_NAME, DATA_TYPE, COLUMN_TYPE, IS_NULLABLE, COLUMN_DEFAULT, COLUMN_COMMENT, EXTRA
            FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
            ORDER BY ORDINAL_POSITION
            """
            cursor.execute(sql, (DB_CONFIG['database'], table_name))
            columns = cursor.fetchall()
            
            # 获取主键信息
            cursor.execute(f"SHOW KEYS FROM {table_name} WHERE Key_name = 'PRIMARY'")
            pk_info = cursor.fetchall()
            pk_columns = [info['Column_name'] for info in pk_info]
            
            return columns, pk_columns
    finally:
        connection.close()

def generate_mysql_ddl(table_name, columns, pk_columns):
    ddl = [f"CREATE TABLE `{table_name}` ("]
    col_defs = []
    for col in columns:
        line = f"  `{col['COLUMN_NAME']}` {col['COLUMN_TYPE']}"
        if col['IS_NULLABLE'] == 'NO':
            line += " NOT NULL"
        if col['COLUMN_DEFAULT'] is not None:
            if isinstance(col['COLUMN_DEFAULT'], str):
                line += f" DEFAULT '{col['COLUMN_DEFAULT']}'"
            else:
                line += f" DEFAULT {col['COLUMN_DEFAULT']}"
        if col['EXTRA']:
            line += f" {col['EXTRA']}"
        if col['COLUMN_COMMENT']:
            line += f" COMMENT '{col['COLUMN_COMMENT']}'"
        col_defs.append(line)
    
    if pk_columns:
        col_defs.append(f"  PRIMARY KEY ({', '.join([f'`{c}`' for c in pk_columns])})")
    
    ddl.append(",\n".join(col_defs))
    ddl.append(") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;")
    return "\n".join(ddl)

def generate_doris_ddl(table_name, columns, pk_columns):
    # Doris 建议主键放在前面
    ddl = [f"CREATE TABLE `{table_name}` ("]
    col_defs = []
    
    # 确定 Unique Key 列表（如果没有主键，则取第一个字段，但用户要求唯一主键）
    keys = pk_columns if pk_columns else [columns[0]['COLUMN_NAME']]
    
    for col in columns:
        data_type = col['DATA_TYPE'].lower()
        doris_type = TYPE_MAPPING.get(data_type, "STRING")
        
        # 处理长度信息
        if doris_type in ["VARCHAR", "CHAR"]:
            # 从 COLUMN_TYPE 中提取长度，如 varchar(255)
            match = re.search(r'\((\d+)\)', col['COLUMN_TYPE'])
            if match:
                length = match.group(1)
                # Doris VARCHAR 最大支持 65533
                if int(length) > 65533:
                    doris_type = "STRING"
                else:
                    doris_type = f"{doris_type}({length})"
            else:
                doris_type = f"{doris_type}(255)"
        elif doris_type == "DECIMAL":
            match = re.search(r'\((\d+,\d+)\)', col['COLUMN_TYPE'])
            if match:
                doris_type = f"DECIMAL({match.group(1)})"

        line = f"  `{col['COLUMN_NAME']}` {doris_type}"
        if col['IS_NULLABLE'] == 'NO':
            line += " NOT NULL"
        # Doris 默认值处理较简化，暂时只处理简单字符串和数字
        if col['COLUMN_DEFAULT'] is not None:
            if isinstance(col['COLUMN_DEFAULT'], str) and '(' not in col['COLUMN_DEFAULT']:
                 line += f" DEFAULT '{col['COLUMN_DEFAULT']}'"
        
        if col['COLUMN_COMMENT']:
            line += f" COMMENT '{col['COLUMN_COMMENT']}'"
        col_defs.append(line)
    
    ddl.append(",\n".join(col_defs))
    ddl.append(")")
    ddl.append("UNIQUE KEY(" + ", ".join([f"`{k}`" for k in keys]) + ")")
    ddl.append(f"DISTRIBUTED BY HASH(`{keys[0]}`) BUCKETS 1")
    ddl.append('PROPERTIES (')
    ddl.append('  "replication_num" = "1"')
    ddl.append(');')
    
    return "\n".join(ddl)

def main():
    output_dir = os.path.join(os.path.dirname(__file__), "output")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 自动获取表名列表
    if len(sys.argv) < 2:
        print("未指定表名，将自动从 json 目录获取...")
        json_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "json"))
        if not os.path.exists(json_dir):
            print(f"错误: 找不到目录 {json_dir}")
            sys.exit(1)
        
        tables = [f.replace('.json', '') for f in os.listdir(json_dir) if f.endswith('.json')]
        if not tables:
            print(f"错误: 在 {json_dir} 中没有找到 .json 文件")
            sys.exit(1)
        print(f"发现 {len(tables)} 张表: {', '.join(tables)}")
    else:
        tables = sys.argv[1:]
    
    for table in tables:
        print(f"正在获取表 {table} 的结构...")
        try:
            columns, pk_columns = get_table_schema(table)
            
            mysql_ddl = generate_mysql_ddl(table, columns, pk_columns)
            doris_ddl = generate_doris_ddl(table, columns, pk_columns)
            
            mysql_file = os.path.join(output_dir, f"{table}_mysql.sql")
            doris_file = os.path.join(output_dir, f"{table}_doris.sql")
            
            with open(mysql_file, "w", encoding="utf-8") as f:
                f.write(mysql_ddl)
            
            with open(doris_file, "w", encoding="utf-8") as f:
                f.write(doris_ddl)
            
            print(f"  - MySQL DDL 已生成: {mysql_file}")
            print(f"  - Doris DDL 已生成: {doris_file}")
            
        except Exception as e:
            print(f"  - 处理表 {table} 时出错: {e}")

if __name__ == "__main__":
    main()