#!/usr/bin/env python
"""
天璇增量更新应用脚本 — 在离线机上运行

用于将增量更新包应用到离线机上的项目目录。

Usage:
    runtime\\python\\python.exe scripts\\apply_update.py <update.zip>
    runtime\\python\\python.exe scripts\\apply_update.py --rollback

安全保证:
1. 更新前完整备份当前配置
2. 永不覆盖 %APPDATA%/TianXuan/ 下的用户数据
3. config.yaml 用户编辑通过合并保留
4. 完整 SHA256 校验防止传输损坏
5. 可回滚到最近一次备份
"""
import hashlib
import json
import os
import shutil
import subprocess
import sys
import zipfile
from datetime import datetime
from pathlib import Path

# ── Constants ──────────────────────────────────────────────────────────
PROJECT_DIR = Path(__file__).resolve().parent.parent
PYTHON = str(PROJECT_DIR / 'runtime' / 'python' / 'python.exe')

if not os.path.exists(PYTHON):
    # Fallback: use system python if portable not found
    PYTHON = sys.executable

APPDATA_DIR = Path(os.environ.get('APPDATA', str(Path.home() / '.tianxuan'))) / 'TianXuan'
BACKUP_DIR = APPDATA_DIR / 'backups'
VERSION_FILE = PROJECT_DIR / 'VERSION'
CONFIG_FILE = PROJECT_DIR / 'config' / 'config.yaml'
SCRIPTS_DIR = PROJECT_DIR / 'scripts'
MERGE_SCRIPT = SCRIPTS_DIR / 'merge_config.py'

# ── Helper Functions ──────────────────────────────────────────────────


def log(msg: str):
    print(f"  {msg}")


def step(n: int, total: int, msg: str):
    print(f"\n[{n}/{total}] {msg}")


def verify_sha256(filepath: Path, expected: str) -> bool:
    """Verify file SHA256 matches expected hash."""
    h = hashlib.sha256()
    with open(filepath, 'rb') as f:
        for chunk in iter(lambda: f.read(65536), b''):
            h.update(chunk)
    return h.hexdigest() == expected


def clean_pycache(root_dir: Path):
    """Recursively remove all __pycache__ directories."""
    count = 0
    for dirpath, dirnames, _ in os.walk(str(root_dir)):
        if '__pycache__' in dirnames:
            cache_dir = os.path.join(dirpath, '__pycache__')
            try:
                shutil.rmtree(cache_dir)
                count += 1
            except PermissionError:
                log(f"  [WARN] 无法删除 {cache_dir} (权限拒绝)")
    if count > 0:
        log(f"  清理了 {count} 个 __pycache__ 目录")


def get_current_version() -> str:
    """Read current version from VERSION file."""
    if VERSION_FILE.exists():
        return VERSION_FILE.read_text(encoding='utf-8').strip()
    return 'v0.0.0'


# ── Main Logic ────────────────────────────────────────────────────────


def do_rollback():
    """Restore from the most recent backup."""
    if not BACKUP_DIR.exists():
        print("[ERROR] 没有找到备份目录")
        return 1

    backups = sorted(BACKUP_DIR.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True)
    if not backups:
        print("[ERROR] 备份目录为空")
        return 1

    latest = backups[0]
    config_backup = latest / 'config.yaml'
    version_backup = latest / 'VERSION'

    print(f"从 {latest.name} 回滚...")

    if config_backup.exists():
        shutil.copy2(str(config_backup), str(CONFIG_FILE))
        log("配置已恢复")

    if version_backup.exists():
        shutil.copy2(str(version_backup), str(VERSION_FILE))
        log("VERSION 已恢复")

    # Restore project snapshot (if exists)
    snapshot_dir = latest / 'project_snapshot'
    if snapshot_dir.exists():
        log("恢复项目文件快照...")
        for item in snapshot_dir.iterdir():
            dest = PROJECT_DIR / item.name
            if item.is_dir():
                if dest.exists():
                    shutil.rmtree(str(dest))
                shutil.copytree(str(item), str(dest))
            else:
                shutil.copy2(str(item), str(dest))
        log("项目文件已恢复")

    print("\n✅ 回滚完成！请重启天璇服务")
    return 0


