"""
Module for exporting OME-Arrow data to other formats.
"""

from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq

from ome_arrow.meta import OME_ARROW_STRUCT, OME_ARROW_TAG_TYPE, OME_ARROW_TAG_VERSION


def to_numpy(
    data: Dict[str, Any] | pa.StructScalar,
    dtype: np.dtype = np.uint16,
    strict: bool = True,
    clamp: bool = False,
) -> np.ndarray:
    """
    Convert an OME-Arrow record into a NumPy array shaped (T,C,Z,Y,X).

    The OME-Arrow "planes" are flattened YX slices indexed by (z, t, c).
    This function reconstitutes them into a dense TCZYX ndarray.

    Args:
        data:
            OME-Arrow data as a Python dict or a `pa.StructScalar`.
        dtype:
            Output dtype (default: np.uint16). If different from plane
            values, a cast (and optional clamp) is applied.
        strict:
            When True, raise if a plane has wrong pixel length. When
            False, truncate/pad that plane to the expected length.
        clamp:
            If True, clamp values to the valid range of the target
            dtype before casting.

    Returns:
        np.ndarray: Dense array with shape (T, C, Z, Y, X).

    Raises:
        KeyError: If required OME-Arrow fields are missing.
        ValueError: If dimensions are invalid or planes are malformed.

    Examples:
        >>> arr = ome_arrow_to_tczyx(my_row)  # (T, C, Z, Y, X)
        >>> arr.shape
        (1, 2, 1, 512, 512)
    """
    # Unwrap Arrow scalar to plain Python dict if needed.
    if isinstance(data, pa.StructScalar):
        data = data.as_py()

    pm = data["pixels_meta"]
    sx, sy = int(pm["size_x"]), int(pm["size_y"])
    sz, sc, st = int(pm["size_z"]), int(pm["size_c"]), int(pm["size_t"])
    if sx <= 0 or sy <= 0 or sz <= 0 or sc <= 0 or st <= 0:
        raise ValueError("All size_* fields must be positive integers.")

    expected_len = sx * sy

    # Prepare target array (T,C,Z,Y,X), zero-filled by default.
    out = np.zeros((st, sc, sz, sy, sx), dtype=dtype)

    # Helper: cast (with optional clamp) to the output dtype.
    if np.issubdtype(dtype, np.integer):
        info = np.iinfo(dtype)
        lo, hi = info.min, info.max
    elif np.issubdtype(dtype, np.floating):
        lo, hi = -np.inf, np.inf
    else:
        # Rare dtypes: no clamping logic; rely on astype.
        lo, hi = -np.inf, np.inf

    def _cast_plane(a: np.ndarray) -> np.ndarray:
        if clamp:
            a = np.clip(a, lo, hi)
        return a.astype(dtype, copy=False)

    # Fill planes.
    for i, p in enumerate(data.get("planes", [])):
        z = int(p["z"])
        t = int(p["t"])
        c = int(p["c"])

        if not (0 <= z < sz and 0 <= t < st and 0 <= c < sc):
            raise ValueError(f"planes[{i}] index out of range: (z,t,c)=({z},{t},{c})")

        pix = p["pixels"]
        # Ensure sequence-like and correct length.
        try:
            n = len(pix)
        except Exception as e:
            raise ValueError(f"planes[{i}].pixels is not a sequence") from e

        if n != expected_len:
            if strict:
                raise ValueError(
                    f"planes[{i}].pixels length {n} != size_x*size_y {expected_len}"
                )
            # Lenient mode: fix length by truncation or zero-pad.
            if n > expected_len:
                pix = pix[:expected_len]
            else:
                pix = list(pix) + [0] * (expected_len - n)

        # Reshape to (Y,X) and cast.
        arr2d = np.asarray(pix).reshape(sy, sx)
        arr2d = _cast_plane(arr2d)
        out[t, c, z] = arr2d

    return out


