"""Yaozarrs-based backends for OME-Zarr v0.5."""

from __future__ import annotations

import json
import threading
import warnings
from abc import abstractmethod
from collections.abc import Iterator, MutableMapping
from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, cast

from ome_writers import __version__
from ome_writers._backends._backend import ArrayBackend
from ome_writers._backends._chunk_buffer import ChunkBuffer
from ome_writers._schema import is_channel_dim

try:
    from yaozarrs import DimSpec, v05
    from yaozarrs.write.v05 import Bf2RawBuilder, PlateBuilder, prepare_image
except ImportError as e:
    raise ImportError(
        f"{__name__} requires yaozarrs with write support: "
        "`pip install yaozarrs[write-zarr]`."
    ) from e


if TYPE_CHECKING:
    from collections.abc import Callable, Sequence
    from typing import Protocol

    import numpy as np
    from yaozarrs.write.v05._write import CompressionName

    from ome_writers._router import FrameRouter
    from ome_writers._schema import (
        AcquisitionSettings,
        Dimension,
        Plate,
        Position,
    )

    class ArrayLike(Protocol):
        """Protocol for array-like objects that support shape attribute."""

        @property
        def shape(self) -> tuple[int, ...]: ...

        def __setitem__(self, index: Any, val: Any) -> Any: ...

        def resize(self, new_shape: Sequence[int]) -> None: ...


# Array type generated by yaozarrs writers
_AT = TypeVar("_AT", bound="ArrayLike")
YAOZARRS_COMPRESSION_OPTIONS = {"blosc-zstd", "blosc-lz4", "zstd", "none"}


