"""Storage backends for different data formats.

Layers:
  - I/O layer:       save_* / load_* functions, read/write raw files (HDF5/bin)
                      return Dict[str, List[Tensor]] — format-specific, no state
  - Store (ABC):     central abstraction, normalizes multi-segment into
                      Dict[str, List[Tensor]] per key via _normalize(),
                      fetch() uses bisect across segments — no forced concat
  - Dataset layer:   BaseDataset owns a Store, only calls store.fetch(begin, end, key)

Key properties:
  - Multi-segment:   segments kept as-is, no forced concatenation — safe for
                      datasets larger than RAM
  - Explicit length: _length = min(total elements across keys), set at load,
                      __len__ returns O(1)
  - Zero-copy mmap:  MmapStore wraps np.memmap(mode="r"), all DataLoader
                       workers share OS page-cache pages
"""

import bisect
import glob
import json
import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Dict, List, Optional, Union

import torch
from torch import Tensor

from astrai.factory import BaseFactory
from astrai.preprocessing.transform import TokenizeTransform
from astrai.serialization import (
    load_bin,
    load_bin_offsets,
    load_h5,
)

logger = logging.getLogger(__name__)


def detect_format(load_path: str) -> str:
    """Auto-detect storage format from files in the directory.

    Args:
        load_path: Directory or file path

    Returns:
        Format string ("h5", "bin", or "jsonl")

    Raises:
        FileNotFoundError: If no supported data files are found
    """
    root = Path(load_path)
    if root.is_file():
        suffix = root.suffix.lower()
        if suffix in (".h5", ".hdf5"):
            return "h5"
        if suffix == ".jsonl":
            return "jsonl"
        raise ValueError(f"Unsupported file format: {suffix}")

    h5_files = [
        Path(p)
        for pattern in ("*.h5", "*.hdf5")
        for p in glob.glob(str(root / "**" / pattern), recursive=True)
    ]
    if h5_files:
        return "h5"
    bin_files = [Path(p) for p in glob.glob(str(root / "**" / "*.bin"), recursive=True)]
    if bin_files:
        has_meta = (root / "meta.json").exists() or len(
            [Path(p) for p in glob.glob(str(root / "**" / "meta.json"), recursive=True)]
        ) > 0
        if has_meta:
            return "bin"
    jsonl_files = [
        Path(p) for p in glob.glob(str(root / "**" / "*.jsonl"), recursive=True)
    ]
    if jsonl_files:
        return "jsonl"
    json_files = [
        Path(p) for p in glob.glob(str(root / "**" / "*.json"), recursive=True)
    ]
    if json_files:
        return "jsonl"
    raise FileNotFoundError(f"No supported data files found at {load_path}")


