#!/usr/bin/env python3
"""
将 taiji 数据库的表结构和数据复制到 postgres 数据库

使用方法:
    python scripts/copy_database.py

注意: 需要安装 psycopg2-binary
    pip install psycopg2-binary
"""

import psycopg2
from urllib.parse import quote_plus
import sys

# 数据库连接配置
DB_HOST = "taijipda.postgres.database.azure.com"
DB_USER = "taiji"
DB_PASSWORD = "By@123456."
DB_PORT = 5432

# 源数据库和目标数据库
SOURCE_DB = "taiji"
TARGET_DB = "postgres"


def get_connection(database):
    """获取数据库连接"""
    return psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        database=database,
        sslmode="require"
    )


def get_all_tables(conn):
    """获取所有用户表"""
    cursor = conn.cursor()
    cursor.execute("""
        SELECT table_name 
        FROM information_schema.tables 
        WHERE table_schema = 'public' 
        AND table_type = 'BASE TABLE'
        ORDER BY table_name
    """)
    tables = [row[0] for row in cursor.fetchall()]
    cursor.close()
    return tables


def get_table_ddl(conn, table_name):
    """获取表的 DDL 语句"""
    cursor = conn.cursor()
    
    # 获取列定义
    cursor.execute("""
        SELECT 
            column_name,
            data_type,
            character_maximum_length,
            numeric_precision,
            numeric_scale,
            is_nullable,
            column_default,
            udt_name
        FROM information_schema.columns
        WHERE table_schema = 'public' AND table_name = %s
        ORDER BY ordinal_position
    """, (table_name,))
    
    columns = cursor.fetchall()
    
    if not columns:
        cursor.close()
        return None
    
    # 构建列定义
    column_defs = []
    for col in columns:
        col_name, data_type, char_max_len, num_precision, num_scale, is_nullable, col_default, udt_name = col
        
        # 处理数据类型
        if data_type == 'character varying':
            if char_max_len:
                type_str = f"VARCHAR({char_max_len})"
            else:
                type_str = "VARCHAR"
        elif data_type == 'character':
            type_str = f"CHAR({char_max_len})" if char_max_len else "CHAR"
        elif data_type == 'numeric':
            if num_precision and num_scale:
                type_str = f"NUMERIC({num_precision},{num_scale})"
            elif num_precision:
                type_str = f"NUMERIC({num_precision})"
            else:
                type_str = "NUMERIC"
        elif data_type == 'ARRAY':
            type_str = f"{udt_name.lstrip('_')}[]"
        elif data_type == 'USER-DEFINED':
            type_str = udt_name
        else:
            type_str = data_type.upper()
        
        # 构建列定义
        col_def = f'    "{col_name}" {type_str}'
        
        if is_nullable == 'NO':
            col_def += " NOT NULL"
        
        if col_default:
            col_def += f" DEFAULT {col_default}"
        
        column_defs.append(col_def)
    
    # 获取主键约束
    cursor.execute("""
        SELECT kcu.column_name
        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
        WHERE tc.constraint_type = 'PRIMARY KEY'
        AND tc.table_schema = 'public'
        AND tc.table_name = %s
        ORDER BY kcu.ordinal_position
    """, (table_name,))
    
    pk_columns = [row[0] for row in cursor.fetchall()]
    
    if pk_columns:
        pk_def = f'    PRIMARY KEY ("{"\", \"".join(pk_columns)}")'
        column_defs.append(pk_def)
    
    ddl = f'CREATE TABLE IF NOT EXISTS "{table_name}" (\n'
    ddl += ",\n".join(column_defs)
    ddl += "\n);"
    
    cursor.close()
    return ddl


def get_indexes(conn, table_name):
    """获取表的索引"""
    cursor = conn.cursor()
    cursor.execute("""
        SELECT indexdef
        FROM pg_indexes
        WHERE schemaname = 'public'
        AND tablename = %s
        AND indexname NOT LIKE '%%_pkey'
    """, (table_name,))
    
    indexes = [row[0] for row in cursor.fetchall()]
    cursor.close()
    return indexes


def get_sequences(conn):
    """获取所有序列"""
    cursor = conn.cursor()
    cursor.execute("""
        SELECT sequence_name
        FROM information_schema.sequences
        WHERE sequence_schema = 'public'
    """)
    sequences = [row[0] for row in cursor.fetchall()]
    cursor.close()
    return sequences


def get_sequence_value(conn, sequence_name):
    """获取序列当前值"""
    cursor = conn.cursor()
    try:
        cursor.execute(f'SELECT last_value FROM "{sequence_name}"')
        value = cursor.fetchone()[0]
    except:
        value = 1
    cursor.close()
    return value