class YaozarrsBackend(ArrayBackend, Generic[_AT]):
    """Base backend using yaozarrs for OME-Zarr v0.5 structure.

    Subclasses must define the `writer` class attribute to specify which
    yaozarrs writer to use (e.g., "zarr", "tensorstore").
    """

    def __init__(self) -> None:
        self._arrays: list[_AT] = []
        self._image_group_paths: list[str] = []  # Parallel to _arrays, for metadata
        self._meta_mirrors: dict[str, JsonDocumentMirror] = {}  # grouppath -> mirror
        self._finalized = False
        self._root: Path | None = None
        self._chunk_buffers: list[ChunkBuffer] | None = None

    @abstractmethod
    def _get_yaozarrs_writer(
        self,
    ) -> Literal["zarr", "tensorstore"] | Callable[..., Any]:
        """Return the writer to use for array creation.

        Subclasses can override to provide a custom CreateArrayFunc.
        """

    def _post_prepare(self, settings: AcquisitionSettings) -> None:
        """Hook called after yaozarrs creates the structure, before metadata caching.

        Subclasses can override to do additional setup (e.g., create streams).
        """

    def is_incompatible(self, settings: AcquisitionSettings) -> Literal[False] | str:
        """Check if settings are incompatible with yaozarrs backend.

        Returns a string describing the incompatibility, or False if compatible.
        """
        cls_name = self.__class__.__name__

        # Check for supported compression
        compression = settings.compression
        if compression and compression not in YAOZARRS_COMPRESSION_OPTIONS:
            return (  # pragma: no cover
                f"Compression '{compression}' is not supported by {cls_name}. "
                f"Supported: {YAOZARRS_COMPRESSION_OPTIONS}."
            )

        return False

    def prepare(self, settings: AcquisitionSettings, router: FrameRouter) -> None:
        """Initialize OME-Zarr storage structure."""

        self._finalized = False
        self._root = root = Path(settings.output_path).expanduser().resolve()
        positions = settings.positions

        # Build storage metadata
        storage_dims = settings.array_storage_dimensions
        shape = tuple(d.count if d.count is not None else 1 for d in storage_dims)
        dtype = settings.dtype
        chunks, shards = _get_chunks_and_shards(storage_dims)
        # this single image model is reused for all positions
        # (the underlying assumption is that we currently don't support inhomogeneous
        # shapes/dtypes across positions)
        image = _build_yaozarrs_image_model(storage_dims)

        compression = cast("CompressionName", settings.compression or "none")

        # Get writer from hook (subclasses can override)
        writer = self._get_yaozarrs_writer()

        # Plate mode
        if settings.plate is not None:
            # mapping of {(row, column): [(position_index, Position), ...]}
            well_positions = {}
            for i, pos in enumerate(positions):
                key = (pos.plate_row, pos.plate_column)
                well_positions.setdefault(key, []).append((i, pos))

            # Build plate metadata and arrays
            builder = PlateBuilder(
                root,
                plate=_build_yaozarrs_plate_model(settings.plate, well_positions),
                overwrite=settings.overwrite,
                chunks=chunks,
                shards=shards,
                writer=writer,
                compression=compression,
            )

            for (row, col), pos_list in well_positions.items():
                images_dict = {
                    pos.name: (image, [(shape, dtype)]) for _idx, pos in pos_list
                }
                builder.add_well(row=row, col=col, images=images_dict)

            _, all_arrays = builder.prepare()
            # Map position index to array path and store arrays
            self._image_group_paths = [
                f"{pos.plate_row}/{pos.plate_column}/{pos.name}" for pos in positions
            ]
            self._arrays = [
                all_arrays[f"{parent}/0"] for parent in self._image_group_paths
            ]

        # Single position
        elif len(positions) == 1:
            _, all_arrays = prepare_image(
                root,
                image,
                datasets=[(shape, dtype)],
                overwrite=settings.overwrite,
                chunks=chunks,
                shards=shards,
                writer=writer,
                compression=compression,
            )
            self._image_group_paths = ["."]
            self._arrays = [all_arrays["0"]]  # ty: ignore[invalid-assignment]

        # Multi-position (bf2raw layout)
        else:
            builder = Bf2RawBuilder(
                root,
                overwrite=settings.overwrite,
                chunks=chunks,
                shards=shards,
                writer=writer,
                compression=compression,
            )
            for pos in positions:
                builder.add_series(_get_series_name(pos), image, [(shape, dtype)])

            _, all_arrays = builder.prepare()
            self._image_group_paths = [_get_series_name(pos) for pos in positions]
            self._arrays = [
                all_arrays[f"{parent}/0"] for parent in self._image_group_paths
            ]

        # Post-prepare hook (subclasses can do additional work)
        self._post_prepare(settings)

        # Cache metadata immediately after creation
        # This is used later for get_metadata() and update_metadata()
        self._cache_metadata_from_arrays(root)

        # Initialize chunk buffering if beneficial
        if self._should_use_chunk_buffering(storage_dims):
            # Extract shapes for buffer initialization
            index_dims = storage_dims[:-2]  # Exclude frame dims
            frame_dims = storage_dims[-2:]  # Y, X

            index_shape = tuple(d.count for d in index_dims)
            chunk_shape = tuple(
                d.chunk_size if d.chunk_size is not None else 1 for d in index_dims
            )
            frame_shape = tuple(d.count for d in frame_dims)

            # Create one buffer per position
            self._chunk_buffers = [
                ChunkBuffer(
                    index_shape=index_shape,
                    chunk_shape=chunk_shape,
                    frame_shape=frame_shape,
                    dtype=dtype,
                )
                for _ in range(len(positions))
            ]
        else:
            self._chunk_buffers = None

    def _ensure_array_size(
        self, position_index: int, indices: Sequence[tuple[int, ...]]
    ) -> None:
        """Resize array if needed to accommodate the given indices.

        This method calculates the maximum extent across all provided indices
        and resizes the array once if needed. It also updates chunk buffer
        index_shape for buffered backends.
        """
        array = self._arrays[position_index]
        current_shape = list(array.shape)
        new_shape = current_shape.copy()

        # Find maximum extent across all indices
        for storage_idx in indices:
            for i, idx_val in enumerate(storage_idx):
                new_shape[i] = max(new_shape[i], idx_val + 1)

        # Resize only once if needed
        if new_shape != current_shape:
            self._resize(array, new_shape)
            if self._chunk_buffers:
                self._chunk_buffers[position_index].index_shape = new_shape[:-2]

    def write(
        self,
        position_index: int,
        index: tuple[int, ...],
        frame: np.ndarray,
        *,
        frame_metadata: dict[str, Any] | None = None,
    ) -> None:
        """Write frame to specified location with auto-resize for unlimited dims."""
        if self._finalized:  # pragma: no cover
            raise RuntimeError("Cannot write after finalize().")

        # Resize array if needed to accommodate index
        self._ensure_array_size(position_index, [index])

        # Route to buffered or direct write
        array = self._arrays[position_index]
        if self._chunk_buffers:
            self._write_with_buffering(
                self._chunk_buffers[position_index], array, index, frame
            )
        else:
            self._write(array, index, frame)

        self._store_frame_metadata(position_index, index, frame_metadata)

    def advance(self, indices: Sequence[tuple[int, tuple[int, ...]]]) -> None:
        """Efficiently resize arrays once based on maximum extent of skipped indices.

        For Zarr backends, skipping frames only requires resizing arrays to
        accommodate the new indices. No data is written - skipped regions use
        the array's fill_value.
        """
        if self._finalized:  # pragma: no cover
            raise RuntimeError("Cannot advance after finalize().")

        # Group indices by position
        position_indices: dict[int, list[tuple[int, ...]]] = {}
        for pos_idx, storage_idx in indices:
            position_indices.setdefault(pos_idx, []).append(storage_idx)

        # For each position, resize once to accommodate all indices
        for pos_idx, storage_indices in position_indices.items():
            self._ensure_array_size(pos_idx, storage_indices)

    def _store_frame_metadata(
        self,
        position_index: int,
        index: tuple[int, ...],
        frame_metadata: dict[str, Any] | None,
    ) -> None:
        # Accumulate frame metadata with storage index
        if frame_metadata is not None:
            mirror = self._meta_mirrors[self._image_group_paths[position_index]]
            mirror.append_frame_metadata({**frame_metadata, "storage_index": index})

    def _write_with_buffering(
        self,
        buffer: ChunkBuffer,
        array: _AT,
        index: tuple[int, ...],
        frame: np.ndarray,
    ) -> None:
        """Write frame using chunk buffering.

        If desired, subclasses can override this to customize buffering behavior.
        """
        chunk_coords = buffer.add_frame(index, frame)

        # If chunk is complete, flush it
        if chunk_coords is not None:
            storage_start, chunk_data = buffer.get_chunk_for_flush(chunk_coords)
            self._write_chunk(array, storage_start, chunk_data)

    def _write_chunk(
        self, array: _AT, start_index: tuple[int, ...], chunk_data: np.ndarray
    ) -> None:
        """Write a complete chunk to the array.

        This is called when the chunk buffer indicates a chunk is complete and ready
        to be flushed to storage.
        """
        slices = tuple(
            slice(start, start + size)
            for start, size in zip(start_index, chunk_data.shape, strict=False)
        )
        self._write(array, slices, chunk_data)

    def _write(self, array: _AT, index: tuple[int, ...], frame: np.ndarray) -> None:
        """Write frame to array at specified index."""
        array[index] = frame

    def get_metadata(self) -> dict[str, dict]:
        """Get metadata from all array groups in the zarr hierarchy.

        Returns a dict mapping group paths (relative to root) to the *attributes*
        section of the corresponding zarr.json files (not the entire file).

        Use `update_metadata()` to modify and write back.

        Note: frame_metadata (appended during writes) is not included in the returned
        value here, but is written to disk when metadata is flushed.

        Returns
        -------
        dict[str, dict]
            Mapping of group paths to .zattrs dicts, or empty dict if not prepared.

        Examples
        --------
        For single position:
            {".": {"ome": {...}}}  # root group attrs

        For multi-position:
            {"pos0": {"ome": {...}}, "pos1": {"ome": {...}}}

        For plates:
            {"A/1/field_0": {"ome": {...}}, "A/2/field_0": {"ome": {...}}}
        """
        if not self._meta_mirrors:  # pragma: no cover
            return {}

        return {
            path: deepcopy(mirror.get("attributes", {}))
            for path, mirror in self._meta_mirrors.items()
        }

    def update_metadata(self, metadata: dict[str, dict]) -> None:
        """Update metadata for all array groups in the zarr hierarchy.

        Parameters
        ----------
        metadata : dict[str, dict]
            Mapping of group paths to .zattrs dicts. Keys should match those
            returned by get_metadata(). Values are dicts that will be written
            to the corresponding .zattrs files.

        Raises
        ------
        KeyError
            If a path in metadata doesn't correspond to a group.
        RuntimeError
            If backend is not prepared.
        """
        if not self._meta_mirrors:  # pragma: no cover
            raise RuntimeError("Backend not prepared. Call prepare() first.")

        if self._root is None:  # pragma: no cover
            warnings.warn("Root path is unknown. Cannot update metadata.", stacklevel=2)
            return

        for path, attrs in metadata.items():
            if path not in self._meta_mirrors:
                raise KeyError(f"Unknown path: {path}")

            mirror = self._meta_mirrors[path]
            mirror["attributes"] = deepcopy(attrs)
            mirror.flush()

    def finalize(self) -> None:
        """Flush and release resources."""
        if not self._finalized:
            self._finalize_chunk_buffers()
            for mirror in self._meta_mirrors.values():
                mirror.flush()
            self._arrays.clear()
            self._image_group_paths.clear()
            self._finalized = True

    def _finalize_chunk_buffers(self) -> None:
        # Flush any remaining partial chunks.
        # NOTE:
        # This is defensive code that should never execute in normal operation.
        # With sequential frame writing via OMEStream.append(), chunks are either:
        # (1) fully complete and auto-flushed immediately via _write_with_buffering
        # (2) never started (no frames written yet). This code only executes if
        # frames are written out-of-order (e.g., direct backend.write() calls) or if
        # there's a bug in chunk completion logic. It ensures no buffered data is
        # silently lost during finalization.
        if self._chunk_buffers:  # pragma: no cover
            for pos_idx, buffer in enumerate(self._chunk_buffers):
                if self._arrays:  # Arrays may have been cleared already
                    array = self._arrays[pos_idx]
                    for storage_start, chunk_data in buffer.flush_all_partial():
                        self._write_chunk(array, storage_start, chunk_data)
        self._chunk_buffers = None

    def _resize(self, array: _AT, new_shape: Sequence[int]) -> None:
        """Resize array to new shape."""
        array.resize(new_shape)

    def _cache_metadata_from_arrays(self, root: Path) -> None:
        """Cache metadata from parent groups.

        Reads zarr.json and caches their attributes in position order.
        """
        for group_path in self._image_group_paths:
            if (zarr_json := root / group_path / "zarr.json").exists():
                self._meta_mirrors[group_path] = doc = JsonDocumentMirror(zarr_json)
                doc.load()

    def _should_use_chunk_buffering(self, storage_dims: list[Dimension]) -> bool:
        """Check if chunk buffering would be beneficial.

        Returns True if any index dimension (excluding last 2 frame dims)
        has chunk_size > 1.
        """
        # Only check index dimensions (exclude frame dimensions)
        for dim in storage_dims[:-2]:
            if dim.chunk_size is not None and dim.chunk_size > 1:
                return True
        return False