class Store(ABC):
    """String keys -> segmented tensors with two access modes.

    Stream mode (SEQ/SFT):
        ``fetch(begin, end, keys)`` slices across concatenated segments,
        transparently ``torch.cat``-ing across segment boundaries.
        ``len(store)`` returns total token count.

    Record mode (DPO/GRPO):
        ``fetch_record(index, keys)`` returns the i-th record without
        cross-record concatenation.  ``num_records`` returns the record
        count.  Backed by either per-record segment lists (H5/JSONL) or
        a single concatenated segment plus per-record offsets (bin).

    Subclasses declare ``segments_are_records`` to indicate whether their
    segments are inherently per-record (H5/JSONL) or opaque shards (bin).
    This is a format-level property, not a per-call decision.
    """

    segments_are_records: bool = False

    def __init__(self):
        self._data: Dict[str, List[Tensor]] = {}
        self._cum: Dict[str, List[int]] = {}
        self._offsets: Dict[str, List[int]] = {}
        self._length: int = 0
        self._num_records: int = 0

    @abstractmethod
    def load(self, path: str) -> None:
        raise NotImplementedError

    @property
    def keys(self) -> List[str]:
        return list(self._data.keys())

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

    def fetch(
        self,
        begin: int,
        end: int,
        keys: Union[str, List[str]],
    ):
        if not self._data:
            raise RuntimeError("Store not loaded")
        if not (0 <= begin < self._length and 0 <= end <= self._length):
            raise ValueError(
                f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
            )
        if isinstance(keys, str):
            return self._fetch_key(keys, begin, end)
        return {k: self._fetch_key(k, begin, end) for k in keys}

    def _fetch_key(self, key: str, begin: int, end: int) -> Tensor:
        """Fetch slice [begin, end) across potentially multiple segments."""
        segments = self._data[key]
        cum = self._cum[key]
        seg_start = bisect.bisect_right(cum, begin)
        seg_end = bisect.bisect_left(cum, end)

        results = []
        for i in range(seg_start, seg_end + 1):
            prev = cum[i - 1] if i > 0 else 0
            s = max(begin - prev, 0)
            e = min(end - prev, segments[i].shape[0])
            results.append(segments[i][s:e])

        return results[0] if len(results) == 1 else torch.cat(results, dim=0)

    @property
    def num_records(self) -> int:
        return self._num_records

    def fetch_record(
        self,
        index: int,
        keys: Union[str, List[str]],
    ):
        """Fetch the *index*-th record without cross-record concatenation.

        Returns a tensor (flat key) or ``List[Tensor]`` (nested key such as
        GRPO ``responses``).
        """
        if not self._data:
            raise RuntimeError("Store not loaded")
        if not 0 <= index < self._num_records:
            raise ValueError(
                f"Record index out of bounds: {index}, num_records={self._num_records}"
            )
        if isinstance(keys, str):
            return self._fetch_record_key(keys, index)
        return {k: self._fetch_record_key(k, index) for k in keys}

    def _fetch_record_key(self, key: str, index: int):
        """Return the *index*-th record for *key*.

        Two storage layouts are supported:

        - **bin + offsets**: ``_data[key]`` is ``[single_long_segment]``;
          ``_offsets[key]`` holds cumulative per-record offsets.  The record
          is sliced as ``segment[offsets[i]:offsets[i+1]]``.
        - **h5 / jsonl**: ``_data[key]`` is ``[t0, t1, ...]`` with one tensor
          (or nested list of tensors for GRPO) per record.  Direct indexing.
        """
        offsets = self._offsets.get(key)
        if offsets:
            start = offsets[index]
            end = (
                offsets[index + 1]
                if index + 1 < len(offsets)
                else self._data[key][0].shape[0]
            )
            return self._data[key][0][start:end]
        return self._data[key][index]

    def _normalize(
        self,
        raw: Dict[str, list],
        offsets: Optional[Dict[str, List[int]]] = None,
    ):
        """Register segments and pre-compute indices for both access modes.

        Stream mode: ``_cum[key]`` accumulates per-segment lengths so
        ``_fetch_key`` can bisect across segments without concatenation.

        Record mode: if *offsets* is provided (bin layout), ``_offsets[key]``
        stores cumulative per-record offsets into the single concatenated
        segment.  Otherwise, when ``segments_are_records`` is True
        (H5/JSONL), ``_data[key]`` is a per-record list and
        ``fetch_record`` indexes it directly.

        Nested keys (GRPO ``responses``/``masks`` as ``List[List[Tensor]]``)
        are stored as-is and excluded from both cumulative bookkeepings —
        they are only accessed record-by-record.
        """
        flat_lengths = []
        for key, tensors in raw.items():
            self._data[key] = tensors
            if not tensors:
                self._cum[key] = []
                flat_lengths.append(0)
                continue
            # Skip nested lists (GRPO responses/masks) — record-level access
            if isinstance(tensors[0], list):
                self._cum[key] = []
                continue
            cum = []
            total = 0
            for t in tensors:
                total += t.shape[0]
                cum.append(total)
            self._cum[key] = cum
            flat_lengths.append(cum[-1] if cum else 0)
        self._length = min(flat_lengths) if flat_lengths else 0

        # Record-mode offsets (bin layout).  Only valid when each key is a
        # single concatenated segment — multi-shard bin + offsets is not
        # supported (merge shards or use H5/JSONL instead).
        valid_offsets: Dict[str, List[int]] = {}
        if offsets:
            for key, off in offsets.items():
                segs = self._data.get(key, [])
                if len(segs) == 1 and len(off) > 1:
                    valid_offsets[key] = off
                elif len(segs) > 1:
                    logger.warning(
                        "Key '%s' has %d segments with offsets — record mode "
                        "disabled for this key (multi-shard bin+offsets not "
                        "supported). Merge shards or use H5/JSONL.",
                        key,
                        len(segs),
                    )
        self._offsets = valid_offsets
        if valid_offsets:
            record_counts = [len(v) - 1 for v in valid_offsets.values()]
            self._num_records = min(record_counts) if record_counts else 0
        elif self.segments_are_records:
            per_record_counts = []
            for key, tensors in self._data.items():
                if not tensors or isinstance(tensors[0], list):
                    continue
                per_record_counts.append(len(tensors))
            self._num_records = min(per_record_counts) if per_record_counts else 0
        else:
            self._num_records = 0


