"""
FlaxKV2 LMDB 后端包装层

提供与 plyvel (LevelDB) 兼容的 API，让上层代码无感迁移。
核心优势：LMDB 原生支持多进程并发读（mmap + MVCC）。

三个核心类：
- LmdbDatabase: 替代 plyvel.DB
- PrefixedLmdbDB: 替代 plyvel.prefixed_db()
- LmdbWriteBatch: 替代 plyvel.WriteBatch
"""

import os
import threading
from typing import Optional

import lmdb

from flaxkv2.utils.log import get_logger

logger = get_logger(__name__)

# 默认 map_size: 1GB（mmap 是稀疏的，设大不会立即占磁盘）
_DEFAULT_MAP_SIZE = 1 << 30


class LmdbWriteBatch:
    """替代 plyvel.WriteBatch：收集 put/delete 操作，write() 时原子提交"""

    def __init__(self, db: "LmdbDatabase"):
        self._db = db
        self._ops: list = []

    def put(self, key: bytes, value: bytes):
        self._ops.append(("put", key, value))

    def delete(self, key: bytes):
        self._ops.append(("delete", key))

    def write(self):
        if not self._ops:
            return
        self._db._write_with_retry(self._apply)
        self._ops.clear()

    def _apply(self, txn: lmdb.Transaction):
        for op in self._ops:
            if op[0] == "put":
                txn.put(op[1], op[2])
            else:
                txn.delete(op[1])


class PrefixedLmdbDB:
    """
    替代 plyvel.prefixed_db()

    持有 parent LmdbDatabase + prefix bytes，所有操作自动拼接/去除前缀。
    """

    def __init__(self, parent: "LmdbDatabase", prefix: bytes):
        self._parent = parent
        self._prefix = prefix

    def get(self, key: bytes) -> Optional[bytes]:
        return self._parent.get(self._prefix + key)

    def put(self, key: bytes, value: bytes):
        self._parent.put(self._prefix + key, value)

    def delete(self, key: bytes):
        self._parent.delete(self._prefix + key)

    def write_batch(self) -> "LmdbWriteBatch":
        """返回一个前缀感知的 WriteBatch"""
        return _PrefixedWriteBatch(self._parent, self._prefix)

    def prefixed_db(self, sub_prefix: bytes) -> "PrefixedLmdbDB":
        """支持递归嵌套"""
        return PrefixedLmdbDB(self._parent, self._prefix + sub_prefix)

    def __iter__(self):
        """迭代所有以 prefix 开头的 key-value，返回去掉前缀的 key"""
        return _prefixed_iterator(
            self._parent._env, self._prefix,
            strip_prefix=self._prefix, reverse=False
        )

    def iterator(self, start=None, stop=None, prefix=None, reverse=False,
                 include_start=True, include_stop=False):
        """兼容 plyvel PrefixedDB.iterator() 接口"""
        if prefix is not None:
            full_prefix = self._prefix + prefix
            return _prefixed_iterator(
                self._parent._env, full_prefix, self._prefix, reverse=reverse
            )
        full_start = self._prefix + start if start is not None else self._prefix
        # stop: 如果未指定，用 prefix 的上界
        if stop is not None:
            full_stop = self._prefix + stop
        else:
            full_stop = _prefix_upper_bound(self._prefix)
        return _range_iterator(
            self._parent._env, full_start, full_stop,
            strip_prefix=self._prefix, reverse=reverse
        )


class _PrefixedWriteBatch:
    """前缀感知的 WriteBatch"""

    def __init__(self, parent: "LmdbDatabase", prefix: bytes):
        self._parent = parent
        self._prefix = prefix
        self._ops: list = []

    def put(self, key: bytes, value: bytes):
        self._ops.append(("put", self._prefix + key, value))

    def delete(self, key: bytes):
        self._ops.append(("delete", self._prefix + key))

    def write(self):
        if not self._ops:
            return
        self._parent._write_with_retry(self._apply)
        self._ops.clear()

    def _apply(self, txn: lmdb.Transaction):
        for op in self._ops:
            if op[0] == "put":
                txn.put(op[1], op[2])
            else:
                txn.delete(op[1])


