from typing import Any, Callable, Dict, Hashable, Iterable, Optional, Union

import xarray as xr

from .base import BaseSchema, SchemaError
from .components import AttrSchema, AttrsSchema
from .dataarray import CoordsSchema, DataArraySchema


class DatasetSchema(BaseSchema):
    '''A light-weight xarray.Dataset validator

    Parameters
    ----------
    data_vars : mapping of variable names and DataArraySchemas, optional
        Per-variable DataArraySchema's, by default None
    coords: Union[CoordsSchema, Dict[Hashable, DataArraySchema]], optional
        Coordinates schema to validate against the Dataset, optional
    attrs: Union[AttrsSchema, Dict[Hashable, AttrsSchema]], optional
        Attribute names and their values, optional.
    checks : Iterable[Callable], optional
        Dataset wide checks, by default None
    '''

    _json_schema = {
        'type': 'object',
        'properties': {
            'data_vars': {'type': 'object'},
            'coords': {'type': 'object'},
            'attrs': {'type': 'object'},
        },
        'required': [],
        'additionalProperties': False,
    }

    def __init__(
        self,
        data_vars: Optional[Dict[Hashable, Optional[DataArraySchema]]] = None,
        coords: Union[CoordsSchema, Dict[Hashable, DataArraySchema], None] = None,
        attrs: Union[AttrsSchema, Dict[Hashable, AttrSchema], None] = None,
        checks: Optional[Iterable[Callable]] = None,
    ) -> None:
        self.data_vars = data_vars  # type: ignore
        self.coords = coords  # type: ignore
        self.attrs = attrs  # type: ignore
        self.checks = checks

    @classmethod
    def from_json(cls, obj: dict):
        kwargs = {}
        if 'data_vars' in obj:
            kwargs['data_vars'] = {
                k: DataArraySchema.from_json(v) for k, v in obj['data_vars'].items()
            }
        if 'coords' in obj:
            kwargs['coords'] = CoordsSchema.from_json(obj['coords'])
        if 'attrs' in obj:
            kwargs['attrs'] = AttrsSchema.from_json(obj['attrs'])

        return cls(**kwargs)

    def validate(self, ds: xr.Dataset) -> None:
        '''Check if the Dataset complies with the Schema.

        Parameters
        ----------
        ds : xr.Dataset
            Dataset to be validated

        Returns
        -------
        xr.Dataset
            Validated Dataset

        Raises
        ------
        SchemaError
        '''

        if self.data_vars is not None:
            for key, da_schema in self.data_vars.items():
                if da_schema is not None:
                    if key not in ds.data_vars:
                        raise SchemaError(f'data variable {key} not in ds')
                    else:
                        da_schema.validate(ds.data_vars[key])

        if self.coords is not None:
            self.coords.validate(ds.coords)

        if self.attrs is not None:
            self.attrs.validate(ds.attrs)

        if self.checks:
            for check in self.checks:
                check(ds)

    @property
    def attrs(self) -> Union[AttrsSchema, None]:
        return self._attrs

    @attrs.setter
    def attrs(self, value: Union[AttrsSchema, Dict[Hashable, Any], None]):
        if value is None or isinstance(value, AttrsSchema):
            self._attrs = value
        else:
            self._attrs = AttrsSchema(value)

    @property
    def data_vars(self) -> Optional[Dict[Hashable, Optional[DataArraySchema]]]:
        return self._data_vars  # type: ignore

    @data_vars.setter
    def data_vars(self, value: Optional[Dict[Hashable, Optional[DataArraySchema]]]):
        if isinstance(value, dict):
            self._data_vars = {}
            for k, v in value.items():
                if isinstance(v, DataArraySchema):
                    self._data_vars[k] = v
                else:
                    self._data_vars[k] = DataArraySchema(**v)  # type: ignore
        elif value is None:
            self._data_vars = None  # type: ignore
        else:
            raise ValueError('must set data_vars with a dict')
        self.check_dims_consistency()

    def check_dims_consistency(self):
        if not hasattr(self, '_coords'):
            return
        das = []
        if self.data_vars is not None:
            das.extend(self.data_vars.items())
        if self.coords is not None:
            das.extend(self.coords.coords.items())
        dims_shapes = {}
        for name, da in das:
            if da is None:
                continue
            if da.dims is None or da.shape is None:
                continue
            if da.dims.dims is None or da.shape.shape is None:
                continue
            for dim, shape in zip(da.dims.dims, da.shape.shape):
                if dim is None:
                    continue
                if dim not in dims_shapes:
                    dims_shapes[dim] = shape
                elif dims_shapes[dim] != shape:
                    raise ValueError(
                        f"Inconsistent shape for dimension {dim!r} in {name!r}: "
                        f"expected {dims_shapes[dim]}, got {shape}"
                    )

    @property
    def coords(self) -> Optional[CoordsSchema]:
        return self._coords  # type: ignore

    @coords.setter
    def coords(self, value: Optional[Union[CoordsSchema, Dict[Hashable, DataArraySchema]]]):
        if value is None or isinstance(value, CoordsSchema):
            self._coords = value
        else:
            self._coords = CoordsSchema(value)
        self.check_dims_consistency()

    @property
    def json(self):
        obj = {}
        if self.data_vars:
            obj['data_vars'] = {}
            for key, var in self.data_vars.items():
                obj['data_vars'][key] = var.json
        if self.coords:
            obj['coords'] = self.coords.json
        if self.attrs is not None:
            obj['attrs'] = self.attrs.json
        return obj
