import ctypes
import gc
import logging
import os
import random

import numpy as np
import pyarrow as pa
import pyarrow.feather as feather
import torch

logger = logging.getLogger(__name__)


def _release_memory():
    """Force Python GC and return freed memory to OS via glibc malloc_trim."""
    gc.collect()
    try:
        ctypes.CDLL("libc.so.6").malloc_trim(0)
    except (OSError, AttributeError):
        pass


class _CompactBatch:
    """Memory-efficient batch backed by numpy arrays (optionally mmap'd via Arrow).

    Replaces the old ``list[dict]`` representation.  Per-sample access via
    ``__getitem__`` returns the same dict shape the rest of the pipeline
    expects, but ``values`` is a numpy 1-D view instead of a Python list,
    cutting per-shard memory from ~1.2 GB of Python objects to ~256 MB of
    contiguous float64 (or zero additional bytes when mmap-backed).
    """

    __slots__ = (
        "_values",
        "_lengths",
        "_generator_types",
        "_starts",
        "_frequencies",
        "_num_channels",
        "_series_ids",
        "_size",
        "_table_ref",
    )

    def __init__(
        self,
        values: np.ndarray,
        lengths: list,
        generator_types: list,
        starts: list,
        frequencies: list,
        num_channels: list,
        series_ids: list,
        table_ref=None,
    ):
        self._values = values            # (n, seq_len) float64
        self._lengths = lengths
        self._generator_types = generator_types
        self._starts = starts
        self._frequencies = frequencies
        self._num_channels = num_channels
        self._series_ids = series_ids
        self._size = len(series_ids)
        self._table_ref = table_ref      # prevents mmap GC while views exist

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

    def __getitem__(self, idx: int) -> dict:
        nc = self._num_channels[idx]
        if nc == 1:
            vals = self._values[idx]     # numpy 1-D view – zero-copy
        else:
            seq_len = self._values.shape[1] // nc
            vals = self._values[idx].reshape(nc, seq_len).tolist()  # rare path
        return {
            "series_id": self._series_ids[idx],
            "values": vals,
            "length": self._lengths[idx],
            "generator_type": self._generator_types[idx],
            "start": self._starts[idx],
            "frequency": self._frequencies[idx],
            "num_channels": nc,
        }

    def __del__(self):
        self._table_ref = None
        self._values = None


