"""CIFAR-10 dataset that stays within NumPy/JAX dependencies."""
from __future__ import annotations

import pickle
import tarfile
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

import jax
import jax.numpy as jnp
import numpy as np

from ..dataset_protocol import DatasetProtocol
from ..sources import DiskSampleSource

CIFAR10_URL = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"


def _to_host_jax_array(array: np.ndarray) -> jax.Array:
    cpu_devices = jax.devices("cpu")
    if cpu_devices:
        with jax.default_device(cpu_devices[0]):
            return jnp.asarray(array)
    return jnp.asarray(array)


def _resolve_cache_dir(cache_dir: str | Path | None) -> Path:
    base_dir = Path(cache_dir) if cache_dir is not None else Path.home() / ".cache" / "cyreal_cifar10"
    base_dir.mkdir(parents=True, exist_ok=True)
    return base_dir


def _download(url: str, path: Path) -> Path:
    if path.exists():
        return path
    path.parent.mkdir(parents=True, exist_ok=True)
    urllib.request.urlretrieve(url, path)
    return path


def _ensure_extracted(archive: Path, extract_root: Path) -> Path:
    target = extract_root / "cifar-10-batches-py"
    if target.exists():
        return target
    with tarfile.open(archive, "r:gz") as tar:
        tar.extractall(path=extract_root)
    return target


def _batch_names(split: Literal["train", "test"]) -> list[str]:
    if split == "train":
        return [f"data_batch_{i}" for i in range(1, 6)]
    if split == "test":
        return ["test_batch"]
    raise ValueError("split must be 'train' or 'test'.")


def _load_split_numpy(split: Literal["train", "test"], extract_dir: Path) -> tuple[np.ndarray, np.ndarray]:
    images = []
    labels = []
    for name in _batch_names(split):
        batch_path = extract_dir / name
        if not batch_path.exists():
            raise FileNotFoundError(f"Missing CIFAR-10 batch '{name}'.")
        with open(batch_path, "rb") as f:
            batch = pickle.load(f, encoding="latin1")
        data = batch["data"].reshape(-1, 3, 32, 32)
        images.append(np.transpose(data, (0, 2, 3, 1)).astype(np.uint8))
        labels.append(np.asarray(batch["labels"], dtype=np.int32))

    images_np = np.concatenate(images, axis=0)
    labels_np = np.concatenate(labels, axis=0)
    return images_np, labels_np


def _ensure_split_numpy_cache(
    split: Literal["train", "test"],
    base_dir: Path,
    extract_dir: Path,
) -> tuple[Path, Path]:
    cache_root = base_dir / "disk_cache"
    cache_root.mkdir(parents=True, exist_ok=True)
    images_path = cache_root / f"{split}_images.npy"
    labels_path = cache_root / f"{split}_labels.npy"
    if not images_path.exists() or not labels_path.exists():
        images_np, labels_np = _load_split_numpy(split, extract_dir)
        np.save(images_path, images_np)
        np.save(labels_path, labels_np)
    return images_path, labels_path


@dataclass
class CIFAR10Dataset(DatasetProtocol):
    """Download-free CIFAR-10 access ready for `ArraySampleSource`."""

    split: Literal["train", "test"] = "train"
    cache_dir: str | Path | None = None

    def __post_init__(self) -> None:
        base_dir = _resolve_cache_dir(self.cache_dir)
        archive_path = base_dir / "cifar-10-python.tar.gz"
        _download(CIFAR10_URL, archive_path)
        extract_dir = _ensure_extracted(archive_path, base_dir)

        images_np, labels_np = _load_split_numpy(self.split, extract_dir)
        self._images = _to_host_jax_array(images_np)
        self._labels = _to_host_jax_array(labels_np)

    def __len__(self) -> int:
        return int(self._images.shape[0])

    def __getitem__(self, index: int):
        return {
            "image": self._images[index],
            "label": self._labels[index],
        }

    def as_array_dict(self) -> dict[str, jax.Array]:
        """Expose the full dataset as a PyTree of JAX arrays."""

        return {
            "image": self._images,
            "label": self._labels,
        }

    @classmethod
    def make_disk_source(
        cls,
        *,
        split: Literal["train", "test"] = "train",
        cache_dir: str | Path | None = None,
        ordering: Literal["sequential", "shuffle"] = "shuffle",
        prefetch_size: int = 64,
    ) -> DiskSampleSource:
        base_dir = _resolve_cache_dir(cache_dir)
        archive_path = base_dir / "cifar-10-python.tar.gz"
        _download(CIFAR10_URL, archive_path)
        extract_dir = _ensure_extracted(archive_path, base_dir)

        images_path, labels_path = _ensure_split_numpy_cache(split, base_dir, extract_dir)
        images_memmap = np.load(images_path, mmap_mode="r")
        labels_memmap = np.load(labels_path, mmap_mode="r")

        if images_memmap.shape[0] != labels_memmap.shape[0]:
            raise ValueError("CIFAR-10 image and label counts do not match.")

        def _read_sample(index: int | np.ndarray) -> dict[str, np.ndarray]:
            idx = int(np.asarray(index))
            image = np.asarray(images_memmap[idx], dtype=np.uint8)
            label = np.asarray(labels_memmap[idx], dtype=np.int32)
            return {"image": image, "label": label}

        sample_spec = {
            "image": jax.ShapeDtypeStruct(shape=tuple(images_memmap.shape[1:]), dtype=jnp.uint8),
            "label": jax.ShapeDtypeStruct(shape=(), dtype=jnp.int32),
        }

        return DiskSampleSource(
            length=int(labels_memmap.shape[0]),
            sample_fn=_read_sample,
            sample_spec=sample_spec,
            ordering=ordering,
            prefetch_size=prefetch_size,
        )


@dataclass
class CIFAR10DiskSource:
    """Stream CIFAR-10 samples directly from cached numpy arrays on disk."""

    split: Literal["train", "test"] = "train"
    cache_dir: str | Path | None = None
    ordering: Literal["sequential", "shuffle"] = "shuffle"
    prefetch_size: int = 64

    def __post_init__(self) -> None:
        self._disk_source = CIFAR10Dataset.make_disk_source(
            split=self.split,
            cache_dir=self.cache_dir,
            ordering=self.ordering,
            prefetch_size=self.prefetch_size,
        )
        self.steps_per_epoch = self._disk_source.steps_per_epoch

    def element_spec(self):
        return self._disk_source.element_spec()

    def init_state(self, key=None):
        return self._disk_source.init_state(key)

    def next(self, state):
        return self._disk_source.next(state)