def apply_update(zip_path: Path):
    """Apply an update package."""
    total_steps = 10

    print("=" * 50)
    print("  天璇 增量更新 — 应用工具")
    print("=" * 50)
    log(f"更新包: {zip_path.name}")
    log(f"当前版本: {get_current_version()}")
    log(f"备份目录: {BACKUP_DIR}")

    # Step 1: Validate package
    step(1, total_steps, "验证更新包完整性")

    if not zip_path.exists():
        print(f"[ERROR] 更新包不存在: {zip_path}")
        return 1

    try:
        with zipfile.ZipFile(zip_path, 'r') as zf:
            if zf.testzip() is not None:
                print(f"[CORRUPT] 更新包损坏！第一个损坏文件: {zf.testzip()}")
                return 1
            # Read manifest
            if 'files/update_info.json' not in zf.namelist():
                print("[ERROR] 更新包缺少 files/update_info.json")
                return 1
            manifest = json.loads(zf.read('files/update_info.json'))
    except Exception as e:
        print(f"[ERROR] 无法读取更新包: {e}")
        return 1

    log(f"  {zip_path.stat().st_size / 1024:.1f} KB — 校验通过")

    # Step 2: Check version compatibility
    step(2, total_steps, "检查版本兼容性")

    current_ver = get_current_version()
    from_ver = manifest.get('from_version', '')

    if current_ver != from_ver:
        print(f"[VERSION_MISMATCH] 版本不匹配！")
        print(f"  当前: {current_ver}")
        print(f"  更新包起始: {from_ver}")
        print(f"  更新包目标: {manifest.get('to_version', '?')}")
        print(f"\n提示: 如需强制应用，请更新 VERSION 文件后重试")
        return 1

    log(f"  版本一致: {current_ver} → {manifest.get('to_version', '?')}")

    # Step 3: Verify file checksums
    step(3, total_steps, "校验文件 SHA256")

    with zipfile.ZipFile(zip_path, 'r') as zf:
        for rel_path, file_info in manifest.get('files', {}).items():
            arcname = f'files/{rel_path}'
            if arcname not in zf.namelist():
                print(f"[ERROR] 包中缺少文件: {arcname}")
                return 1
            # Extract to temp and verify
            zf.extract(arcname, str(PROJECT_DIR))
            extracted = PROJECT_DIR / arcname
            if extracted.exists() and not verify_sha256(extracted, file_info['sha256']):
                print(f"[CORRUPT] 文件校验失败: {rel_path}")
                extracted.unlink()
                return 1
            extracted.unlink()

    log(f"  {len(manifest.get('files', {}))} 个文件校验通过")

    # Step 4: Backup current config
    step(4, total_steps, "备份当前配置")

    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    version_dir = BACKUP_DIR / f'{manifest.get("to_version", "unknown")}_{timestamp}'
    version_dir.mkdir(parents=True, exist_ok=True)

    if CONFIG_FILE.exists():
        shutil.copy2(str(CONFIG_FILE), str(version_dir / 'config.yaml'))
        log("config.yaml 已备份")

    if VERSION_FILE.exists():
        shutil.copy2(str(VERSION_FILE), str(version_dir / 'VERSION'))
        log("VERSION 已备份")

    # Step 5: Delete removed files
    step(5, total_steps, "删除已废弃的文件")

    for rel_path in manifest.get('deleted', []):
        target = PROJECT_DIR / rel_path
        if target.exists():
            try:
                if target.is_dir():
                    shutil.rmtree(str(target))
                else:
                    target.unlink()
                log(f"  删除: {rel_path}")
            except Exception as e:
                log(f"  [WARN] 删除失败 {rel_path}: {e}")

    # Also process __delete__ (explicit removal list, may overlap with deleted)
    for rel_path in manifest.get('__delete__', []):
        target = PROJECT_DIR / rel_path
        if target.exists():
            try:
                if target.is_dir():
                    shutil.rmtree(str(target))
                else:
                    target.unlink()
                log(f"  删除: {rel_path}")
            except Exception as e:
                log(f"  [WARN] 删除失败 {rel_path}: {e}")

    # Step 6: Extract changed files
    # config.yaml is extracted but NOT moved — left in files/ for the merge step
    step(6, total_steps, "解压更新文件")

    extracted_count = 0
    new_config_yaml = None  # will hold path to extracted new config.yaml
    with zipfile.ZipFile(zip_path, 'r') as zf:
        for name in zf.namelist():
            if not name.startswith('files/'):
                continue
            rel_path = name[6:]  # strip 'files/' prefix
            zf.extract(name, str(PROJECT_DIR))

            if rel_path == 'config/config.yaml':
                # Don't move — keep in files/ for merge step 7
                new_config_yaml = PROJECT_DIR / name
                extracted_count += 1
                continue

            # Move from files/ subdir to project root
            src = PROJECT_DIR / name
            dst = PROJECT_DIR / rel_path
            if src != dst:
                dst.parent.mkdir(parents=True, exist_ok=True)
                if src.exists():
                    shutil.move(str(src), str(dst))
            extracted_count += 1

    log(f"  已解压 {extracted_count} 个文件")

    # Step 7: Merge config.yaml
    step(7, total_steps, "合并配置文件")

    user_config = version_dir / 'config.yaml'  # backed-up user config with edits

    if user_config.exists() and new_config_yaml and new_config_yaml.exists() and MERGE_SCRIPT.exists():
        log("合并用户配置与更新配置...")
        # --old: user's backed-up config (preserves their edits)
        # --new: config.yaml from the update package (new keys to add)
        # --output: the live config file
        result = subprocess.run(
            [sys.executable, str(MERGE_SCRIPT),
             '--old', str(user_config),
             '--new', str(new_config_yaml),
             '--output', str(CONFIG_FILE)],
            capture_output=True, text=True, cwd=str(PROJECT_DIR)
        )
        print(result.stdout.strip())
        if result.returncode != 0:
            log(f"  [WARN] 配置合并异常: {result.stderr.strip()}")
    elif user_config.exists() and MERGE_SCRIPT.exists():
        # New config not in package — keep user config unchanged
        log("  更新包不含 config.yaml，保留现有配置")
        # Restore the backup since nothing new to merge
        shutil.copy2(str(user_config), str(CONFIG_FILE))
    else:
        log("  无备份配置，保留现有 config.yaml")

    # Clean up extracted files/ directory
    files_dir = PROJECT_DIR / 'files'
    if files_dir.exists():
        shutil.rmtree(str(files_dir))

    # Step 8: Clean __pycache__
    step(8, total_steps, "清理 Python 缓存")

    clean_pycache(PROJECT_DIR)
    # Also clean in runtime (portable python)
    runtime_dir = PROJECT_DIR / 'runtime'
    if runtime_dir.exists():
        clean_pycache(runtime_dir)

    # Step 9: Run Django migrations (always, since migrate is idempotent)
    step(9, total_steps, "运行数据库迁移")

    python_exe = PYTHON
    migrate_cmd = [python_exe, 'manage.py', 'migrate']
    log(f"  执行: {' '.join(migrate_cmd)}")

    result = subprocess.run(migrate_cmd, capture_output=True, text=True, cwd=str(PROJECT_DIR))
    print(result.stdout.strip())
    if result.returncode != 0:
        print(f"[MIGRATE_ERROR] 迁移失败: {result.stderr.strip()}")
        print("  请运行 python scripts/apply_update.py --rollback 回滚")
        return 1

    if 'Applying' in result.stdout or manifest.get('run_migrate'):
        log("  数据库 schema 已更新 ✓")
    else:
        log("  数据库已是最新 ✓")

    # ── Ensure tls_ref_data table ─────────────────────────────────────
    # import_tlsdb is idempotent — safe to run every update.
    log("  检查 tls_ref_data 引用数据表...")
    import_result = subprocess.run(
        [python_exe, 'manage.py', 'import_tlsdb'],
        capture_output=True, text=True, cwd=str(PROJECT_DIR)
    )
    if import_result.returncode == 0:
        log("  引用数据已就绪 ✓")
    else:
        # import_tlsdb may warn about existing data; that's fine
        if 'already exists' in import_result.stderr.lower() or '0 row' in import_result.stdout.lower():
            log("  引用数据已存在 ✓")
        else:
            log(f"  [INFO] import_tlsdb 输出: {import_result.stdout.strip()[:200]}")

    # Step 10: Update VERSION and finalize
    step(10, total_steps, "完成更新")

    to_version = manifest.get('to_version', '')
    if to_version:
        VERSION_FILE.write_text(to_version + '\n', encoding='utf-8')
        log(f"VERSION → {to_version}")

    # Print runtime warning if needed
    if manifest.get('runtime_changed'):
        print("\n" + "!" * 50)
        print("  ⚠️  依赖已变更！")
        print("  请手动将开发机的 runtime/ 目录同步到离线机")
        print("  (使用 U 盘复制 %APPDATA%/TianXuan/runtime/ 或直接覆盖)")
        print("!" * 50)

    print(f"\n{'=' * 50}")
    print(f"✅ 更新完成！")
    print(f"  {from_ver} → {to_version}")
    print(f"  {extracted_count} 个文件已更新")
    print(f"  {len(manifest.get('deleted', []))} 个文件已删除")
    print(f"\n  请关闭当前窗口，重新双击 run.bat 启动新版本")
    print(f"{'=' * 50}")
    return 0


def main():
    if len(sys.argv) < 2:
        print("用法:")
        print("  python scripts/apply_update.py <update.zip>  — 应用更新")
        print("  python scripts/apply_update.py --rollback    — 回滚到最近备份")
        return 1

    if sys.argv[1] == '--rollback':
        return do_rollback()

    zip_path = Path(sys.argv[1])
    if not zip_path.exists():
        # Try relative to project dir
        zip_path = PROJECT_DIR / sys.argv[1]
    if not zip_path.exists():
        print(f"[ERROR] 找不到更新包: {sys.argv[1]}")
        return 1

    return apply_update(zip_path.resolve())


if __name__ == '__main__':
    sys.exit(main())