class LmdbDatabase:
    """
    替代 plyvel.DB

    提供与 plyvel 兼容的 API：
    - get/put/delete
    - write_batch()
    - iterator(start, stop, prefix, reverse)
    - prefixed_db(prefix)
    - __iter__
    - close()
    - stat() / compact() 等维护方法
    """

    def __init__(self, path: str, map_size: int = _DEFAULT_MAP_SIZE,
                 create_if_missing: bool = True, max_readers: int = 126,
                 **kwargs):
        """
        Args:
            path: 数据库目录路径
            map_size: 初始 mmap 大小（遇到 MapFullError 自动翻倍）
            create_if_missing: 不存在时创建
            max_readers: 最大并发读者数
            **kwargs: 忽略 LevelDB 特有参数（兼容性）
        """
        self._path = path
        self._map_size = map_size
        self._lock = threading.Lock()  # 保护 map_size 扩容

        # 确保目录存在
        if create_if_missing:
            os.makedirs(path, exist_ok=True)

        self._env = lmdb.open(
            path,
            map_size=map_size,
            max_readers=max_readers,
            # subdir=True 是默认行为：path 作为目录
            subdir=True,
            # 允许多进程读
            lock=True,
            readahead=False,  # 随机读场景关闭预读
        )

    def begin(self, write=False):
        """返回 LMDB 事务上下文管理器，用于单事务内多次读写"""
        return self._env.begin(write=write)

    def get(self, key: bytes) -> Optional[bytes]:
        with self._env.begin() as txn:
            return txn.get(key)

    def put(self, key: bytes, value: bytes):
        self._write_with_retry(lambda txn: txn.put(key, value))

    def delete(self, key: bytes):
        self._write_with_retry(lambda txn: txn.delete(key))

    def write_batch(self, **kwargs) -> LmdbWriteBatch:
        """返回 WriteBatch（忽略 plyvel 的 transaction 参数）"""
        return LmdbWriteBatch(self)

    def prefixed_db(self, prefix: bytes) -> PrefixedLmdbDB:
        return PrefixedLmdbDB(self, prefix)

    def iterator(self, start=None, stop=None, prefix=None, reverse=False,
                 include_start=True, include_stop=False, **kwargs):
        """
        兼容 plyvel.DB.iterator() 接口

        Args:
            start: 范围起始键（包含）
            stop: 范围结束键（不包含）
            prefix: 键前缀过滤
            reverse: 是否反向迭代
        """
        if prefix is not None:
            return _prefixed_iterator(self._env, prefix, reverse=reverse)
        if start is not None or stop is not None:
            return _range_iterator(self._env, start, stop, reverse=reverse)
        return _full_iterator(self._env, reverse=reverse)

    def __iter__(self):
        """全量迭代所有 key-value 对"""
        return _full_iterator(self._env)

    def compact_range(self, start=None, stop=None):
        """LMDB 无内置 compact_range，使用 copy(compact=True) 替代"""
        # LMDB 不需要手动 compact，B+树结构自动管理
        # 如需真正压缩可用 env.copy(compact=True) 但会复制整个数据库
        pass

    def get_property(self, name: bytes) -> Optional[bytes]:
        """兼容 plyvel 的 get_property，返回 LMDB stat 信息"""
        try:
            stat = self._env.stat()
            info = self._env.info()
            lines = [
                "LMDB Statistics",
                "--------------------",
                f"  Page size: {stat['psize']}",
                f"  Tree depth: {stat['depth']}",
                f"  Branch pages: {stat['branch_pages']}",
                f"  Leaf pages: {stat['leaf_pages']}",
                f"  Overflow pages: {stat['overflow_pages']}",
                f"  Entries: {stat['entries']}",
                f"  Map size: {info['map_size']}",
                f"  Last page: {info['last_pgno']}",
                f"  Last txn: {info['last_txnid']}",
                f"  Max readers: {info['max_readers']}",
                f"  Num readers: {info['num_readers']}",
            ]
            return "\n".join(lines).encode("utf-8")
        except Exception:
            return None

    def approximate_size(self, start: bytes, stop: bytes) -> int:
        """近似数据库大小（字节）"""
        try:
            stat = self._env.stat()
            # 近似值：页数 * 页大小
            total_pages = (
                stat["branch_pages"] + stat["leaf_pages"] + stat["overflow_pages"]
            )
            return total_pages * stat["psize"]
        except Exception:
            return 0

    def close(self):
        if self._env is not None:
            try:
                self._env.close()
            except Exception:
                pass
            self._env = None

    def stat(self) -> dict:
        """返回 LMDB 环境统计信息"""
        return self._env.stat()

    def info(self) -> dict:
        """返回 LMDB 环境信息"""
        return self._env.info()

    def _write_with_retry(self, fn):
        """
        执行写操作，遇到 MapFullError 自动翻倍 map_size 后重试

        LMDB 的 set_mapsize() 可在运行时调整，无需重启。
        """
        while True:
            try:
                with self._env.begin(write=True) as txn:
                    fn(txn)
                return
            except lmdb.MapFullError:
                with self._lock:
                    self._map_size *= 2
                    self._env.set_mapsize(self._map_size)
                    logger.warning(
                        f"LMDB map full, doubled to {self._map_size / (1024*1024):.0f} MB"
                    )