def to_ome_tiff(
    data: Dict[str, Any] | pa.StructScalar,
    out_path: str,
    *,
    dtype: np.dtype = np.uint16,
    clamp: bool = False,
    dim_order: str = "TCZYX",
    compression: Optional[str] = "zlib",  # "zlib","lzma","jpegxl", or None
    compression_level: int = 6,
    tile: Optional[Tuple[int, int]] = None,  # (Y, X)
    use_channel_colors: bool = False,
) -> None:
    """
    Export an OME-Arrow record to OME-TIFF using BioIO's OmeTiffWriter.

    Notes
    -----
    - No 'bigtiff' kwarg is passed (invalid for tifffile.TiffWriter.write()).
      BigTIFF selection is automatic based on file size.
    """
    from ome_arrow.export import to_numpy  # your existing function

    try:
        from bioio.writers import OmeTiffWriter
    except Exception:
        from bioio_ome_tiff.writers import OmeTiffWriter  # type: ignore

    # PhysicalPixelSizes (robust import or shim)
    try:
        from bioio import PhysicalPixelSizes  # modern bioio
    except Exception:
        try:
            from bioio.types import PhysicalPixelSizes
        except Exception:
            try:
                from aicsimageio.types import PhysicalPixelSizes
            except Exception:
                from typing import NamedTuple
                from typing import Optional as _Opt

                class PhysicalPixelSizes(NamedTuple):  # type: ignore
                    Z: _Opt[float] = None
                    Y: _Opt[float] = None
                    X: _Opt[float] = None

    # 1) Dense array (T,C,Z,Y,X)
    arr = to_numpy(data, dtype=dtype, clamp=clamp)

    # 2) Metadata
    row = data.as_py() if isinstance(data, pa.StructScalar) else data
    pm = row["pixels_meta"]
    _st, sc, _sz, _sy, _sx = arr.shape

    # Channel names
    chs: Sequence[Dict[str, Any]] = pm.get("channels", []) or []
    channel_names = [f"C{i}" for i in range(sc)]
    if len(chs) == sc:
        for i, ch in enumerate(chs):
            nm = ch.get("name")
            if nm is not None:
                channel_names[i] = str(nm)

    # Optional channel colors (guarded)
    channel_colors_for_writer = None
    if use_channel_colors and len(chs) == sc:

        def _rgba_to_rgb(rgba: int) -> int:
            r = (rgba >> 24) & 0xFF
            g = (rgba >> 16) & 0xFF
            b = (rgba >> 8) & 0xFF
            return (r << 16) | (g << 8) | b

        flat_colors: list[int] = []
        for ch in chs:
            rgba = ch.get("color_rgba")
            flat_colors.append(
                _rgba_to_rgb(int(rgba)) if isinstance(rgba, int) else 0xFFFFFF
            )
        if len(flat_colors) == sc:
            channel_colors_for_writer = [flat_colors]  # list-per-image

    # Physical sizes (µm) in Z, Y, X order for BioIO
    p_dx = float(pm.get("physical_size_x", 1.0) or 1.0)
    p_dy = float(pm.get("physical_size_y", 1.0) or 1.0)
    p_dz = float(pm.get("physical_size_z", 1.0) or 1.0)
    pps_list = [PhysicalPixelSizes(Z=p_dz, Y=p_dy, X=p_dx)]

    # tifffile passthrough (NO 'bigtiff' here)
    tifffile_kwargs: Dict[str, Any] = {}
    if compression is not None:
        tifffile_kwargs["compression"] = compression
        if compression == "zlib":
            tifffile_kwargs["compressionargs"] = {"level": int(compression_level)}
    if tile is not None:
        tifffile_kwargs["tile"] = (int(tile[0]), int(tile[1]))

    # list-per-image payloads
    data_list = [arr]
    dim_order_list = [dim_order]
    image_name_list = [str(row.get("name") or row.get("id") or "image")]
    ch_names_list = [channel_names]

    # 3) Write
    OmeTiffWriter.save(
        data_list,
        out_path,
        dim_order=dim_order_list,
        image_name=image_name_list,
        channel_names=ch_names_list,
        channel_colors=channel_colors_for_writer,  # None or [flat list len=sc]
        physical_pixel_sizes=pps_list,
        tifffile_kwargs=tifffile_kwargs,
    )


