import logging
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import Tensor

from pipeline.utils import error_handler

logger = logging.getLogger(__name__)


class SequencePacker:
    """
    Stream-concatenation packer for LLM training sequences.

    Algorithm (streaming concat):

        Input:  sequences = [A(len=3), B(len=5), C(len=2)], pack_size = 6

        1. Validate & Normalize
           - Check 1D dimension, unify dtype, warn on overlong sequences
           - Result: [A, B, C]

        2. Stream into buffer, slice off full chunks
           - buffer += A(3) -> [a1 a2 a3], pos=3
           - buffer += B(5) -> [a1 a2 a3 b1 b2 b3 b4 b5], pos=8
             pos >= 6 -> flush [a1 a2 a3 b1 b2 b3], buffer=[b4 b5], pos=2
           - buffer += C(2) -> [b4 b5 c1 c2], pos=4
             loop ends -> flush tail [b4 b5 c1 c2 PAD PAD]

        Output: [[a1 a2 a3 b1 b2 b3], [b4 b5 c1 c2 PAD PAD]]

    Samples may be split across chunks — this is intentional and standard
    practice in LLM training (TRL, Megatron-LM, etc.).

    Cross-group consistency:
        Different tensor groups (e.g. input_ids, loss_masks) packed with
        separate packer instances on samples with matching lengths produce
        identical chunk boundaries. Element-level correspondence is preserved.
    """

    def __init__(
        self,
        pack_size: int,
        pad_value: Union[int, bool] = 0,
        dtype: Optional[torch.dtype] = None,
    ):
        self.pack_size = pack_size
        self.pad_value = pad_value
        self.dtype = dtype
        self._buffer: List = []
        self._pos: int = 0
        self._packages: List[Tensor] = []

    def reset(self) -> None:
        """Reset packer state for instance reuse."""
        self._buffer = []
        self._pos = 0
        self._packages = []

    @error_handler()
    def pack(self, sequences: List[Tensor]) -> List[Tensor]:
        """
        Pack sequences via streaming concatenation into fixed-size chunks.

        Sequences are concatenated in order and sliced at pack_size boundaries.
        The final chunk is padded with pad_value.

        When dtype is not set at init, it is inferred from the first input tensor.

        Args:
            sequences: List of 1D input tensors.

        Returns:
            List of packed tensors, each with length equal to pack_size.
        """
        if not sequences:
            return []

        # --- auto-infer dtype from first sequence ---
        if self.dtype is None:
            self.dtype = sequences[0].dtype

        # --- validate & normalize ---
        normalized: List[Tensor] = []
        for i, seq in enumerate(sequences):
            if seq.dim() != 1:
                raise ValueError(
                    f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}"
                )
            if seq.dtype != self.dtype:
                seq = seq.to(self.dtype)
            normalized.append(seq)

        # --- stream into buffer, slice off full chunks ---
        self._buffer = []
        self._packages = []
        pack_size = self.pack_size
        buf = self._buffer

        for seq in normalized:
            buf.extend(seq.tolist())
            while len(buf) >= pack_size:
                self._packages.append(torch.tensor(buf[:pack_size], dtype=self.dtype))
                buf = buf[pack_size:]

        # flush tail with padding
        if buf:
            padded = buf + [self.pad_value] * (pack_size - len(buf))
            self._packages.append(torch.tensor(padded, dtype=self.dtype))

        self._pos = len(buf)
        return self._packages


def pack_tensors(
    tensors: Dict[str, List[Tensor]],
    pack_size: int,
    pad_value: Union[int, bool] = 0,
    dtypes: Optional[Dict[str, torch.dtype]] = None,
) -> Dict[str, List[Tensor]]:
    """
    Pack multiple named tensor groups in parallel.

    Each group is packed independently with its own SequencePacker instance.
    When dtypes is provided, packers use the declared dtype per key;
    otherwise dtype is auto-inferred from the first tensor in each group.

    Args:
        tensors: Dict mapping key names to lists of 1D tensors.
        pack_size: Fixed chunk length.
        pad_value: Padding value for non-bool tensors.
        dtypes: Optional per-key dtype declarations.

    Returns:
        Dict mapping key names to lists of packed tensors.
    """
    if dtypes is None:
        dtypes = {}

    output: Dict[str, List[Tensor]] = {}
    for key, seqs in tensors.items():
        dtype = dtypes.get(key)
        packer = SequencePacker(pack_size, pad_value, dtype=dtype)
        output[key] = packer.pack(seqs)
    return output