"""
FlaxKV2 FlaxQueue — 持久化 FIFO 队列

基于 LMDB，O(1) 入队/出队，支持 max_size 自动丢弃。
"""

import os
from typing import Any, Iterator, Optional

from flaxkv2.core.flax_list import _LmdbStorage
from flaxkv2.serialization.value_meta import ValueWithMeta

_SENTINEL = object()
_HEAD_KEY = b'\xff' * 7 + b'\xfd'
_TAIL_KEY = b'\xff' * 7 + b'\xfc'


class FlaxQueue:
    """
    持久化 FIFO 队列，基于 LMDB。

    使用 head/tail 指针实现 O(1) 入队和出队，无需移位。

    用法:
        q = FlaxQueue("tasks", "./data")
        q.put("job_1")
        q.put("job_2")
        item = q.get()  # "job_1"
    """

    def __init__(
        self,
        name: str,
        path: str = ".",
        max_size: int | None = None,
        map_size: int = 2 ** 30,
        mp_safe: bool = False,
    ):
        self.name = name
        self.path = os.path.abspath(path)
        self.db_path = os.path.join(self.path, name)
        os.makedirs(self.db_path, exist_ok=True)
        self._db = _LmdbStorage(self.db_path, map_size=map_size, mp_safe=mp_safe)
        self._max_size = max_size
        self._closed = False

        # 恢复 head/tail 指针
        raw_head = self._db.get(_HEAD_KEY)
        raw_tail = self._db.get(_TAIL_KEY)
        self._head = int.from_bytes(raw_head, 'big') if raw_head else 0
        self._tail = int.from_bytes(raw_tail, 'big') if raw_tail else 0

    def _data_key(self, idx: int) -> bytes:
        return idx.to_bytes(8, 'big')

    def put(self, value: Any):
        """入队。若有 max_size 且满，自动丢弃队首。"""
        if self._max_size and self.qsize() >= self._max_size:
            self.get()  # 丢弃队首
        val_bytes = ValueWithMeta.encode_value(value, ttl_seconds=None)
        wb = self._db.write_batch()
        wb.put(self._data_key(self._tail), val_bytes)
        self._tail += 1
        wb.put(_TAIL_KEY, self._tail.to_bytes(8, 'big'))
        wb.write()

    def put_batch(self, values):
        """批量入队。"""
        values = list(values)
        if not values:
            return
        # 若批量超过 max_size，只取最后 max_size 个
        if self._max_size and len(values) > self._max_size:
            values = values[-self._max_size:]
        # 先丢弃溢出的队首
        if self._max_size:
            overflow = self.qsize() + len(values) - self._max_size
            if overflow > 0:
                self.get_batch(overflow)
        wb = self._db.write_batch()
        for v in values:
            val_bytes = ValueWithMeta.encode_value(v, ttl_seconds=None)
            wb.put(self._data_key(self._tail), val_bytes)
            self._tail += 1
        wb.put(_TAIL_KEY, self._tail.to_bytes(8, 'big'))
        wb.write()

    def get(self, default=_SENTINEL) -> Any:
        """出队。队列空时返回 default 或抛 IndexError。"""
        if self._head >= self._tail:
            if default is not _SENTINEL:
                return default
            raise IndexError("get from empty FlaxQueue")
        raw = self._db.get(self._data_key(self._head))
        if raw is None:
            if default is not _SENTINEL:
                return default
            raise IndexError("get from empty FlaxQueue")
        value, _, _ = ValueWithMeta.decode_value(raw)
        wb = self._db.write_batch()
        wb.delete(self._data_key(self._head))
        self._head += 1
        wb.put(_HEAD_KEY, self._head.to_bytes(8, 'big'))
        wb.write()
        return value

    def get_batch(self, n: int) -> list:
        """批量出队，最多取 n 个。"""
        n = min(n, self.qsize())
        if n <= 0:
            return []
        results = []
        wb = self._db.write_batch()
        for _ in range(n):
            raw = self._db.get(self._data_key(self._head))
            if raw is None:
                break
            value, _, _ = ValueWithMeta.decode_value(raw)
            results.append(value)
            wb.delete(self._data_key(self._head))
            self._head += 1
        wb.put(_HEAD_KEY, self._head.to_bytes(8, 'big'))
        wb.write()
        return results

    def peek(self, default=_SENTINEL) -> Any:
        """查看队首，不移除。"""
        if self._head >= self._tail:
            if default is not _SENTINEL:
                return default
            raise IndexError("peek at empty FlaxQueue")
        raw = self._db.get(self._data_key(self._head))
        if raw is None:
            if default is not _SENTINEL:
                return default
            raise IndexError("peek at empty FlaxQueue")
        value, _, _ = ValueWithMeta.decode_value(raw)
        return value

    def qsize(self) -> int:
        return self._tail - self._head

    def empty(self) -> bool:
        return self._head >= self._tail

    def full(self) -> bool:
        """仅 max_size 时有意义。"""
        if self._max_size is None:
            return False
        return self.qsize() >= self._max_size

    def clear(self):
        wb = self._db.write_batch()
        for i in range(self._head, self._tail):
            wb.delete(self._data_key(i))
        self._head = 0
        self._tail = 0
        wb.put(_HEAD_KEY, self._head.to_bytes(8, 'big'))
        wb.put(_TAIL_KEY, self._tail.to_bytes(8, 'big'))
        wb.write()

    def compact(self) -> int:
        """重新编号，回收空洞空间。返回移动的元素数。"""
        if self._head == 0:
            return 0
        size = self.qsize()
        if size == 0:
            self._head = 0
            self._tail = 0
            wb = self._db.write_batch()
            wb.put(_HEAD_KEY, (0).to_bytes(8, 'big'))
            wb.put(_TAIL_KEY, (0).to_bytes(8, 'big'))
            wb.write()
            return 0
        # 读取所有元素的原始字节
        raw_items = []
        for i in range(self._head, self._tail):
            raw = self._db.get(self._data_key(i))
            raw_items.append(raw)
        # 重写
        wb = self._db.write_batch()
        # 删除旧键
        for i in range(self._head, self._tail):
            wb.delete(self._data_key(i))
        # 写入新键
        for i, raw in enumerate(raw_items):
            wb.put(self._data_key(i), raw)
        self._head = 0
        self._tail = size
        wb.put(_HEAD_KEY, (0).to_bytes(8, 'big'))
        wb.put(_TAIL_KEY, size.to_bytes(8, 'big'))
        wb.write()
        return size

    def __len__(self) -> int:
        return self.qsize()

    def __bool__(self) -> bool:
        return not self.empty()

    def __iter__(self) -> Iterator:
        """迭代队列（不消费）。"""
        for i in range(self._head, self._tail):
            raw = self._db.get(self._data_key(i))
            if raw is not None:
                value, _, _ = ValueWithMeta.decode_value(raw)
                yield value

    def close(self):
        if self._closed:
            return
        self._db.close()
        self._closed = True

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()
        return False

    def __repr__(self):
        return f"FlaxQueue(name={self.name!r}, size={self.qsize()})"