def to_ome_zarr(
    data: Dict[str, Any] | pa.StructScalar,
    out_path: str,
    *,
    dtype: np.dtype = np.uint16,
    clamp: bool = False,
    # Axes order for the on-disk array — must match arr shape (T,C,Z,Y,X)
    dim_order: str = "TCZYX",
    # NGFF / multiscale
    multiscale_levels: int = 1,  # 1 = no pyramid; >1 builds levels
    downscale_spatial_by: int = 2,  # per-level factor for Z,Y,X
    zarr_format: int = 3,  # 3 (NGFF 0.5) or 2 (NGFF 0.4)
    # Storage knobs
    chunks: Optional[Tuple[int, int, int, int, int]] = None,  # (T,C,Z,Y,X) or None
    shards: Optional[Tuple[int, int, int, int, int]] = None,  # v3 only, optional
    compressor: Optional[str] = "zstd",  # "zstd","lz4","gzip", or None
    compressor_level: int = 3,
    # Optional display metadata (carried through if you later enrich channels/rdefs)
    image_name: Optional[str] = None,
) -> None:
    """
    Write OME-Zarr using your `OMEZarrWriter` (instance API).

    - Builds arr as (T,C,Z,Y,X) using your `to_numpy`.
    - Creates level shapes for a multiscale pyramid (if multiscale_levels>1).
    - Chooses Blosc codec compatible with zarr_format (v2 vs v3).
    - Populates axes names/types/units and physical pixel sizes from pixels_meta.
    """
    # --- local import to avoid hard deps at module import time
    # Use the class you showed
    from bioio_ome_zarr.writers import OMEZarrWriter

    from ome_arrow.export import to_numpy  # your existing function

    # Optional compressors for v2 vs v3
    compressor_obj = None
    if compressor is not None:
        if zarr_format == 2:
            # numcodecs Blosc (v2 path)
            from numcodecs import Blosc as BloscV2

            cname = {"zstd": "zstd", "lz4": "lz4", "gzip": "zlib"}.get(
                compressor, "zstd"
            )
            compressor_obj = BloscV2(
                cname=cname, clevel=int(compressor_level), shuffle=BloscV2.BITSHUFFLE
            )
        else:
            # zarr v3 codec
            from zarr.codecs import BloscCodec, BloscShuffle

            cname = {"zstd": "zstd", "lz4": "lz4", "gzip": "zlib"}.get(
                compressor, "zstd"
            )
            compressor_obj = BloscCodec(
                cname=cname,
                clevel=int(compressor_level),
                shuffle=BloscShuffle.bitshuffle,
            )

    # 1) Dense pixel data (T,C,Z,Y,X)
    arr = to_numpy(data, dtype=dtype, clamp=clamp)

    # 2) Unwrap OME-Arrow metadata
    row = data.as_py() if isinstance(data, pa.StructScalar) else data
    pm = row["pixels_meta"]
    st, sc, sz, sy, sx = arr.shape

    # 3) Axis metadata (names/types/units aligned with T,C,Z,Y,X)
    axes_names = [a.lower() for a in dim_order]  # ["t","c","z","y","x"]
    axes_types = ["time", "channel", "space", "space", "space"]
    # Units: micrometers for spatial, leave T/C None
    axes_units = [
        None,
        None,
        pm.get("physical_size_z_unit") or "µm",
        pm.get("physical_size_y_unit") or "µm",
        pm.get("physical_size_x_unit") or "µm",
    ]
    # Physical pixel sizes at level 0 in axis order
    p_dx = float(pm.get("physical_size_x", 1.0) or 1.0)
    p_dy = float(pm.get("physical_size_y", 1.0) or 1.0)
    p_dz = float(pm.get("physical_size_z", 1.0) or 1.0)
    physical_pixel_size = [1.0, 1.0, p_dz, p_dy, p_dx]  # T,C,Z,Y,X

    # 4) Multiscale level shapes (level 0 first). Only spatial dims are downscaled.
    def _down(a: int, f: int) -> int:
        return max(1, a // f)

    def _level_shapes_tcxyz(levels: int) -> List[Tuple[int, int, int, int, int]]:
        shapes = [(st, sc, sz, sy, sx)]
        for _ in range(levels - 1):
            t, c, z, y, x = shapes[-1]
            shapes.append(
                (
                    t,
                    c,
                    _down(z, downscale_spatial_by),
                    _down(y, downscale_spatial_by),
                    _down(x, downscale_spatial_by),
                )
            )
        return shapes

    multiscale_levels = max(1, int(multiscale_levels))
    level_shapes: List[Tuple[int, int, int, int, int]] = _level_shapes_tcxyz(
        multiscale_levels
    )

    # 5) Chunking / shards (can be single-shape or per-level;
    # we pass single-shape if provided)
    chunk_shape: Optional[List[Tuple[int, ...]]] = None
    if chunks is not None:
        chunk_shape = [tuple(int(v) for v in chunks)] * multiscale_levels

    shard_shape: Optional[List[Tuple[int, ...]]] = None
    if shards is not None and zarr_format == 3:
        shard_shape = [tuple(int(v) for v in shards)] * multiscale_levels

    # 6) Image name default
    img_name = image_name or str(row.get("name") or row.get("id") or "Image")

    # 7) Instantiate writer with your class constructor
    writer = OMEZarrWriter(
        store=out_path,
        level_shapes=level_shapes,
        dtype=dtype,
        chunk_shape=chunk_shape,
        shard_shape=shard_shape,
        compressor=compressor_obj,
        zarr_format=3 if int(zarr_format) == 3 else 2,
        image_name=img_name,
        channels=None,  # you can map your channel metadata here later
        rdefs=None,  # optional OMERO display metadata
        creator_info=None,  # optional "creator" block
        root_transform=None,  # optional NGFF root transform
        axes_names=axes_names,
        axes_types=axes_types,
        axes_units=axes_units,
        physical_pixel_size=physical_pixel_size,
    )

    # 8) Write full-resolution; writer will build & fill lower levels
    writer.write_full_volume(arr)


def to_ome_parquet(
    data: Dict[str, Any] | pa.StructScalar,
    out_path: str,
    column_name: str = "image",
    file_metadata: Optional[Dict[str, str]] = None,
    compression: Optional[str] = "zstd",
    row_group_size: Optional[int] = None,
) -> None:
    """
    Export an OME-Arrow record to a Parquet file as a single-row, single-column table.
    The single column holds a struct with the OME-Arrow schema.
    """

    # 1) Normalize to a plain Python dict (works better with pyarrow builders,
    #    especially when the struct has a `null`-typed field like "masks").
    if isinstance(data, pa.StructScalar):
        record_dict = data.as_py()
    else:
        # Validate by round-tripping through a typed scalar, then back to dict.
        record_dict = pa.scalar(data, type=OME_ARROW_STRUCT).as_py()

    # 2) Build a single-row struct array from the dict, explicitly passing the schema
    struct_array = pa.array([record_dict], type=OME_ARROW_STRUCT)  # len=1

    # 3) Wrap into a one-column table
    table = pa.table({column_name: struct_array})

    # 4) Attach optional file-level metadata
    meta: Dict[bytes, bytes] = dict(table.schema.metadata or {})
    try:
        meta[b"ome.arrow.type"] = str(OME_ARROW_TAG_TYPE).encode("utf-8")
        meta[b"ome.arrow.version"] = str(OME_ARROW_TAG_VERSION).encode("utf-8")
    except Exception:
        pass
    if file_metadata:
        for k, v in file_metadata.items():
            meta[str(k).encode("utf-8")] = str(v).encode("utf-8")
    table = table.replace_schema_metadata(meta)

    # 5) Write Parquet (single row, single column)
    pq.write_table(
        table,
        out_path,
        compression=compression,
        row_group_size=row_group_size,
    )
