"""Code related to loading and dumping."""

import inspect
import re
from typing import Any, Callable, Dict, Tuple, Type

import yaml

from ._common import load_value_mode, parent_parser
from ._optionals import import_jsonnet, omegaconf_support
from ._type_checking import ArgumentParser

__all__ = [
    "set_loader",
    "set_dumper",
]


class DefaultLoader(getattr(yaml, "CSafeLoader", yaml.SafeLoader)):  # type: ignore[misc]
    pass


# https://stackoverflow.com/a/37958106/2732151
def remove_implicit_resolver(cls, tag_to_remove):
    if "yaml_implicit_resolvers" not in cls.__dict__:
        cls.yaml_implicit_resolvers = cls.yaml_implicit_resolvers.copy()

    for first_letter, mappings in cls.yaml_implicit_resolvers.items():
        cls.yaml_implicit_resolvers[first_letter] = [(tag, regexp) for tag, regexp in mappings if tag != tag_to_remove]


remove_implicit_resolver(DefaultLoader, "tag:yaml.org,2002:timestamp")
remove_implicit_resolver(DefaultLoader, "tag:yaml.org,2002:float")


DefaultLoader.add_implicit_resolver(
    "tag:yaml.org,2002:float",
    re.compile(
        """^(?:
     [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)?
    |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
    |\\.[0-9_]+(?:[eE][-+][0-9]+)?
    |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]*
    |[-+]?\\.(?:inf|Inf|INF)
    |\\.(?:nan|NaN|NAN))$""",
        re.X,
    ),
    list("-+0123456789."),
)


def yaml_load(stream):
    if stream.strip() == "-":
        value = stream
    else:
        value = yaml.load(stream, Loader=DefaultLoader)
    if isinstance(value, dict) and value and all(v is None for v in value.values()):
        if len(value) == 1 and stream.strip() == next(iter(value.keys())) + ":":
            value = stream
        else:
            keys = set(stream.strip(" {}").replace(" ", "").split(","))
            if len(keys) > 0 and keys == set(value.keys()):
                value = stream
    return value


def jsonnet_load(stream, path="", ext_vars=None):
    from ._jsonnet import ActionJsonnet

    ext_vars, ext_codes = ActionJsonnet.split_ext_vars(ext_vars)
    _jsonnet = import_jsonnet("jsonnet_load")
    try:
        val = _jsonnet.evaluate_snippet(path, stream, ext_vars=ext_vars, ext_codes=ext_codes)
    except RuntimeError as ex:
        try:
            return yaml_load(stream)
        except pyyaml_exceptions:
            raise ValueError(str(ex)) from ex
    return yaml_load(val)


loaders: Dict[str, Callable] = {
    "yaml": yaml_load,
    "jsonnet": jsonnet_load,
}

pyyaml_exceptions = (yaml.YAMLError,)
jsonnet_exceptions = pyyaml_exceptions + (ValueError,)

loader_exceptions: Dict[str, Tuple[Type[Exception], ...]] = {
    "yaml": pyyaml_exceptions,
    "jsonnet": jsonnet_exceptions,
}


def get_load_value_mode() -> str:
    mode = load_value_mode.get()
    if mode is None:
        parser = parent_parser.get()
        assert parser is not None
        mode = parser.parser_mode
    return mode


def get_loader_exceptions():
    return loader_exceptions[get_load_value_mode()]


def load_value(value: str, simple_types: bool = False, **kwargs):
    loader = loaders[get_load_value_mode()]
    if kwargs:
        params = set(list(inspect.signature(loader).parameters)[1:])
        kwargs = {k: v for k, v in kwargs.items() if k in params}
    loaded_value = loader(value, **kwargs)
    if not simple_types and isinstance(loaded_value, (int, float, bool, str)):
        loaded_value = value
    return loaded_value


dump_yaml_kwargs = {
    "default_flow_style": False,
    "allow_unicode": True,
    "sort_keys": False,
}

dump_json_kwargs = {
    "ensure_ascii": False,
    "sort_keys": False,
}


def yaml_dump(data):
    return yaml.safe_dump(data, **dump_yaml_kwargs)


def yaml_comments_dump(data, parser):
    dump = dumpers["yaml"](data)
    formatter = parser.formatter_class(parser.prog)
    return formatter.add_yaml_comments(dump)


def json_dump(data):
    import json

    return json.dumps(data, separators=(",", ":"), **dump_json_kwargs)


def json_indented_dump(data):
    import json

    return json.dumps(data, indent=2, **dump_json_kwargs) + "\n"


dumpers: Dict[str, Callable] = {
    "yaml": yaml_dump,
    "yaml_comments": yaml_comments_dump,
    "json": json_dump,
    "json_indented": json_indented_dump,
    "jsonnet": json_indented_dump,
}

comment_prefix: Dict[str, str] = {
    "yaml": "# ",
    "yaml_comments": "# ",
    "jsonnet": "// ",
}


def check_valid_dump_format(dump_format: str):
    if dump_format not in {"parser_mode"}.union(set(dumpers)):
        raise ValueError(f'Unknown output format "{dump_format}".')


def dump_using_format(parser: "ArgumentParser", data: dict, dump_format: str) -> str:
    if dump_format == "parser_mode":
        dump_format = parser.parser_mode if parser.parser_mode in dumpers else "yaml"
    args = (data, parser) if dump_format == "yaml_comments" else (data,)
    dump = dumpers[dump_format](*args)
    if parser.dump_header and comment_prefix.get(dump_format):
        prefix = comment_prefix[dump_format]
        header = "\n".join(prefix + line for line in parser.dump_header)
        dump = f"{header}\n{dump}"
    return dump


def set_loader(mode: str, loader_fn: Callable[[str], Any], exceptions: Tuple[Type[Exception], ...] = pyyaml_exceptions):
    """Sets the value loader function to be used when parsing with a certain mode.

    The ``loader_fn`` function must accept as input a single str type parameter
    and return any of the basic types {str, bool, int, float, list, dict, None}.
    If this function is not based on PyYAML for things to work correctly the
    exceptions types that can be raised when parsing a value fails should be
    provided.

    Args:
        mode: The parser mode for which to set its loader function. Example: "yaml".
        loader_fn: The loader function to set. Example: ``yaml.safe_load``.
        exceptions: Exceptions that the loader can raise when load fails.
            Example: (yaml.parser.ParserError, yaml.scanner.ScannerError).
    """
    loaders[mode] = loader_fn
    loader_exceptions[mode] = exceptions


def set_dumper(format_name: str, dumper_fn: Callable[[Any], str]):
    """Sets the dumping function for a given format name.

    Args:
        format_name: Name to use for dumping with this function. Example: "yaml_custom".
        dumper_fn: The dumper function to set. Example: ``yaml.safe_dump``.
    """
    dumpers[format_name] = dumper_fn


def set_omegaconf_loader():
    if omegaconf_support and "omegaconf" not in loaders:
        from ._optionals import get_omegaconf_loader

        set_loader("omegaconf", get_omegaconf_loader())
