"""
FlaxKV2 FlaxTimeSeries — 时间序列存储

基于 FlaxList sorted 模式，每条数据自动打时间戳，支持高效时间范围查询。
"""

import math
import os
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Generator, Iterator, Optional

from flaxkv2.core.flax_list import FlaxList


@dataclass
class TimeSeriesStats:
    """时间范围内的聚合统计"""
    count: int = 0
    min: float = float('inf')
    max: float = float('-inf')
    sum: float = 0.0
    mean: float = 0.0
    variance: float = 0.0
    std: float = 0.0
    start_time: float | None = None
    end_time: float | None = None

    def __repr__(self):
        lines = [
            f"  count    = {self.count}",
            f"  min      = {self.min}",
            f"  max      = {self.max}",
            f"  sum      = {self.sum}",
            f"  mean     = {self.mean:.6f}",
            f"  std      = {self.std:.6f}",
        ]
        if self.start_time is not None:
            lines.append(f"  start_time = {self.start_time}")
            lines.append(f"  end_time   = {self.end_time}")
        return "TimeSeriesStats(\n" + "\n".join(lines) + "\n)"


class FlaxTimeSeries:
    """
    持久化时间序列存储。

    每条数据存为 (timestamp, value) 元组，sorted 模式以 timestamp 排序，
    支持 O(log n) 时间范围查询。

    用法:
        ts = FlaxTimeSeries("metrics", "./data", retention=86400, min_size=100)
        ts.append({"cpu": 0.8, "mem": 0.6})
        recent = ts.latest(10)
        hourly = ts.query(start=time.time() - 3600)
    """

    def __init__(
        self,
        name: str,
        path: str = ".",
        retention: float | None = None,
        min_size: int = 0,
        max_size: int | None = None,
        map_size: int = 2 ** 30,
        mp_safe: bool = False,
    ):
        self.name = name
        self.path = os.path.abspath(path)
        self._retention = retention
        self._min_size = min_size
        self._closed = False

        self._list = FlaxList(
            name, path,
            sorted=True,
            key=lambda x: x[0],  # 按 timestamp 排序
            max_size=max_size,
            map_size=map_size,
            mp_safe=mp_safe,
        )

    # ========== 定位辅助 ==========

    def _locate(self, start: float | None, end: float | None) -> tuple[int, int]:
        """返回 [lo, hi) 索引范围，O(log n)。"""
        n = len(self._list)
        if n == 0:
            return 0, 0
        lo = self._list._bisect_left((start, None)) if start is not None else 0
        hi = self._list._bisect_right((end, None)) if end is not None else n
        return lo, hi

    # ========== 写入 ==========

    def append(self, value: Any, timestamp: float | None = None):
        """追加数据点，默认使用当前时间戳。"""
        if timestamp is None:
            timestamp = time.time()
        self._list.add((timestamp, value))

    def extend(self, items):
        """
        批量追加数据点。

        items 中每个元素可以是：
        - value（自动使用当前时间戳，递增微秒避免重复）
        - (value, timestamp) 元组
        """
        now = time.time()
        for i, item in enumerate(items):
            if isinstance(item, (list, tuple)) and len(item) == 2:
                value, ts = item
                if ts is None:
                    ts = now + i * 1e-6
            else:
                value = item
                ts = now + i * 1e-6
            self._list.add((ts, value))

    # ========== 查询 ==========

    def query(self, start: float | None = None, end: float | None = None) -> list:
        """
        时间范围查询。O(log n) 定位 + O(k) 读取。

        返回 [(timestamp, value), ...] 列表。
        """
        lo, hi = self._locate(start, end)
        if lo >= hi:
            return []
        return self._list._range_read(lo, hi)

    def query_values(self, start: float | None = None, end: float | None = None) -> list:
        """时间范围查询，只返回 values（不含 timestamp）。"""
        return [item[1] for item in self.query(start, end)]

    def count(self, start: float | None = None, end: float | None = None) -> int:
        """O(log n) 计数，不加载数据。"""
        lo, hi = self._locate(start, end)
        return max(0, hi - lo)

    def latest(self, n: int = 1) -> list:
        """最近 N 条数据。"""
        length = len(self._list)
        if length == 0:
            return []
        n = min(n, length)
        start = length - n
        return self._list._range_read(start, length)

    def oldest(self, n: int = 1) -> list:
        """最早 N 条数据。"""
        length = len(self._list)
        if length == 0:
            return []
        n = min(n, length)
        return self._list._range_read(0, n)

    def time_range(self) -> tuple[float, float] | None:
        """返回 (最早时间戳, 最新时间戳)，空序列返回 None。"""
        if len(self._list) == 0:
            return None
        oldest = self._list._get_by_index(0)
        newest = self._list._get_by_index(len(self._list) - 1)
        return (oldest[0], newest[0])

    # ========== 统计 ==========

    def stats(self, start: float | None = None, end: float | None = None,
              key: Callable | None = None) -> TimeSeriesStats:
        """
        时间范围内的聚合统计（Welford 在线算法，O(1) 内存）。

        Args:
            start/end: 时间范围（默认全部）
            key: 从 value 提取数值的函数，如 lambda x: x["cpu"]。
                 默认 None 表示 value 本身是数值。
        """
        lo, hi = self._locate(start, end)
        result = TimeSeriesStats()
        if lo >= hi:
            return result

        n = 0
        mean = 0.0
        m2 = 0.0
        total = 0.0
        vmin = float('inf')
        vmax = float('-inf')
        t_start = None
        t_end = None

        for item in self._list.iter_range(lo, hi):
            ts, val = item[0], item[1]
            v = float(key(val)) if key else float(val)
            n += 1
            total += v
            if v < vmin:
                vmin = v
            if v > vmax:
                vmax = v
            delta = v - mean
            mean += delta / n
            delta2 = v - mean
            m2 += delta * delta2
            if t_start is None:
                t_start = ts
            t_end = ts

        result.count = n
        result.min = vmin
        result.max = vmax
        result.sum = total
        result.mean = mean
        result.variance = m2 / n if n > 0 else 0.0
        result.std = math.sqrt(result.variance)
        result.start_time = t_start
        result.end_time = t_end
        return result

    # ========== 降采样 ==========

    def downsample(self, interval: float,
                   start: float | None = None, end: float | None = None,
                   agg: str = "mean") -> list[tuple[float, float]]:
        """
        按固定时间间隔聚合数据。

        Args:
            interval: 聚合间隔（秒），如 60 表示每分钟一个点
            start/end: 时间范围
            agg: 聚合方式 "mean" | "sum" | "min" | "max" | "count"

        Returns:
            [(bucket_start_time, aggregated_value), ...] 列表
        """
        lo, hi = self._locate(start, end)
        if lo >= hi:
            return []

        buckets: list[tuple[float, float]] = []
        bucket_start = None
        bucket_vals: list[float] = []

        for item in self._list.iter_range(lo, hi):
            ts, val = item[0], item[1]
            v = float(val)

            if bucket_start is None:
                bucket_start = ts - (ts % interval)

            current_bucket = ts - (ts % interval)
            if current_bucket != bucket_start:
                # 输出上一个 bucket
                buckets.append((bucket_start, self._agg(bucket_vals, agg)))
                bucket_start = current_bucket
                bucket_vals = []

            bucket_vals.append(v)

        # 最后一个 bucket
        if bucket_vals:
            buckets.append((bucket_start, self._agg(bucket_vals, agg)))

        return buckets

    @staticmethod
    def _agg(vals: list[float], agg: str) -> float:
        if agg == "mean":
            return sum(vals) / len(vals)
        elif agg == "sum":
            return sum(vals)
        elif agg == "min":
            return min(vals)
        elif agg == "max":
            return max(vals)
        elif agg == "count":
            return float(len(vals))
        raise ValueError(f"unsupported agg: {agg!r}")

    # ========== 滑动窗口 ==========

    def iter_windows(self, window_size: float, step: float | None = None,
                     start: float | None = None,
                     end: float | None = None) -> Generator[list, None, None]:
        """
        滑动窗口迭代。

        Args:
            window_size: 窗口大小（秒）
            step: 步长（秒），默认等于 window_size（滚动窗口）
            start/end: 时间范围

        Yields:
            每个窗口内的 [(timestamp, value), ...] 列表
        """
        if step is None:
            step = window_size

        tr = self.time_range()
        if tr is None:
            return
        series_start, series_end = tr

        win_start = start if start is not None else series_start
        win_end = end if end is not None else series_end

        t = win_start
        while t <= win_end:
            window = self.query(t, t + window_size)
            if window:
                yield window
            t += step

    # ========== 删除 ==========

    def delete_range(self, start: float, end: float) -> int:
        """删除指定时间范围内的数据，返回删除数。"""
        lo, hi = self._locate(start, end)
        count = hi - lo
        if count <= 0:
            return 0
        # 从后往前删避免索引偏移
        for i in range(hi - 1, lo - 1, -1):
            del self._list[i]
        return count

    def cleanup(self) -> int:
        """删除超出 retention 的旧数据，返回删除数。受 min_size 保护。"""
        if self._retention is None or len(self._list) == 0:
            return 0
        cutoff = time.time() - self._retention
        pos = self._list._bisect_left((cutoff, None))
        if pos <= 0:
            return 0
        # min_size 保护：保证至少保留 min_size 条
        if self._min_size and len(self._list) - pos < self._min_size:
            pos = max(0, len(self._list) - self._min_size)
        if pos <= 0:
            return 0
        self._list._trim_front(pos)
        return pos

    # ========== 基础 ==========

    def __len__(self) -> int:
        return len(self._list)

    def __iter__(self) -> Iterator:
        return iter(self._list)

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

    def __enter__(self):
        return self

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

    def __repr__(self):
        return f"FlaxTimeSeries(name={self.name!r}, length={len(self._list)})"