class CyclicalBatchDataset:
    """
    Dataset class that loads saved batches from continuous generation script.
    Maintains a pointer and provides cyclical access to individual samples.
    Includes enhanced logging to track data shard cycling during training.
    Supports per-rank file sharding for large-scale distributed training.
    """

    def __init__(
        self,
        batches_dir: str,
        generator_type: str,
        device: torch.device | None = None,
        prefetch_next: bool = True,
        prefetch_threshold: int = 32,
        rank: int = 0,
        world_size: int = 1,
    ):
        self.batches_dir = batches_dir
        self.generator_type = generator_type
        self.device = device
        self.prefetch_next = prefetch_next
        self.prefetch_threshold = prefetch_threshold
        self.rank = rank
        self.world_size = world_size

        self.batch_files = self._find_batch_files()
        if not self.batch_files:
            raise ValueError(f"No batch files found in {batches_dir}")

        self.current_batch_idx = 0
        self.current_sample_idx = 0
        self.current_batch_data: _CompactBatch | None = None
        self.next_batch_data: _CompactBatch | None = None
        self.prefetching_in_progress = False

        self.visited_batch_indices: set[int] = set()
        self.full_cycles_completed = 0

        self._load_current_batch()
        self.visited_batch_indices.add(self.current_batch_idx)

        logger.info(
            f"Initialized '{self.generator_type}' dataset with {len(self.batch_files)} batches. "
            f"Current batch file: '{os.path.basename(self.batch_files[self.current_batch_idx])}' "
            f"has {len(self.current_batch_data)} samples."
        )

    # ------------------------------------------------------------------
    # File discovery
    # ------------------------------------------------------------------

    def _find_batch_files(self) -> list[str]:
        import glob

        pattern = os.path.join(self.batches_dir, "batch_*.arrow")
        all_files = sorted(glob.glob(pattern))

        if not all_files:
            return []

        rank_files = [f for i, f in enumerate(all_files) if i % self.world_size == self.rank]
        random.shuffle(rank_files)

        logger.info(
            f"[Rank {self.rank}] '{self.generator_type}': Sharded {len(all_files)} files → "
            f"{len(rank_files)} files for this rank ({len(rank_files) / len(all_files) * 100:.1f}%)"
        )
        return rank_files

    # ------------------------------------------------------------------
    # Arrow → numpy zero-copy loading
    # ------------------------------------------------------------------

    def _load_batch_from_file(self, batch_file: str) -> _CompactBatch:
        """Load a batch from an Arrow/Feather file using memory-mapped zero-copy reads.

        The ``values`` column is extracted as a contiguous numpy array via
        Arrow's mmap and ``to_numpy``.  When the data lives on tmpfs
        (``/dev/shm``), the OS shares the physical pages across all workers
        that mmap the same file — eliminating duplicate copies of the time
        series data.
        """
        try:
            table = feather.read_table(batch_file, memory_map=True)
            n = len(table)
            has_num_channels = "num_channels" in table.column_names

            # --- values column → contiguous numpy (the big win) ----------
            values_col = table.column("values")
            # ChunkedArray → single array (usually a no-op for Feather)
            if values_col.num_chunks > 1:
                values_arr = values_col.combine_chunks()
            else:
                values_arr = values_col.chunk(0)

            keep_table = False  # assume we'll copy; flip to True for zero-copy

            if isinstance(values_arr.type, pa.ListType) and not isinstance(
                values_arr.type.value_type, pa.ListType
            ):
                # Univariate: list<double> — flat child is contiguous
                flat_child = values_arr.values
                try:
                    flat_np = flat_child.to_numpy(zero_copy_only=True)
                    keep_table = True
                except (pa.ArrowInvalid, pa.ArrowTypeError, pa.ArrowNotImplementedError):
                    flat_np = flat_child.to_numpy(zero_copy_only=False)

                offsets_np = values_arr.offsets.to_numpy()
                row_lengths = np.diff(offsets_np)

                if n > 0 and np.all(row_lengths == row_lengths[0]):
                    seq_len = int(row_lengths[0])
                    values_np = flat_np.reshape(n, seq_len)
                else:
                    max_len = int(row_lengths.max()) if n > 0 else 0
                    values_np = np.full((n, max_len), np.nan, dtype=np.float64)
                    for i in range(n):
                        s, e = int(offsets_np[i]), int(offsets_np[i + 1])
                        values_np[i, : e - s] = flat_np[s:e]
                    keep_table = False
            else:
                # Multivariate or unexpected schema — safe fallback
                values_list = values_col.to_pylist()
                max_flat = max(
                    (len(v) if not isinstance(v[0], list) else sum(len(c) for c in v))
                    for v in values_list
                ) if n > 0 else 0
                values_np = np.full((n, max_flat), np.nan, dtype=np.float64)
                for i, v in enumerate(values_list):
                    flat = v if not isinstance(v[0], list) else [x for ch in v for x in ch]
                    values_np[i, : len(flat)] = flat
                keep_table = False

            # --- lightweight metadata (tiny compared to values) ----------
            series_ids = table.column("series_id").to_pylist()
            lengths = table.column("length").to_pylist()
            generator_types = table.column("generator_type").to_pylist()
            starts = table.column("start").to_pylist()
            frequencies = table.column("frequency").to_pylist()
            num_channels = table.column("num_channels").to_pylist() if has_num_channels else [1] * n

            batch = _CompactBatch(
                values=values_np,
                lengths=lengths,
                generator_types=generator_types,
                starts=starts,
                frequencies=frequencies,
                num_channels=num_channels,
                series_ids=series_ids,
                table_ref=table if keep_table else None,
            )

            if not keep_table:
                del table
                _release_memory()

            return batch

        except Exception as e:
            logger.error(f"Error loading batch from {batch_file}: {e}")
            raise

    # ------------------------------------------------------------------
    # Batch lifecycle
    # ------------------------------------------------------------------

    def _load_current_batch(self):
        if self.current_batch_data is not None:
            del self.current_batch_data
            self.current_batch_data = None
            _release_memory()
        batch_file = self.batch_files[self.current_batch_idx]
        self.current_batch_data = self._load_batch_from_file(batch_file)
        self.current_sample_idx = 0
        logger.debug(
            f"Loaded batch {self.current_batch_idx} for {self.generator_type} "
            f"with {len(self.current_batch_data)} samples"
        )

    def _trigger_smart_prefetch(self):
        if not self.prefetch_next or len(self.batch_files) <= 1:
            return
        remaining_samples = self.get_remaining_samples_in_current_batch()
        should_prefetch = (
            remaining_samples <= self.prefetch_threshold
            and self.next_batch_data is None
            and not self.prefetching_in_progress
        )
        if should_prefetch:
            self._prefetch_next_batch()

    def _prefetch_next_batch(self):
        if self.prefetching_in_progress:
            return
        self.prefetching_in_progress = True
        next_batch_idx = (self.current_batch_idx + 1) % len(self.batch_files)
        next_batch_file = self.batch_files[next_batch_idx]
        try:
            self.next_batch_data = self._load_batch_from_file(next_batch_file)
            logger.debug(f"Prefetched next batch {next_batch_idx} for {self.generator_type}")
        except Exception as e:
            logger.warning(f"Failed to prefetch batch {next_batch_idx}: {e}")
            self.next_batch_data = None
        finally:
            self.prefetching_in_progress = False

    def _advance_to_next_batch(self):
        if self.current_batch_data is not None:
            del self.current_batch_data
            self.current_batch_data = None
            _release_memory()

        previous_batch_idx = self.current_batch_idx
        self.current_batch_idx = (self.current_batch_idx + 1) % len(self.batch_files)

        if self.next_batch_data is not None:
            self.current_batch_data = self.next_batch_data
            self.next_batch_data = None
        else:
            self._load_current_batch()

        self.current_sample_idx = 0
        self.prefetching_in_progress = False

        self.visited_batch_indices.add(self.current_batch_idx)

        total_files = len(self.batch_files)
        visited_count = len(self.visited_batch_indices)
        progress_percent = (visited_count / total_files) * 100

        logger.info(
            f"\nDATA SHARD CYCLED for '{self.generator_type}': "
            f"Moved from file index {previous_batch_idx} to {self.current_batch_idx}. "
            f"Unique files visited: {visited_count}/{total_files} ({progress_percent:.1f}%)."
        )

        if visited_count == total_files:
            self.full_cycles_completed += 1
            logger.info(
                f"FULL CYCLE #{self.full_cycles_completed} COMPLETED for '{self.generator_type}'! "
                f"All {total_files} data files have been visited at least once. "
                "Resetting visited set to track the next cycle."
            )
            self.visited_batch_indices.clear()
            self.visited_batch_indices.add(self.current_batch_idx)

    # ------------------------------------------------------------------
    # Public sample access
    # ------------------------------------------------------------------

    def get_sample(self) -> dict:
        if self.current_batch_data is None:
            self._load_current_batch()
        if self.current_batch_data is None:
            raise RuntimeError("No batch data loaded")
        if self.current_sample_idx >= len(self.current_batch_data):
            self._advance_to_next_batch()
        self._trigger_smart_prefetch()
        sample = self.current_batch_data[self.current_sample_idx]
        self.current_sample_idx += 1
        return sample

    def get_samples(self, num_samples: int) -> list[dict]:
        return [self.get_sample() for _ in range(num_samples)]

    def get_total_samples_in_current_batch(self) -> int:
        if self.current_batch_data is None:
            return 0
        return len(self.current_batch_data)

    def get_remaining_samples_in_current_batch(self) -> int:
        if self.current_batch_data is None:
            return 0
        return max(0, len(self.current_batch_data) - self.current_sample_idx)

    def get_info(self) -> dict:
        total_files = len(self.batch_files)
        visited_count = len(self.visited_batch_indices)
        return {
            "generator_type": self.generator_type,
            "total_batch_files": total_files,
            "current_batch_idx": self.current_batch_idx,
            "current_sample_idx": self.current_sample_idx,
            "current_batch_size": self.get_total_samples_in_current_batch(),
            "remaining_in_batch": self.get_remaining_samples_in_current_batch(),
            "unique_files_visited": visited_count,
            "cycle_progress_percent": (visited_count / total_files) * 100 if total_files > 0 else 0,
            "full_cycles_completed": self.full_cycles_completed,
        }