"""
Functions to import and export :py:class:`xradio.schema.metamodel.DatasetSchema`
as JSON representation. This can be used to externalise schema checks, or
generate documentation from schemas in JSON representation.
"""

import dataclasses
import json

from xradio.schema import (
    bases,
    metamodel,
    xarray_dataclass_to_array_schema,
    xarray_dataclass_to_dataset_schema,
    xarray_dataclass_to_dict_schema,
)

__all__ = ["export_schema_json_file", "import_schema_json_file"]

CLASS_ATTR = "$class"


class DataclassEncoder(json.JSONEncoder):
    """
    General-purpose encoder that represents data classes as
    dictionaries, omitting defaults and annotating the original class
    as a ``'$class'`` attribute.
    """

    def default(self, o):
        if dataclasses.is_dataclass(o):
            res = {CLASS_ATTR: o.__class__.__name__}
            for fld in dataclasses.fields(type(o)):
                if (
                    getattr(o, fld.name) is not fld.default
                    and getattr(o, fld.name) is not dataclasses.MISSING
                ):
                    res[fld.name] = getattr(o, fld.name)
            return res
        return super().default(o)


DATACLASS_MAP = {
    cls.__name__: cls
    for cls in [
        metamodel.DictSchema,
        metamodel.ValueSchema,
        metamodel.AttrSchemaRef,
        metamodel.ArraySchema,
        metamodel.ArraySchemaRef,
        metamodel.DatasetSchema,
    ]
}


class DataclassDecoder(json.JSONDecoder):
    """
    General-purpose decoder that reads JSON as generated by
    :py:class:`DataclassEncoder`.
    """

    def __init__(self, dataclass_map, *args, **kwargs):
        self._dataclass_map = dataclass_map
        super().__init__(*args, object_hook=self.object_hook, **kwargs)

    def object_hook(self, obj):

        # Detect dictionaries with '$class' annotation
        if isinstance(obj, dict) and CLASS_ATTR in obj:

            # Identify the class
            cls_name = obj[CLASS_ATTR]
            cls = self._dataclass_map.get(cls_name)
            if not cls:
                raise ValueError(
                    f"Unknown $dataclass encountered while decoding JSON: {cls_name}"
                )

            # Instantiate
            del obj[CLASS_ATTR]
            obj = cls(**obj)

        return obj


def export_schema_json_file(schema: "DatasetSchema", fname: str):
    """
    Exports given schema as a JSON file

    :param schema: Dataset schema. Dataclasses will be converted automatically.
    :param fname: File name to write serialised schema to
    """

    # Check that this is actually a Dataset
    if bases.is_dataset_schema(schema):
        schema = xarray_dataclass_to_dataset_schema(schema)
    if not isinstance(schema, metamodel.DatasetSchema):
        raise TypeError(
            f"export_schema_json_file: Expected DatasetSchema, but got {type(schema)}!"
        )

    # Perform export
    with open(fname, "w", encoding="utf8") as f:
        json.dump(schema, f, cls=DataclassEncoder, ensure_ascii=False, indent="  ")


def import_schema_json_file(fname: str):
    """
    Imports a schema from a JSON file

    For JSON files generated by
    :py:func:`export_schema_json_file`, this will return a
    :py:class:`~xradio.schema.metamodel.DatasetSchema`.

    :param fname: File name to load
    :returns: Deserialised object
    """

    with open(fname, "r", encoding="utf8") as f:
        return json.load(f, cls=DataclassDecoder, dataclass_map=DATACLASS_MAP)
