#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
数据库迁移管理器
提供安全的数据库结构更新和数据迁移功能
"""

import os
import re
import pymysql
import hashlib
from datetime import datetime
from typing import List, Dict, Optional, Tuple
import logging
import json

class DatabaseMigrationManager:
    """数据库迁移管理器"""
    
    def __init__(self, db_config: Dict, migrations_dir: str = "migrations"):
        """
        初始化迁移管理器
        
        Args:
            db_config: 数据库配置
            migrations_dir: 迁移文件目录
        """
        self.db_config = db_config
        self.migrations_dir = migrations_dir
        self.logger = self._setup_logger()
        
        # 确保迁移目录存在
        os.makedirs(migrations_dir, exist_ok=True)
        
        # 初始化迁移历史表
        self._init_migration_table()
    
    def _setup_logger(self) -> logging.Logger:
        """设置日志记录器"""
        logger = logging.getLogger('migration_manager')
        logger.setLevel(logging.INFO)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def _get_connection(self):
        """获取数据库连接"""
        return pymysql.connect(**self.db_config)
    
    def _init_migration_table(self):
        """初始化迁移历史表"""
        try:
            connection = self._get_connection()
            cursor = connection.cursor()
            
            # 创建迁移历史表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS migration_history (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    version VARCHAR(50) NOT NULL UNIQUE,
                    filename VARCHAR(255) NOT NULL,
                    checksum VARCHAR(64) NOT NULL,
                    executed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    execution_time_ms INT,
                    status ENUM('SUCCESS', 'FAILED', 'ROLLED_BACK') DEFAULT 'SUCCESS',
                    error_message TEXT,
                    INDEX idx_version (version),
                    INDEX idx_executed_at (executed_at)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """)
            
            connection.commit()
            self.logger.info("✅ 迁移历史表初始化完成")
            
        except Exception as e:
            self.logger.error(f"❌ 初始化迁移历史表失败: {e}")
            raise
        finally:
            if 'connection' in locals():
                connection.close()
    
    def create_migration(self, name: str, description: str = "") -> str:
        """
        创建新的迁移文件
        
        Args:
            name: 迁移名称
            description: 迁移描述
            
        Returns:
            创建的迁移文件路径
        """
        # 生成版本号（时间戳格式）
        version = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 清理文件名
        clean_name = re.sub(r'[^\w\-_]', '_', name.lower())
        filename = f"{version}_{clean_name}.sql"
        filepath = os.path.join(self.migrations_dir, filename)
        
        # 创建迁移文件模板
        template = f"""-- Migration: {name}
-- Version: {version}
-- Description: {description}
-- Created: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

-- ==========================================
-- UP Migration (执行迁移)
-- ==========================================

-- 在这里添加你的迁移SQL语句
-- 例如：
-- ALTER TABLE users ADD COLUMN email VARCHAR(255);
-- CREATE INDEX idx_users_email ON users(email);

-- ==========================================
-- DOWN Migration (回滚迁移)
-- ==========================================
-- 注意：DOWN部分用于回滚，请确保操作的可逆性

-- 在这里添加回滚SQL语句
-- 例如：
-- DROP INDEX idx_users_email ON users;
-- ALTER TABLE users DROP COLUMN email;

-- ==========================================
-- 验证脚本 (可选)
-- ==========================================
-- 在这里添加验证SQL，确保迁移成功
-- 例如：
-- SELECT COUNT(*) FROM information_schema.columns 
-- WHERE table_name = 'users' AND column_name = 'email';
"""
        
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(template)
        
        self.logger.info(f"✅ 创建迁移文件: {filepath}")
        return filepath
    
    def get_pending_migrations(self) -> List[str]:
        """获取待执行的迁移文件"""
        try:
            connection = self._get_connection()
            cursor = connection.cursor()
            
            # 获取已执行的迁移版本
            cursor.execute("SELECT version FROM migration_history WHERE status = 'SUCCESS'")
            executed_versions = {row[0] for row in cursor.fetchall()}
            
            # 获取所有迁移文件
            migration_files = []
            for filename in os.listdir(self.migrations_dir):
                if filename.endswith('.sql'):
                    version = filename.split('_')[0] + '_' + filename.split('_')[1]
                    if version not in executed_versions:
                        migration_files.append(filename)
            
            # 按版本排序
            migration_files.sort()
            return migration_files
            
        except Exception as e:
            self.logger.error(f"❌ 获取待执行迁移失败: {e}")
            raise
        finally:
            if 'connection' in locals():
                connection.close()
    
    def _calculate_checksum(self, content: str) -> str:
        """计算文件内容的校验和"""
        return hashlib.sha256(content.encode('utf-8')).hexdigest()
    
    def _parse_migration_file(self, filepath: str) -> Dict:
        """解析迁移文件"""
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 提取UP和DOWN部分
        up_match = re.search(r'-- UP Migration.*?\n(.*?)-- DOWN Migration', content, re.DOTALL)
        down_match = re.search(r'-- DOWN Migration.*?\n(.*?)-- ==========================================', content, re.DOTALL)
        
        up_sql = up_match.group(1).strip() if up_match else ""
        down_sql = down_match.group(1).strip() if down_match else ""
        
        # 清理SQL（移除注释和空行）
        up_sql = self._clean_sql(up_sql)
        down_sql = self._clean_sql(down_sql)
        
        return {
            'content': content,
            'up_sql': up_sql,
            'down_sql': down_sql,
            'checksum': self._calculate_checksum(content)
        }
    
    def _clean_sql(self, sql: str) -> str:
        """清理SQL语句"""
        lines = []
        for line in sql.split('\n'):
            line = line.strip()
            if line and not line.startswith('--'):
                lines.append(line)
        return '\n'.join(lines)

    def _create_backup(self) -> str:
        """创建数据库备份"""
        backup_dir = "database_backups"
        os.makedirs(backup_dir, exist_ok=True)

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        backup_file = os.path.join(backup_dir, f"migration_backup_{timestamp}.sql")

        try:
            # 使用mysqldump创建备份
            import subprocess
            cmd = [
                'mysqldump',
                f'--host={self.db_config["host"]}',
                f'--port={self.db_config["port"]}',
                f'--user={self.db_config["user"]}',
                f'--password={self.db_config["password"]}',
                '--single-transaction',
                '--routines',
                '--triggers',
                self.db_config['database']
            ]

            with open(backup_file, 'w') as f:
                subprocess.run(cmd, stdout=f, check=True)

            self.logger.info(f"✅ 数据库备份完成: {backup_file}")
            return backup_file

        except Exception as e:
            self.logger.error(f"❌ 数据库备份失败: {e}")
            raise

    def execute_migration(self, filename: str, dry_run: bool = False) -> bool:
        """
        执行单个迁移文件

        Args:
            filename: 迁移文件名
            dry_run: 是否为试运行模式

        Returns:
            执行是否成功
        """
        filepath = os.path.join(self.migrations_dir, filename)
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"迁移文件不存在: {filepath}")

        # 解析迁移文件
        migration_data = self._parse_migration_file(filepath)
        version = filename.split('_')[0] + '_' + filename.split('_')[1]

        if dry_run:
            self.logger.info(f"🔍 试运行模式 - 迁移: {filename}")
            self.logger.info(f"SQL内容:\n{migration_data['up_sql']}")
            return True

        # 创建备份
        backup_file = self._create_backup()

        connection = None
        start_time = datetime.now()

        try:
            connection = self._get_connection()
            connection.begin()  # 开始事务
            cursor = connection.cursor()

            self.logger.info(f"🚀 执行迁移: {filename}")

            # 执行UP SQL
            if migration_data['up_sql']:
                # 分割SQL语句（处理多个语句）
                statements = [stmt.strip() for stmt in migration_data['up_sql'].split(';') if stmt.strip()]

                for stmt in statements:
                    self.logger.debug(f"执行SQL: {stmt[:100]}...")
                    cursor.execute(stmt)

            # 记录迁移历史
            execution_time = int((datetime.now() - start_time).total_seconds() * 1000)
            cursor.execute("""
                INSERT INTO migration_history
                (version, filename, checksum, execution_time_ms, status)
                VALUES (%s, %s, %s, %s, 'SUCCESS')
            """, (version, filename, migration_data['checksum'], execution_time))

            connection.commit()
            self.logger.info(f"✅ 迁移执行成功: {filename} (耗时: {execution_time}ms)")
            return True

        except Exception as e:
            if connection:
                connection.rollback()

            # 记录失败信息
            try:
                cursor.execute("""
                    INSERT INTO migration_history
                    (version, filename, checksum, status, error_message)
                    VALUES (%s, %s, %s, 'FAILED', %s)
                """, (version, filename, migration_data['checksum'], str(e)))
                connection.commit()
            except:
                pass

            self.logger.error(f"❌ 迁移执行失败: {filename} - {e}")
            return False

        finally:
            if connection:
                connection.close()

    def migrate(self, target_version: Optional[str] = None, dry_run: bool = False) -> bool:
        """
        执行所有待执行的迁移

        Args:
            target_version: 目标版本（可选）
            dry_run: 是否为试运行模式

        Returns:
            是否全部执行成功
        """
        pending_migrations = self.get_pending_migrations()

        if not pending_migrations:
            self.logger.info("✅ 没有待执行的迁移")
            return True

        # 过滤到目标版本
        if target_version:
            filtered_migrations = []
            for migration in pending_migrations:
                version = migration.split('_')[0] + '_' + migration.split('_')[1]
                filtered_migrations.append(migration)
                if version == target_version:
                    break
            pending_migrations = filtered_migrations

        self.logger.info(f"📋 待执行迁移数量: {len(pending_migrations)}")

        success_count = 0
        for migration in pending_migrations:
            if self.execute_migration(migration, dry_run):
                success_count += 1
            else:
                self.logger.error(f"❌ 迁移失败，停止执行: {migration}")
                break

        if success_count == len(pending_migrations):
            self.logger.info(f"🎉 所有迁移执行完成! ({success_count}/{len(pending_migrations)})")
            return True
        else:
            self.logger.error(f"⚠️ 部分迁移失败: ({success_count}/{len(pending_migrations)})")
            return False

    def rollback_migration(self, version: str) -> bool:
        """
        回滚指定版本的迁移

        Args:
            version: 要回滚的版本号

        Returns:
            回滚是否成功
        """
        try:
            connection = self._get_connection()
            cursor = connection.cursor()

            # 检查迁移是否存在且已执行
            cursor.execute("""
                SELECT filename, checksum FROM migration_history
                WHERE version = %s AND status = 'SUCCESS'
            """, (version,))

            result = cursor.fetchone()
            if not result:
                self.logger.error(f"❌ 未找到已执行的迁移版本: {version}")
                return False

            filename, checksum = result
            filepath = os.path.join(self.migrations_dir, filename)

            if not os.path.exists(filepath):
                self.logger.error(f"❌ 迁移文件不存在: {filepath}")
                return False

            # 解析迁移文件
            migration_data = self._parse_migration_file(filepath)

            # 验证校验和
            if migration_data['checksum'] != checksum:
                self.logger.error(f"❌ 迁移文件校验和不匹配，可能已被修改: {filename}")
                return False

            if not migration_data['down_sql']:
                self.logger.error(f"❌ 迁移文件缺少DOWN部分，无法回滚: {filename}")
                return False

            # 创建备份
            backup_file = self._create_backup()

            connection.begin()  # 开始事务

            self.logger.info(f"🔄 回滚迁移: {filename}")

            # 执行DOWN SQL
            statements = [stmt.strip() for stmt in migration_data['down_sql'].split(';') if stmt.strip()]

            for stmt in statements:
                self.logger.debug(f"执行回滚SQL: {stmt[:100]}...")
                cursor.execute(stmt)

            # 更新迁移历史状态
            cursor.execute("""
                UPDATE migration_history
                SET status = 'ROLLED_BACK'
                WHERE version = %s
            """, (version,))

            connection.commit()
            self.logger.info(f"✅ 迁移回滚成功: {filename}")
            return True

        except Exception as e:
            if 'connection' in locals() and connection:
                connection.rollback()
            self.logger.error(f"❌ 迁移回滚失败: {version} - {e}")
            return False

        finally:
            if 'connection' in locals():
                connection.close()

    def get_migration_status(self) -> Dict:
        """获取迁移状态信息"""
        try:
            connection = self._get_connection()
            cursor = connection.cursor()

            # 获取迁移历史
            cursor.execute("""
                SELECT version, filename, status, executed_at, execution_time_ms, error_message
                FROM migration_history
                ORDER BY executed_at DESC
            """)

            history = []
            for row in cursor.fetchall():
                history.append({
                    'version': row[0],
                    'filename': row[1],
                    'status': row[2],
                    'executed_at': row[3].strftime('%Y-%m-%d %H:%M:%S') if row[3] else None,
                    'execution_time_ms': row[4],
                    'error_message': row[5]
                })

            # 获取待执行迁移
            pending = self.get_pending_migrations()

            return {
                'executed_count': len([h for h in history if h['status'] == 'SUCCESS']),
                'failed_count': len([h for h in history if h['status'] == 'FAILED']),
                'rolled_back_count': len([h for h in history if h['status'] == 'ROLLED_BACK']),
                'pending_count': len(pending),
                'history': history,
                'pending_migrations': pending
            }

        except Exception as e:
            self.logger.error(f"❌ 获取迁移状态失败: {e}")
            raise
        finally:
            if 'connection' in locals():
                connection.close()

    def validate_migrations(self) -> bool:
        """验证所有迁移文件的完整性"""
        self.logger.info("🔍 验证迁移文件完整性...")

        migration_files = [f for f in os.listdir(self.migrations_dir) if f.endswith('.sql')]
        migration_files.sort()

        valid_count = 0
        for filename in migration_files:
            try:
                filepath = os.path.join(self.migrations_dir, filename)
                migration_data = self._parse_migration_file(filepath)

                # 检查基本结构
                if not migration_data['up_sql']:
                    self.logger.warning(f"⚠️ 迁移文件缺少UP部分: {filename}")
                    continue

                # 检查版本号格式
                version_match = re.match(r'^\d{8}_\d{6}_', filename)
                if not version_match:
                    self.logger.warning(f"⚠️ 迁移文件名格式不正确: {filename}")
                    continue

                valid_count += 1
                self.logger.debug(f"✅ 验证通过: {filename}")

            except Exception as e:
                self.logger.error(f"❌ 验证失败: {filename} - {e}")

        self.logger.info(f"📊 验证完成: {valid_count}/{len(migration_files)} 个文件通过验证")
        return valid_count == len(migration_files)
