import pickle
from pathlib import Path
from typing import Any

from arti.formats.pickle import Pickle
from arti.io import read, write
from arti.storage.local import LocalFilePartition
from arti.types import Type
from arti.views.python import PythonBuiltin


def _read_pickle_file(path: str) -> Any:
    with open(path, "rb") as file:
        return pickle.load(file)


# TODO: Need to handle partitioned lists. Should I separate read and read_partitions?


@read.register
def _read_pickle_localfile_python(
    type: Type,
    format: Pickle,
    storage_partitions: tuple[LocalFilePartition, ...],
    view: PythonBuiltin,
) -> Any:
    return [_read_pickle_file(storage_partition.path) for storage_partition in storage_partitions]


@write.register
def _write_pickle_localfile_python(
    data: Any,
    type: Type,
    format: Pickle,
    storage_partition: LocalFilePartition,
    view: PythonBuiltin,
) -> None:
    path = Path(storage_partition.path)
    path.parent.mkdir(exist_ok=True, parents=True)
    with path.open("wb") as file:
        pickle.dump(data, file)
