"""Time series data storage: write/read/sample metric sequences."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from matyan_backend.fdb_types import transactional

from .encoding import decode_value, encode_value
from .fdb_client import get_directories

if TYPE_CHECKING:
    from matyan_backend.fdb_types import DirectorySubspace, Transaction


def _runs_dir() -> DirectorySubspace:
    return get_directories().runs


def _empty_result(columns: tuple[str, ...]) -> dict[str, list]:
    result: dict[str, list] = {"steps": []}
    for col in columns:
        result[col] = []
    return result


def _pick_evenly_spaced(items: list, num_points: int) -> list:
    """Return up to *num_points* evenly-spaced elements from *items*."""
    total = len(items)
    if total <= num_points:
        return items
    if num_points <= 1:
        return [items[0]]
    step_size = (total - 1) / (num_points - 1)
    indices = list(dict.fromkeys(round(i * step_size) for i in range(num_points)))
    return [items[idx] for idx in indices]


def _read_extra_columns(
    tr: Transaction,
    rd: DirectorySubspace,
    run_hash: str,
    ctx_id: int,
    name: str,
    steps: list[int],
    columns: tuple[str, ...],
) -> dict[str, list]:
    """Fetch non-val columns for the given steps."""
    result: dict[str, list] = {}
    for col in columns:
        vals: list = []
        for step in steps:
            raw = tr[rd.pack((run_hash, "seqs", ctx_id, name, col, step))]
            vals.append(decode_value(raw) if raw.present() else None)
        result[col] = vals
    return result


# ---------------------------------------------------------------------------
# Write
# ---------------------------------------------------------------------------


@transactional
def write_sequence_step(
    tr: Transaction,
    run_hash: str,
    ctx_id: int,
    name: str,
    step: int,
    value: Any,  # noqa: ANN401
    epoch: int | None = None,
    timestamp: float | None = None,
) -> None:
    rd = _runs_dir()
    base = (run_hash, "seqs", ctx_id, name)
    tr[rd.pack((*base, "val", step))] = encode_value(value)
    tr[rd.pack((*base, "step", step))] = encode_value(step)
    if epoch is not None:
        tr[rd.pack((*base, "epoch", step))] = encode_value(epoch)
    if timestamp is not None:
        tr[rd.pack((*base, "time", step))] = encode_value(timestamp)


@transactional
def write_sequence_batch(
    tr: Transaction,
    run_hash: str,
    ctx_id: int,
    name: str,
    steps_data: list[dict],
) -> None:
    """Write multiple steps in a single transaction.

    Each dict in *steps_data* must have ``step`` and ``value``, and may have
    ``epoch`` and ``timestamp``.
    """
    rd = _runs_dir()
    base = (run_hash, "seqs", ctx_id, name)
    for entry in steps_data:
        step = entry["step"]
        tr[rd.pack((*base, "val", step))] = encode_value(entry["value"])
        tr[rd.pack((*base, "step", step))] = encode_value(step)
        if "epoch" in entry and entry["epoch"] is not None:
            tr[rd.pack((*base, "epoch", step))] = encode_value(entry["epoch"])
        if "timestamp" in entry and entry["timestamp"] is not None:
            tr[rd.pack((*base, "time", step))] = encode_value(entry["timestamp"])


# ---------------------------------------------------------------------------
# Read
# ---------------------------------------------------------------------------


@transactional
def read_sequence(
    tr: Transaction,
    run_hash: str,
    ctx_id: int,
    name: str,
    *,
    start_step: int | None = None,
    end_step: int | None = None,
    columns: tuple[str, ...] = ("val",),
) -> dict[str, list]:
    """Read sequence data, optionally bounded by step range.

    Returns ``{column: [values...]}`` plus a special ``"steps"`` key with
    the step integers.
    """
    rd = _runs_dir()
    result = _empty_result(columns)

    val_base = (run_hash, "seqs", ctx_id, name, "val")
    full_range = rd.range(val_base)
    begin = rd.pack((*val_base, start_step)) if start_step is not None else full_range.start
    end = min(rd.pack((*val_base, end_step + 1)), full_range.stop) if end_step is not None else full_range.stop

    for kv in tr.get_range(begin, end):
        step = rd.unpack(kv.key)[-1]
        result["steps"].append(step)
        if "val" in columns:
            result["val"].append(decode_value(kv.value))

    extra_cols = tuple(c for c in columns if c != "val")
    if extra_cols:
        result.update(_read_extra_columns(tr, rd, run_hash, ctx_id, name, result["steps"], extra_cols))

    return result


@transactional
def sample_sequence(
    tr: Transaction,
    run_hash: str,
    ctx_id: int,
    name: str,
    num_points: int,
    *,
    columns: tuple[str, ...] = ("val",),
) -> dict[str, list]:
    """Uniformly sample *num_points* from a sequence."""
    rd = _runs_dir()
    val_base = (run_hash, "seqs", ctx_id, name, "val")
    r = rd.range(val_base)

    all_kvs = list(tr.get_range(r.start, r.stop))
    if not all_kvs:
        return _empty_result(columns)

    sampled_kvs = _pick_evenly_spaced(all_kvs, num_points)

    result = _empty_result(columns)
    for kv in sampled_kvs:
        step = rd.unpack(kv.key)[-1]
        result["steps"].append(step)
        if "val" in columns:
            result["val"].append(decode_value(kv.value))

    extra_cols = tuple(c for c in columns if c != "val")
    if extra_cols:
        result.update(_read_extra_columns(tr, rd, run_hash, ctx_id, name, result["steps"], extra_cols))

    return result


@transactional
def read_and_sample_sequence(
    tr: Transaction,
    run_hash: str,
    ctx_id: int,
    name: str,
    *,
    start_step: int | None = None,
    end_step: int | None = None,
    density: int | None = None,
    columns: tuple[str, ...] = ("val",),
) -> dict[str, list]:
    """Read a sequence with optional range filtering and downsampling."""
    data = read_sequence(
        tr,
        run_hash,
        ctx_id,
        name,
        start_step=start_step,
        end_step=end_step,
        columns=columns,
    )
    if density and density > 0 and len(data["steps"]) > density:
        sampled_steps = _pick_evenly_spaced(data["steps"], density)
        keep = {s: i for i, s in enumerate(data["steps"])}
        indices = [keep[s] for s in sampled_steps]
        for key in data:
            data[key] = [data[key][i] for i in indices]
    return data


@transactional
def get_sequence_step_bounds(
    tr: Transaction,
    run_hash: str,
    ctx_id: int,
    name: str,
) -> tuple[int | None, int | None]:
    """Return ``(first_step, last_step)`` for a sequence, or ``(None, None)`` if empty."""
    rd = _runs_dir()
    val_base = (run_hash, "seqs", ctx_id, name, "val")
    r = rd.range(val_base)

    first_kv = next(iter(tr.get_range(r.start, r.stop, limit=1)), None)
    if first_kv is None:
        return (None, None)
    last_kv = next(iter(tr.get_range(r.start, r.stop, limit=1, reverse=True)), None)
    first_step: int = rd.unpack(first_kv.key)[-1]
    last_step: int = rd.unpack(last_kv.key)[-1] if last_kv else first_step
    return (first_step, last_step)


@transactional
def get_sequence_length(tr: Transaction, run_hash: str, ctx_id: int, name: str) -> int:
    rd = _runs_dir()
    r = rd.range((run_hash, "seqs", ctx_id, name, "val"))
    count = 0
    for _ in tr.get_range(r.start, r.stop):
        count += 1
    return count