# -----------------------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------------------


def _get_chunks_and_shards(
    dims: list[Dimension],
) -> tuple[tuple[int, ...], tuple[int, ...] | None]:
    """Compute chunk and shard sizes from dimensions."""
    chunks = []
    shards = []
    has_shards = False
    n = len(dims)

    for i, dim in enumerate(dims):
        # Chunks
        if dim.chunk_size is not None:
            chunks.append(dim.chunk_size)
        elif i >= n - 2:  # Last 2 dims (spatial)
            chunks.append(dim.count)
        else:
            chunks.append(1)

        # Shards
        if dim.shard_size_chunks is not None:
            # multiply shard_size_chunks by chunk size
            # to get shard shape in pixels (which is what yaozarrs expects)
            shards.append(chunks[-1] * dim.shard_size_chunks)
            has_shards = True
        else:
            shards.append(chunks[-1])  # Default to chunk size

    return tuple(chunks), tuple(shards) if has_shards else None


def _build_yaozarrs_image_model(dims: list[Dimension]) -> v05.Image:
    """Build yaozarrs v05 Image metadata from Dimensions."""
    dim_specs = [
        DimSpec(
            name=dim.name,
            size=dim.count,
            type=dim.type,
            unit=dim.unit,
            scale=1.0 if dim.scale is None else dim.scale,
            translation=dim.translation,
        )
        for dim in dims
    ]
    channel_dim = next((d for d in dims if is_channel_dim(d)), None)
    if channel_dim is not None and channel_dim.coords:
        omero_channels = []
        for c in channel_dim.coords:
            color = (
                c.color.as_hex(format="long").lstrip("#").upper() if c.color else None
            )
            omero_channels.append(v05.OmeroChannel(label=c.name, color=color))

        omero = v05.Omero(channels=omero_channels)
    else:
        omero = None
    return v05.Image(multiscales=[v05.Multiscale.from_dims(dim_specs)], omero=omero)