def copy_table_data(source_conn, target_conn, table_name):
    """复制表数据"""
    source_cursor = source_conn.cursor()
    target_cursor = target_conn.cursor()
    
    # 获取列名
    source_cursor.execute("""
        SELECT column_name
        FROM information_schema.columns
        WHERE table_schema = 'public' AND table_name = %s
        ORDER BY ordinal_position
    """, (table_name,))
    
    columns = [row[0] for row in source_cursor.fetchall()]
    
    if not columns:
        source_cursor.close()
        target_cursor.close()
        return 0
    
    # 获取数据
    columns_str = ', '.join([f'"{c}"' for c in columns])
    source_cursor.execute(f'SELECT {columns_str} FROM "{table_name}"')
    rows = source_cursor.fetchall()
    
    if not rows:
        source_cursor.close()
        target_cursor.close()
        return 0
    
    # 插入数据
    placeholders = ', '.join(['%s'] * len(columns))
    insert_sql = f'INSERT INTO "{table_name}" ({columns_str}) VALUES ({placeholders}) ON CONFLICT DO NOTHING'
    
    for row in rows:
        try:
            target_cursor.execute(insert_sql, row)
        except Exception as e:
            print(f"    警告: 插入数据失败 - {e}")
    
    target_conn.commit()
    
    source_cursor.close()
    target_cursor.close()
    
    return len(rows)


def drop_all_tables(conn):
    """删除目标数据库中的所有表"""
    cursor = conn.cursor()
    
    # 获取所有表
    cursor.execute("""
        SELECT table_name
        FROM information_schema.tables
        WHERE table_schema = 'public'
        AND table_type = 'BASE TABLE'
    """)
    tables = [row[0] for row in cursor.fetchall()]
    
    if tables:
        # 禁用外键检查并删除所有表
        for table in tables:
            try:
                cursor.execute(f'DROP TABLE IF EXISTS "{table}" CASCADE')
                print(f"    ✓ 已删除表: {table}")
            except Exception as e:
                print(f"    ✗ 删除表 {table} 失败: {e}")
        conn.commit()
    
    cursor.close()
    return len(tables)


def main():
    print("=" * 60)
    print("PostgreSQL 数据库复制工具")
    print(f"源数据库: {SOURCE_DB}")
    print(f"目标数据库: {TARGET_DB}")
    print("=" * 60)
    
    # 连接源数据库
    print("\n[1] 连接源数据库...")
    try:
        source_conn = get_connection(SOURCE_DB)
        print(f"    ✓ 成功连接到 {SOURCE_DB}")
    except Exception as e:
        print(f"    ✗ 连接失败: {e}")
        sys.exit(1)
    
    # 连接目标数据库
    print("\n[2] 连接目标数据库...")
    try:
        target_conn = get_connection(TARGET_DB)
        print(f"    ✓ 成功连接到 {TARGET_DB}")
    except Exception as e:
        print(f"    ✗ 连接失败: {e}")
        source_conn.close()
        sys.exit(1)
    
    # 删除目标数据库中的旧表
    print("\n[3] 清理目标数据库旧表...")
    dropped_count = drop_all_tables(target_conn)
    print(f"    共删除 {dropped_count} 个旧表")
    
    # 获取所有表
    print("\n[4] 获取源数据库表列表...")
    tables = get_all_tables(source_conn)
    print(f"    找到 {len(tables)} 个表:")
    for t in tables:
        print(f"      - {t}")
    
    # 复制表结构
    print("\n[5] 复制表结构...")
    target_cursor = target_conn.cursor()
    
    for table in tables:
        print(f"    处理表: {table}")
        
        # 获取并执行 DDL
        ddl = get_table_ddl(source_conn, table)
        if ddl:
            try:
                target_cursor.execute(ddl)
                target_conn.commit()
                print(f"      ✓ 表结构已创建")
            except Exception as e:
                target_conn.rollback()
                if "already exists" in str(e):
                    print(f"      ○ 表已存在，跳过创建")
                else:
                    print(f"      ✗ 创建失败: {e}")
        
        # 获取并创建索引
        indexes = get_indexes(source_conn, table)
        for idx in indexes:
            try:
                # 修改索引名以避免冲突
                target_cursor.execute(idx)
                target_conn.commit()
                print(f"      ✓ 索引已创建")
            except Exception as e:
                target_conn.rollback()
                if "already exists" in str(e):
                    print(f"      ○ 索引已存在，跳过")
                else:
                    print(f"      ✗ 索引创建失败: {e}")
    
    target_cursor.close()
    
    # 复制数据
    print("\n[6] 复制表数据...")
    for table in tables:
        print(f"    复制表: {table}")
        try:
            count = copy_table_data(source_conn, target_conn, table)
            print(f"      ✓ 已复制 {count} 行数据")
        except Exception as e:
            print(f"      ✗ 复制失败: {e}")
    
    # 更新序列
    print("\n[7] 同步序列值...")
    sequences = get_sequences(source_conn)
    target_cursor = target_conn.cursor()
    
    for seq in sequences:
        try:
            value = get_sequence_value(source_conn, seq)
            target_cursor.execute(f'SELECT setval(\'{seq}\', {value}, true)')
            target_conn.commit()
            print(f"    ✓ 序列 {seq} 设置为 {value}")
        except Exception as e:
            target_conn.rollback()
            print(f"    ✗ 序列 {seq} 同步失败: {e}")
    
    target_cursor.close()
    
    # 关闭连接
    source_conn.close()
    target_conn.close()
    
    print("\n" + "=" * 60)
    print("数据库复制完成!")
    print("=" * 60)


if __name__ == "__main__":
    main()