"""
FlaxKV2 RotatingFlaxList — 时间分片日志

管理多个 FlaxList 实例，按时间分片，自动清理过期分片。
"""

import os
import shutil
from datetime import datetime, timezone
from typing import Any, Iterator, Optional

from flaxkv2.core.flax_list import FlaxList


class RotatingFlaxList:
    """
    时间分片的持久化列表，适用于日志轮转场景。

    按时间（小时/天）自动创建新分片，保留最近 N 个分片。

    用法:
        log = RotatingFlaxList("app_log", "./data", partition_by="day", retention=7)
        log.append({"level": "INFO", "msg": "started"})
        log.cleanup()  # 删除超过 7 天的分片
    """

    def __init__(
        self,
        name: str,
        path: str = ".",
        partition_by: str = "day",
        retention: int = 7,
        max_size_per_partition: int | None = None,
        map_size: int = 2 ** 30,
        mp_safe: bool = False,
    ):
        if partition_by not in ("hour", "day"):
            raise ValueError(f"partition_by must be 'hour' or 'day', got {partition_by!r}")
        if retention < 1:
            raise ValueError(f"retention must be >= 1, got {retention}")

        self.name = name
        self.path = os.path.abspath(path)
        self._partition_by = partition_by
        self._retention = retention
        self._max_size_per_partition = max_size_per_partition
        self._map_size = map_size
        self._mp_safe = mp_safe
        self._closed = False

        # 当前活跃分片缓存
        self._current_key: str | None = None
        self._current_list: FlaxList | None = None

    def _partition_key(self, dt: datetime | None = None) -> str:
        """生成分片键，如 'app_log_2026-03-12' 或 'app_log_2026-03-12T14'"""
        if dt is None:
            dt = datetime.now(timezone.utc)
        if self._partition_by == "hour":
            return f"{self.name}_{dt.strftime('%Y-%m-%dT%H')}"
        return f"{self.name}_{dt.strftime('%Y-%m-%d')}"

    def _get_or_create(self, key: str) -> FlaxList:
        """获取或创建指定分片"""
        if self._current_key == key and self._current_list is not None:
            return self._current_list
        # 关闭旧分片
        if self._current_list is not None:
            self._current_list.close()
        self._current_key = key
        self._current_list = FlaxList(
            key, self.path,
            max_size=self._max_size_per_partition,
            map_size=self._map_size,
            mp_safe=self._mp_safe,
        )
        return self._current_list

    def current(self) -> FlaxList:
        """当前分片"""
        key = self._partition_key()
        return self._get_or_create(key)

    def append(self, value: Any) -> int:
        """追加到当前分片"""
        return self.current().append(value)

    def extend(self, values) -> int:
        """批量追加到当前分片"""
        return self.current().extend(values)

    def partitions(self) -> list[str]:
        """列出所有分片名（按时间排序）"""
        prefix = f"{self.name}_"
        result = []
        if not os.path.exists(self.path):
            return result
        for entry in os.listdir(self.path):
            full = os.path.join(self.path, entry)
            if entry.startswith(prefix) and os.path.isdir(full):
                result.append(entry)
        result.sort()
        return result

    def cleanup(self) -> int:
        """删除超出 retention 的旧分片，返回删除的分片数。"""
        parts = self.partitions()
        if len(parts) <= self._retention:
            return 0
        to_delete = parts[:-self._retention]
        deleted = 0
        for name in to_delete:
            # 如果是当前活跃分片，先关闭
            if self._current_key == name and self._current_list is not None:
                self._current_list.close()
                self._current_key = None
                self._current_list = None
            full_path = os.path.join(self.path, name)
            if os.path.exists(full_path):
                shutil.rmtree(full_path)
                deleted += 1
        return deleted

    def __iter__(self) -> Iterator:
        """跨分片迭代（从旧到新）"""
        for name in self.partitions():
            lst = FlaxList(name, self.path, map_size=self._map_size, mp_safe=self._mp_safe)
            try:
                yield from lst
            finally:
                lst.close()

    def __len__(self) -> int:
        total = 0
        for name in self.partitions():
            lst = FlaxList(name, self.path, map_size=self._map_size, mp_safe=self._mp_safe)
            try:
                total += len(lst)
            finally:
                lst.close()
        return total

    def close(self):
        if self._closed:
            return
        if self._current_list is not None:
            self._current_list.close()
            self._current_list = None
        self._closed = True

    def __enter__(self):
        return self

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

    def __repr__(self):
        return (
            f"RotatingFlaxList(name={self.name!r}, "
            f"partition_by={self._partition_by!r}, "
            f"retention={self._retention})"
        )