def _get_series_name(pos: Position) -> str:
    """Get unique series name for a position.

    For positions with grid coordinates, combines name with grid coordinates
    to ensure uniqueness. Otherwise, returns the position name as-is.
    """
    if pos.grid_row is not None and pos.grid_column is not None:
        return f"{pos.name}_{pos.grid_row}_{pos.grid_column}"
    return pos.name


def _build_yaozarrs_plate_model(
    plate: Plate, well_positions: dict[tuple[str, str], list[tuple[int, Position]]]
) -> v05.Plate:
    """Build yaozarrs v05 Plate metadata from ome-writers Plate schema."""
    return v05.Plate(
        plate=v05.PlateDef(
            name=plate.name,
            rows=[v05.Row(name=name) for name in plate.row_names],
            columns=[v05.Column(name=name) for name in plate.column_names],
            wells=[
                v05.PlateWell(
                    path=f"{row}/{col}",
                    rowIndex=plate.row_names.index(row),
                    columnIndex=plate.column_names.index(col),
                )
                for row, col in well_positions
            ],
        )
    )


class JsonDocumentMirror(MutableMapping[str, Any]):
    """In-memory mirror of a .json document with dirty tracking."""

    __slots__ = ("_data", "_dirty", "_frame_metadata", "_lock", "_path")

    def __init__(self, path: Path | str) -> None:
        self._path = Path(path)
        self._data: dict[str, Any] = {}
        self._dirty = False
        self._lock = threading.Lock()  # for thread safety

        # store this outside of _data for easy access and data integrity
        self._frame_metadata: list[dict[str, Any]] = []

    def append_frame_metadata(self, metadata: dict[str, Any]) -> None:
        """Append frame metadata entry."""
        self._frame_metadata.append(metadata)
        self._dirty = True

    def __getitem__(self, key: str) -> Any:
        with self._lock:
            return self._data[key]

    def __setitem__(self, key: str, value: Any) -> None:
        with self._lock:
            self._data[key] = value
            self._dirty = True

    def __delitem__(self, key: str) -> None:  # pragma: no cover
        with self._lock:
            del self._data[key]
            self._dirty = True

    def __iter__(self) -> Iterator[str]:  # pragma: no cover
        return iter(self._data)

    def __len__(self) -> int:  # pragma: no cover
        return len(self._data)

    def flush(self, indent: int | None = None) -> None:
        """Write data to disk if dirty."""
        with self._lock:
            if self._dirty:
                # Ensure frame_metadata is up-to-date in the document
                self.ome_writers["frame_metadata"] = self._frame_metadata

                # NOTE:
                # we could consider attempting to serialize certain sub-sections here
                # (e.g., "ome_writers.frame_metadata", which contains user-data)
                # and pop them from the metadata if they fail to serialize...
                # for now it's all-or-nothing
                try:
                    json_str = json.dumps(self._data, indent=indent)
                except (TypeError, ValueError) as e:  # pragma: no cover
                    raise ValueError(
                        f"Failed to serialize zarr.json for {self._path}:\n{e}"
                    ) from e

                self._path.write_text(json_str)
                self._dirty = False

    def load(self) -> None:
        """Load data from disk."""
        with self._lock:
            if self._dirty:  #  pragma: no cover
                raise RuntimeError("Cannot load while dirty. Flush first.")

            if self._path.exists():
                # NB: could use streaming parser for large files, but zarr.json
                # are typically small enough to read fully into memory
                content = self._path.read_text()
                self._data = cast("dict[str, Any]", json.loads(content))
            else:
                self._data = {}  # pragma: no cover
            self._dirty = False

    @property
    def ome_writers(self) -> dict[str, Any]:
        """Get ome_writers dict from attributes."""
        attrs = cast("dict", self._data.setdefault("attributes", {}))
        return attrs.setdefault("ome_writers", {"version": __version__})
