#
# This file is distributed under the MIT License. See LICENSE.md for details.
#
from typing import Dict

import yaml
from pydantic import BaseModel
from pydantic.main import ModelMetaclass


def get_monkey_patching_metaclass(
    substitutions: Dict[str, type],
    register_global_yaml_helpers=False,
):
    """Returns a metaclass which transparently substitutes some types with others. This way pydantic uses the substitute
    types when creating the validators for its types, without the need to modify the autogenerated classes.
    """

    class MonkeyPatchingMetaClass(ModelMetaclass):
        class YamlLoader(yaml.Loader):
            pass

        class YamlDumper(yaml.Dumper):
            def increase_indent(self, flow=False, indentless=False):
                """Improves indentation"""
                return super(MonkeyPatchingMetaClass.YamlDumper, self).increase_indent(flow, False)

            def analyze_scalar(self, scalar):
                """Ensures literals starting with ? are quoted to make LLVM YAML parser happy"""
                analysis = super(MonkeyPatchingMetaClass.YamlDumper, self).analyze_scalar(scalar)
                if isinstance(analysis.scalar, str) and analysis.scalar.startswith("?"):
                    analysis.allow_flow_plain = False
                    analysis.allow_block_plain = False
                return analysis

        @staticmethod
        def __new__(mcs, clsname, bases, namespace):
            substitution = substitutions.get(clsname)
            if substitution is not None:
                return substitution

            created_class = super(MonkeyPatchingMetaClass, mcs).__new__(
                mcs, clsname, bases, namespace
            )

            if "__root__" in namespace:

                def yaml_representer(dumper: yaml.dumper.Dumper, instance):
                    classname = instance.__root__.__class__.__name__
                    tag = f"!{classname}"
                    return dumper.represent_mapping(
                        tag,
                        {k: v for k, v in instance.__root__._iter(exclude_none=True)},
                    )

            else:

                def yaml_representer(dumper: yaml.dumper.Dumper, instance):
                    return dumper.represent_dict(
                        {k: v for k, v in instance._iter(exclude_none=True)},
                    )

            mcs.YamlDumper.add_representer(
                created_class,
                yaml_representer,
            )
            if register_global_yaml_helpers:
                yaml.add_representer(
                    created_class,
                    yaml_representer,
                )

            def yaml_constructor(loader, node):
                mapping = loader.construct_mapping(node, deep=True)
                return created_class(**mapping)

            tag = f"!{clsname}"

            mcs.YamlLoader.add_constructor(tag, yaml_constructor)
            if register_global_yaml_helpers:
                yaml.add_constructor(tag, yaml_constructor)

            return created_class

    return MonkeyPatchingMetaClass


def get_monkey_patching_base_class(
    substitutions: Dict[str, type],
    register_global_yaml_helpers=False,
):
    """Returns a base class which transparently substitutes some model types with others"""
    MonkeyPatchingMetaClass = get_monkey_patching_metaclass(
        substitutions,
        register_global_yaml_helpers=register_global_yaml_helpers,
    )

    class MonkeyPatchingBaseClass(BaseModel, metaclass=MonkeyPatchingMetaClass):
        class Config:
            # Allows enums to be converted to strings when calling dict() on a model instance
            use_enum_values = True

    return MonkeyPatchingBaseClass