from typing import Literal
import xarray as xr
from xarray import DataTree
import xproj  # noqa ignore
from .metadata import create_level_encoding, create_multiscale_metadata
from .pyramid import Pyramid
from .chunking import (
    DEFAULT_CHUNK_BYTES,
    DEFAULT_CHUNKS_PER_SHARD,
    ChunksPerShard,
    validate_chunks_per_shard,
)

CoarseningMethod = Literal["mean", "max", "min", "sum"]


def get_crs(ds: xr.Dataset) -> str:
    crs = ds.proj.crs
    if not crs:
        raise ValueError(
            "dataset is missing a crs. Use xproj. ex: ds.proj.assign_crs({'EPSG':4326)."
        )
    return str(crs)


def build_coarsened_levels(
    ds: xr.Dataset,
    num_levels: int,
    x_dim: str,
    y_dim: str,
    method: CoarseningMethod,
) -> dict[int, xr.Dataset]:
    levels = [ds]
    for lvl in range(num_levels - 1):
        curr = levels[0]
        if curr.sizes[x_dim] < 2 or curr.sizes[y_dim] < 2:
            raise ValueError(f"cannot coarsen {num_levels}")
        coarsened = curr.coarsen({x_dim: 2, y_dim: 2}, boundary="trim")
        levels.insert(0, getattr(coarsened, method)())

    # zarr-multiscales: lowest levels = highest resolution (level 0 = highest res)
    return dict(enumerate(reversed(levels)))


def create_pyramid(
    ds: xr.Dataset,
    levels: int,
    x_dim: str = "x",
    y_dim: str = "y",
    method: CoarseningMethod = "mean",
    target_chunk_bytes: int = DEFAULT_CHUNK_BYTES,
    chunks_per_shard: ChunksPerShard | None = DEFAULT_CHUNKS_PER_SHARD,
) -> Pyramid:
    if chunks_per_shard is not None:
        validate_chunks_per_shard(chunks_per_shard)
    crs_str = get_crs(ds)
    level_datasets = build_coarsened_levels(ds, levels, x_dim, y_dim, method)

    dt = DataTree(name="root")
    full_encoding = {}

    for idx, ds_level in level_datasets.items():
        name = str(idx)
        path = f"/{idx}"

        level_encoding = create_level_encoding(
            ds_level,
            x_dim,
            y_dim,
            target_chunk_bytes=target_chunk_bytes,
            chunks_per_shard=chunks_per_shard,
        )

        dim_chunks = {}
        for var_name, var_enc in level_encoding.items():
            if var_name in ds_level.data_vars and "chunks" in var_enc:
                dask_chunks = var_enc.get("shards", var_enc["chunks"])
                da = ds_level[var_name]

                for dim, chunk_size in zip(da.dims, dask_chunks):
                    if dim not in dim_chunks:
                        dim_chunks[dim] = chunk_size
                    else:
                        dim_chunks[dim] = min(dim_chunks[dim], chunk_size)

        if dim_chunks:
            ds_level = ds_level.chunk(dim_chunks)

        dt[path] = DataTree(ds_level, name=name)
        full_encoding[path] = level_encoding

    dt.attrs = create_multiscale_metadata(
        ds=ds,
        x_dim=x_dim,
        y_dim=y_dim,
        level_datasets=level_datasets,
        crs=crs_str,
        method=str(method),
    )

    return Pyramid(datatree=dt, encoding=full_encoding)
