import os
import sys
import asyncio
import subprocess
import gzip
import tempfile
import shutil
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from config import settings
from scripts.backup_db import list_backups, BACKUP_DIR

async def restore_database(backup_file):
    """从备份文件恢复数据库"""
    file_path = backup_file
    if not os.path.isabs(backup_file):
        file_path = os.path.join(BACKUP_DIR, backup_file)
    
    if not os.path.exists(file_path):
        print(f"错误: 备份文件 '{file_path}' 不存在")
        return False
    
    confirm = input(f"警告: 这将覆盖当前数据库的所有数据。确定要继续吗? (y/N): ")
    if confirm.lower() != 'y':
        print("操作已取消")
        return False
    
    temp_file = None
    sql_file = file_path
    
    try:
        if file_path.endswith('.gz'):
            print("解压备份文件...")
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.sql')
            temp_file.close()
            
            with gzip.open(file_path, 'rb') as f_in:
                with open(temp_file.name, 'wb') as f_out:
                    shutil.copyfileobj(f_in, f_out)
            
            sql_file = temp_file.name
        
        cmd = [
            "mysql",
            "-h", settings.db_host,
            "-P", str(settings.db_port),
            "-u", settings.db_user,
            f"-p{settings.db_password}",
            settings.db_name
        ]
        
        print(f"开始恢复数据库...")
        with open(sql_file, 'r') as f:
            process = subprocess.Popen(
                cmd, 
                stdin=f,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            stdout, stderr = process.communicate()
            
        if process.returncode != 0:
            print(f"恢复失败: {stderr.decode()}")
            return False
        
        print("数据库恢复成功!")
        return True
    
    except Exception as e:
        print(f"恢复过程中出错: {e}")
        return False
    
    finally:
        if temp_file and os.path.exists(temp_file.name):
            os.unlink(temp_file.name)

async def interactive_restore():
    """交互式选择要恢复的备份文件"""
    backups = await list_backups()
    
    if not backups:
        print("没有找到可用的备份文件")
        return False
    
    print("\n可用的备份文件:")
    for i, backup in enumerate(backups):
        file_path = os.path.join(BACKUP_DIR, backup)
        size = os.path.getsize(file_path) / (1024 * 1024)  # 转换为MB
        print(f"{i+1}. {backup} ({size:.2f} MB)")
    
    try:
        choice = int(input("\n请选择要恢复的备份文件 (输入序号): "))
        if 1 <= choice <= len(backups):
            selected_backup = backups[choice-1]
            return await restore_database(selected_backup)
        else:
            print("无效的选择")
            return False
    except ValueError:
        print("无效的输入")
        return False

async def main():
    if len(sys.argv) > 1:
        await restore_database(sys.argv[1])
    else:
        await interactive_restore()

if __name__ == "__main__":
    asyncio.run(main())