"""
Choosing which slice of a variable to plot.

Two conventions meet here and must not be confused: the interactive prompts are
1-based, because that is what a scientist reading "size 33" expects, while
--extra-indices is 0-based to match every other tool. `get_dimension_indices` is
the single place the conversion happens.
"""

import json

import numpy as np

from .io import materialize, maybe_chunk

# The whole word is accepted alongside the letter, as it is for the yes/no
# prompts: a one-letter answer is a shortcut, not the only way in.
AVERAGE_WORDS = frozenset({'a', 'avg', 'average', 'mean'})
EVERY_WORDS = frozenset({'', 'e', 'all', 'every'})


def get_dimension_indices(src, varname):
    """
    For each dimension of the variable:
     - if size == 1 → automatically select index 0
     - otherwise prompt the user:
         <number>     : take that specific index (1-based)
         'a'          : average over this dimension
         'e' or Enter : take all values
    Returns {dim_name: int index, 'avg', or None}.
    """
    dims, shape = src.dims_and_shape(varname)
    selection = {}
    for dim, size in zip(dims, shape):
        if size == 1:
            selection[dim] = 0
            continue
        prompt = (
                f"Available options for '{dim}' (size {size}):\n"
                f"  > '1–{size}' to pick that index\n"
                "  > 'a' to average over this dimension\n"
                "  > 'e' or Enter to take all values\n"
                "Choose: "
        )
        while True:
            try:
                resp = input(prompt).strip().lower()
            except EOFError:
                print()
                selection[dim] = None
                break
            if resp in EVERY_WORDS:
                selection[dim] = None
                break
            if resp in AVERAGE_WORDS:
                selection[dim] = 'avg'
                break
            if resp.isdigit():
                n = int(resp)
                if 1 <= n <= size:
                    selection[dim] = n - 1
                    break
            print(f"  Invalid entry '{resp}'. Please enter a number, 'a', 'e', or just Enter.")
    return selection


def autoselect_singletons(extra, dims, shape):
    """
    Pin every dimension of size 1 to index 0, leaving explicit choices alone.
    Mutates and returns `extra`.
    """
    for name, size in zip(dims, shape):
        if size == 1 and name not in extra:
            extra[name] = 0
    return extra


def select_lazy(da, extra_indices):
    """
    The slice the user asked for, still lazy: nothing has been read yet.

    An integer pins that dimension and drops it; 'avg' averages over it; anything
    else keeps it whole. Indexing stays lazy, so only the requested hyperslab
    ever leaves the disk, and large variables are averaged through dask instead
    of being read whole.

    Kept apart from `apply_selection` so that a reduction can be applied on top
    of it before anything is materialized - which is what lets `--reduce` stream
    the way `-e '{"dim": "avg"}'` always has.
    """
    isel = {}
    avg_dims = []
    for dim in da.dims:
        sel = extra_indices.get(dim)
        if isinstance(sel, (int, np.integer)) and not isinstance(sel, bool):
            # An integer index drops the dimension on read
            isel[dim] = int(sel)
        elif sel == 'avg':
            avg_dims.append(dim)

    out = da.isel(isel) if isel else da
    if avg_dims:
        # Chunk before anything else: casting first would pull the whole
        # variable into memory and defeat the streaming entirely. Chunking is
        # what keeps a 669-step average over a 51 MB variable bounded.
        out = maybe_chunk(out, over=avg_dims)
        # Accumulate in float64 so a long average over float32 storage does not
        # lose precision. On a dask-backed array this stays lazy.
        out = out.astype('float64').mean(dim=avg_dims, skipna=True, keep_attrs=True)

    return out


def apply_selection(da, extra_indices):
    """
    The same slice, read into memory.

    Returns (data, remaining_dims) where data is a float64 array and
    remaining_dims lists, in order, the dimensions still present.
    """
    out = select_lazy(da, extra_indices)
    return materialize(out), list(out.dims)


def parse_extra_indices(extra_json, dims, shape):
    """
    Parse the --extra-indices JSON object into {dim: 0-based index or 'avg'}.
    Raises ValueError with an explicit message on anything unusable, rather than
    silently plotting a different slice than the user asked for.
    """
    if not extra_json:
        return {}
    try:
        parsed = json.loads(extra_json)
    except json.JSONDecodeError as err:
        raise ValueError(f"--extra-indices is not valid JSON: {err}") from err
    if not isinstance(parsed, dict):
        raise ValueError('--extra-indices must be a JSON object, e.g. \'{"Time": 0}\'')

    sizes = dict(zip(dims, shape))
    extra = {}
    for key, value in parsed.items():
        if key not in sizes:
            raise ValueError(f"--extra-indices: '{key}' is not a dimension of this "
                             f"variable (available: {list(dims)})")
        if isinstance(value, str):
            if value.lower() != 'avg':
                raise ValueError(f"--extra-indices: value for '{key}' must be an "
                                 f'integer or "avg", got {value!r}')
            extra[key] = 'avg'
        elif isinstance(value, bool) or not isinstance(value, int):
            raise ValueError(f"--extra-indices: value for '{key}' must be a 0-based "
                             f'integer or "avg", got {value!r}')
        elif not 0 <= value < sizes[key]:
            raise ValueError(f"--extra-indices: index {value} is out of range for "
                             f"'{key}' (size {sizes[key]}, 0-based)")
        else:
            extra[key] = value
    return extra