# ========== 迭代器辅助函数 ==========


def _prefix_upper_bound(prefix: bytes) -> Optional[bytes]:
    """计算前缀的上界（用于范围查询结束）"""
    # 找到最后一个非 0xFF 的字节，+1 得到上界
    arr = bytearray(prefix)
    for i in range(len(arr) - 1, -1, -1):
        if arr[i] < 0xFF:
            arr[i] += 1
            return bytes(arr[: i + 1])
    return None  # 全是 0xFF，没有上界


def _full_iterator(env: lmdb.Environment, reverse: bool = False):
    """全量迭代"""
    txn = env.begin()
    cursor = txn.cursor()
    try:
        if reverse:
            if not cursor.last():
                return
            for key, value in cursor.iterprev():
                yield key, value
        else:
            if not cursor.first():
                return
            for key, value in cursor:
                yield key, value
    finally:
        txn.abort()


def _range_iterator(env: lmdb.Environment, start=None, stop=None,
                    strip_prefix=None, reverse: bool = False):
    """范围迭代（start 包含，stop 不包含）"""
    txn = env.begin()
    cursor = txn.cursor()
    prefix_len = len(strip_prefix) if strip_prefix else 0

    try:
        if reverse:
            # 反向迭代：从 stop 往回走到 start
            if stop is not None:
                if not cursor.set_range(stop):
                    # stop 超过所有键，从最后开始
                    if not cursor.last():
                        return
                else:
                    # set_range 找到 >= stop 的第一个键，需要往前一步
                    if not cursor.prev():
                        return
            else:
                if not cursor.last():
                    return

            for key, value in cursor.iterprev():
                if start is not None and key < start:
                    break
                if strip_prefix:
                    yield key[prefix_len:], value
                else:
                    yield key, value
        else:
            if start is not None:
                if not cursor.set_range(start):
                    return
            else:
                if not cursor.first():
                    return

            for key, value in cursor:
                if stop is not None and key >= stop:
                    break
                if strip_prefix:
                    yield key[prefix_len:], value
                else:
                    yield key, value
    finally:
        txn.abort()


def _prefixed_iterator(env: lmdb.Environment, prefix: bytes,
                       strip_prefix=None, reverse: bool = False):
    """
    前缀迭代

    Args:
        env: LMDB 环境
        prefix: 搜索前缀（只返回以此开头的键）
        strip_prefix: 需要从返回的 key 中去掉的前缀。
                      None = 不去除（返回完整 key，与 plyvel iterator(prefix=) 一致）
        reverse: 是否反向迭代
    """
    upper = _prefix_upper_bound(prefix)
    txn = env.begin()
    cursor = txn.cursor()
    sp_len = len(strip_prefix) if strip_prefix else 0

    try:
        if reverse:
            if upper is not None:
                if not cursor.set_range(upper):
                    if not cursor.last():
                        return
                else:
                    if not cursor.prev():
                        return
            else:
                if not cursor.last():
                    return

            for key, value in cursor.iterprev():
                if not key.startswith(prefix):
                    break
                yield (key[sp_len:] if sp_len else key), value
        else:
            if not cursor.set_range(prefix):
                return
            for key, value in cursor:
                if not key.startswith(prefix):
                    break
                yield (key[sp_len:] if sp_len else key), value
    finally:
        txn.abort()
