"""
FlaxKV2 数据迁移工具：LevelDB → LMDB

将旧版 plyvel (LevelDB) 数据库迁移到 LMDB 后端。
迁移过程保留原始字节，不解码。

依赖：plyvel（可选依赖，仅迁移时需要）
安装：pip install flaxkv2[migration]
"""

import os
import shutil
from typing import Optional

from flaxkv2.utils.log import get_logger

logger = get_logger(__name__)


def migrate_leveldb_to_lmdb(
    db_path: str,
    output_path: Optional[str] = None,
    map_size: int = 1 << 30,
    batch_size: int = 10000,
) -> str:
    """
    将 LevelDB 数据库迁移到 LMDB

    Args:
        db_path: LevelDB 数据库路径
        output_path: LMDB 输出路径（默认为 db_path + '.lmdb'）
        map_size: 初始 LMDB map_size（遇到 MapFullError 自动翻倍）
        batch_size: 每批写入的键数

    Returns:
        输出路径

    Raises:
        ImportError: 未安装 plyvel
        FileNotFoundError: LevelDB 路径不存在
        FileExistsError: 输出路径已存在
    """
    try:
        import plyvel
    except ImportError:
        raise ImportError(
            "迁移工具需要 plyvel，请安装：pip install flaxkv2[migration]"
        )

    from flaxkv2.core.lmdb_backend import LmdbDatabase

    db_path = os.path.abspath(db_path)
    if not os.path.exists(db_path):
        raise FileNotFoundError(f"LevelDB 数据库不存在: {db_path}")

    if output_path is None:
        output_path = db_path + '.lmdb'

    output_path = os.path.abspath(output_path)
    if os.path.exists(output_path):
        raise FileExistsError(f"输出路径已存在: {output_path}")

    logger.info(f"开始迁移: {db_path} → {output_path}")

    # 打开 LevelDB（只读）
    src_db = plyvel.DB(db_path, create_if_missing=False)

    # 创建 LMDB
    dst_db = LmdbDatabase(output_path, map_size=map_size)

    total = 0
    try:
        wb = dst_db.write_batch()
        count = 0

        for key, value in src_db:
            wb.put(key, value)
            count += 1
            total += 1

            if count >= batch_size:
                wb.write()
                wb = dst_db.write_batch()
                logger.info(f"已迁移 {total} 条记录")
                count = 0

        # 写入剩余数据
        if count > 0:
            wb.write()

        logger.info(f"迁移完成：共 {total} 条记录")

    finally:
        src_db.close()
        dst_db.close()

    return output_path


def migrate_and_replace(
    db_path: str,
    backup: bool = True,
    map_size: int = 1 << 30,
    batch_size: int = 10000,
) -> str:
    """
    迁移并替换原有 LevelDB 数据库

    1. LevelDB → 临时 LMDB
    2. 备份原 LevelDB（可选）
    3. 将 LMDB 移到原路径

    Args:
        db_path: LevelDB 数据库路径
        backup: 是否保留 LevelDB 备份（默认 True）
        map_size: 初始 LMDB map_size
        batch_size: 每批写入的键数

    Returns:
        最终数据库路径
    """
    db_path = os.path.abspath(db_path)
    tmp_path = db_path + '.lmdb_tmp'
    backup_path = db_path + '.leveldb_backup'

    # 迁移到临时路径
    migrate_leveldb_to_lmdb(db_path, tmp_path, map_size, batch_size)

    # 备份或删除原 LevelDB
    if backup:
        os.rename(db_path, backup_path)
        logger.info(f"LevelDB 备份到: {backup_path}")
    else:
        shutil.rmtree(db_path)
        logger.info(f"已删除原 LevelDB: {db_path}")

    # 移动 LMDB 到原路径
    os.rename(tmp_path, db_path)
    logger.info(f"LMDB 已就位: {db_path}")

    return db_path
