from __future__ import annotations

from dataclasses import asdict, is_dataclass
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    # peer dependencies that are used for types only
    from torch.utils.data import DataLoader as TorchDataLoader  # type: ignore
    from torch.utils.data import Dataset as TorchDataset  # type: ignore


def parse_dict_like(item: Any, column_names: list[str] | None = None) -> dict:
    if isinstance(item, dict):
        return item

    if isinstance(item, tuple):
        if column_names is not None:
            if len(item) != len(column_names):
                raise ValueError(
                    f"Tuple length ({len(item)}) does not match number of column names ({len(column_names)})"
                )
            return {column_names[i]: item[i] for i in range(len(item))}
        elif hasattr(item, "_fields") and all(isinstance(field, str) for field in item._fields):  # type: ignore
            return {field: getattr(item, field) for field in item._fields}  # type: ignore
        else:
            raise ValueError("For datasets that return unnamed tuples, please provide column_names argument")

    if is_dataclass(item) and not isinstance(item, type):
        return asdict(item)

    raise ValueError(f"Cannot parse {type(item)}")


def parse_batch(batch: Any, column_names: list[str] | None = None) -> list[dict]:
    if isinstance(batch, list):
        return [parse_dict_like(item, column_names) for item in batch]

    batch = parse_dict_like(batch, column_names)
    keys = list(batch.keys())
    batch_size = len(batch[keys[0]])
    for key in keys:
        if not len(batch[key]) == batch_size:
            raise ValueError(f"Batch must consist of values of the same length, but {key} has length {len(batch[key])}")
    return [{key: batch[key][idx] for key in keys} for idx in range(batch_size)]


def list_from_torch(
    torch_data: TorchDataLoader | TorchDataset,
    column_names: list[str] | None = None,
) -> list[dict]:
    """
    Convert a PyTorch DataLoader or Dataset to a list of dictionaries.

    Params:
        torch_data: A PyTorch DataLoader or Dataset object to convert.
        column_names: Optional list of column names to use for the data. If not provided,
            the column names will be inferred from the data.
    Returns:
        A list of dictionaries containing the data from the PyTorch DataLoader or Dataset.
    """
    # peer dependency that is guaranteed to exist if the user provided a torch dataset
    from torch.utils.data import DataLoader as TorchDataLoader  # type: ignore

    if isinstance(torch_data, TorchDataLoader):
        dataloader = torch_data
    else:
        dataloader = TorchDataLoader(torch_data, batch_size=1, collate_fn=lambda x: x)

    # Collect data from the dataloader into a list
    data_list = []
    try:
        for batch in dataloader:
            data_list.extend(parse_batch(batch, column_names=column_names))
    except ValueError as e:
        raise ValueError(str(e)) from e

    return data_list