class StoreFactory(BaseFactory["Store"]):
    """Factory for creating Store instances by type name.

    Example::

        @StoreFactory.register("custom")
        class CustomStore(Store):
            ...
    """


@StoreFactory.register("h5")
class H5Store(Store):
    """HDF5-based storage backend (pre-tokenized data).

    Each key is stored as a group of per-record datasets (``data_0``,
    ``data_1``, …), so record mode indexes ``_data[key]`` directly.
    Stream mode concatenates across records via ``_cum``.
    """

    segments_are_records = True

    def load(self, path: str):
        self._normalize(load_h5(path))


@StoreFactory.register("bin")
class MmapStore(Store):
    """Memory-mapped binary storage backend.

    Each key is a single .bin file backed by ``np.memmap(mode="r")``.
    No per-process memory duplication — all DataLoader workers share the
    same OS page-cache pages.

    When ``meta.json`` contains per-record ``offsets`` for a key (written
    via ``save_bin(..., record_keys=...)``), record-mode access slices
    individual records from the concatenated memmap.  Legacy bin files
    without offsets only support stream mode.

    Format on disk::

        data_root/
          meta.json          # {key: {shape, dtype, offsets?}, ...}
          <key>.bin          # raw numpy array, one per key
    """

    def load(self, path: str):
        self._mmap_refs = []
        root = Path(path)
        all_raw: Dict[str, List[Tensor]] = {}
        all_offsets: Dict[str, List[int]] = {}
        meta_paths = [
            Path(p) for p in glob.glob(str(root / "**" / "meta.json"), recursive=True)
        ]
        for meta_path in meta_paths:
            raw = load_bin(str(meta_path.parent))
            off = load_bin_offsets(str(meta_path.parent))
            for key, tensors in raw.items():
                if key not in all_raw:
                    all_raw[key] = []
                all_raw[key].extend(tensors)
            for key, o in off.items():
                if key not in all_offsets:
                    all_offsets[key] = []
                all_offsets[key].extend(o)
        if not meta_paths:
            raise FileNotFoundError(f"No meta.json found under {path}")
        self._normalize(all_raw, offsets=all_offsets or None)
        for tensors in self._data.values():
            self._mmap_refs.extend(tensors)


@StoreFactory.register("jsonl")
class JsonlStore(Store):
    """JSONL reader with pluggable tokenization transform.

    A JSONL dataset directory contains ``*.jsonl`` files plus a
    ``dataset_config.json`` describing the tokenization pipeline.

    Responsibilities are split across two layers:

    - **Reader** (this class): reads raw JSON records from disk.
    - **Transform** (:class:`~astrai.preprocessing.transform.TokenizeTransform`):
      tokenizes records into per-key tensors.  When not supplied explicitly,
      a default transform is built from ``dataset_config.json`` so that
      existing callers keep working without changes.

    The Store itself never imports the tokenizer — the dependency lives
    in the Transform layer.
    """

    CONFIG_NAME = "dataset_config.json"
    segments_are_records = True

    def load(self, path: str, transform=None, **kwargs):
        root = Path(path)
        records = self._read_records(root)

        if transform is None:
            config_path = root / self.CONFIG_NAME
            if not config_path.exists():
                raise FileNotFoundError(
                    f"JSONL dataset config not found: {config_path}. "
                    f"Expected {self.CONFIG_NAME} alongside *.jsonl files, "
                    f"or pass an explicit transform."
                )
            transform = TokenizeTransform.from_config_file(str(config_path))

        raw = transform.apply(records)
        self._normalize(raw)

    @staticmethod
    def _read_records(root: Path) -> List[dict]:
        records: List[dict] = []
        for jsonl_path in sorted(root.glob("*.jsonl")):
            with open(jsonl_path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        records.append(json.loads(line))
                    except json.JSONDecodeError:
                        logger.warning(
                            "Failed to parse JSON line in %s, skipping", jsonl_path
                        )
        for json_path in sorted(root.glob("*.json")):
            if json_path.name == JsonlStore.CONFIG_NAME:
                continue
            with open(json_path, "r", encoding="utf-8") as f:
                try:
                    data = json.load(f)
                except json.JSONDecodeError:
                    logger.warning("Failed to parse JSON file %s, skipping", json_path)
                    continue
            if isinstance(data, list):
                records.extend(data)
            elif isinstance(data, dict):
                records.append(data)
        return records