"""Defines serialization and deserialization functions for optplans."""
import copy
import json
from typing import Any, Callable, List, Set, Union
import warnings

from schematics import models
from schematics import types

from spins.invdes.problem_graph import optplan

# Global variable for managing name collisions.
# Maps type names to counters indicat
problem_graph_name_map = {}  # pylint: disable=invalid-name


def reset_graph() -> None:
    """Resets the problem graph.

    Currently this resets the map that tracks graph node names. Calling this
    function is not required when a new problem graph is created.
    """
    global problem_graph_name_map  # pylint: disable=invalid-name
    problem_graph_name_map = {}


def generate_name(model_type: str) -> str:
    """Generates a name for a model with type `model_type`.

    This is used by auto-name generation. The name will be of the form
    `model_type.num` where `num` is a number that is tracked by
    `problem_graph_name_map`. Note that this does NOT guarantee name uniqueness
    unless the following are true:
        1) All nodes created before the most recent call to `reset_graph` are
           not present in the current graph (since `reset_graph` resets the
           counters).
        2) All user-defined names do not carrying the period (.).

    Args:
        model_type: String model type.

    Returns:
        A unique automatically generated name subject the constraints stated
        above.
    """
    global problem_graph_name_map  # pylint: disable=invalid-name

    if model_type not in problem_graph_name_map:
        problem_graph_name_map[model_type] = 0

    name = "{}.{}".format(model_type, problem_graph_name_map[model_type])

    problem_graph_name_map[model_type] += 1

    return name


def top_level_type(model: optplan.optplan.ProblemGraphNode) -> str:
    """Determines the top level type.

    Args:
        model: Model.

    Returns:
        Top-level model type.
    """
    return model.type.split(".")[0]


def autoname(model: models.Model) -> None:
    """Autonamtically names `optplan.ProblemGraphNode` recursively.

    This function recursively searches through fields to find any
    `optplan.ProblemGraphNode` instances. Any unnamed `optplan.ProblemGraphNode` instances are
    given a name generated by `generate_name`.

    Args:
        model: Model to recursively search.
    """
    _autoname_helper(model, set())


def _autoname_helper(value: Any, visited: Set[str]) -> None:
    """Recursively names any `optplan.ProblemGraphNode`.

    Args:
        value: Value to recurse on.
        visited: List of `optplan.ProblemGraphNode` names that were visited.
            This is to prevent recursing on models that have already been
            processed.
    """
    # If we encounter a `optplan.ProblemGraphNode`, name it.
    if isinstance(value, optplan.ProblemGraphNode):
        if not value.name:
            value.name = generate_name(value.type)

        if value.name in visited:
            return

        visited.add(value.name)

    # Recursively search for `optplan.ProblemGraphNode`.
    if isinstance(value, models.Model):
        for field_name in value.keys():
            _autoname_helper(value[field_name], visited)
    elif isinstance(value, list):
        for item in value:
            _autoname_helper(item, visited)
    elif isinstance(value, dict):
        for _, item in value.items():
            _autoname_helper(item, visited)


def _iter_optplan_fields(
        model: models.Model,
        visited: Set[int],
        process_field: Callable[
            [models.Model, Union[str, optplan.ProblemGraphNode]], None],
        pass_field_info: bool = False) -> None:
    """Iterate through every reference field recursively in the model.

    This function operates by performing a depth-first search through every
    model and calling `process_field` on every field in the model. This is
    useful, for example, for finding all the reference nodes and replacing them
    with their names (see `_replace_ref_nodes_with_names`).

    Args:
        model: The model to recrusively search.
        visited: A list of model ids (computed from `id`) to indicate which
            nodes have been visited before.
        process_field: Function to call on every reference field. This function
            should accept the following arguments:
                - A `models.Model` model corresponding to the model with the
                  field.
                - The model (or name of the model) of the reference field.
            The function should return an object that should replace the
            reference field with. If it returns `None`, the optplan will be
            left intact.
        pass_field_info: If `True`, the field information is passed as a third
            argument to `process_field`.
    """
    # Handle only dictionaries, lists, and models. Every other type is assumed
    # to not contain a `optplan.ReferenceType`.
    if not isinstance(model, models.Model):
        return

    # Avoid double processing nodes.
    if id(model) in visited:
        return
    visited.add(id(model))

    # Wrap `process_field` so that returning `None` is same as returning the
    # child.
    def process_field_wrapped(
            parent: models.Model,
            child: Union[str, optplan.ProblemGraphNode],
            field_type: optplan.ReferenceType,
    ) -> optplan.ProblemGraphNode:
        if pass_field_info:
            return_val = process_field(parent, child, field_type)
        else:
            return_val = process_field(parent, child)

        if return_val is None:
            return child
        return return_val

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")

        for field_name, field_type in model.fields.items():
            if model[field_name] is None:
                continue
            # Recursively check every field for `optplan.ReferenceType` in
            # compound types.
            if isinstance(model[field_name], models.Model):
                _iter_optplan_fields(
                    model[field_name],
                    visited,
                    process_field,
                    pass_field_info=pass_field_info)
            elif isinstance(field_type, types.ListType):
                for item in model[field_name]:
                    _iter_optplan_fields(
                        item,
                        visited,
                        process_field,
                        pass_field_info=pass_field_info)
            elif isinstance(field_type, types.DictType):
                for _, value in model[field_name].items():
                    _iter_optplan_fields(
                        value,
                        visited,
                        process_field,
                        pass_field_info=pass_field_info)

            # Actually process any reference fields.
            if isinstance(field_type, optplan.ReferenceType):
                model[field_name] = process_field_wrapped(
                    model, model[field_name], field_type)
            elif (isinstance(field_type, types.ListType) and
                  isinstance(field_type.field, optplan.ReferenceType)):
                model[field_name] = [
                    process_field_wrapped(model, m, field_type.field)
                    for m in model[field_name]
                ]
            elif (isinstance(field_type, types.DictType) and
                  isinstance(field_type.field, optplan.ReferenceType)):
                model[field_name] = {
                    key: process_field_wrapped(model, m, field_type.field)
                    for key, m in model[field_name].items()
                }


def validate(plan: optplan.OptimizationPlan) -> None:
    """Validates an optimization plan.

    Args:
        plan: Plan to validate.

    Raises:
        Exception if plan is not valid.
    """
    # Do schema validation.
    # Rather than validate the entire plan all at once, we first validate
    # node-by-node because `schematics` is very bad at giving informative
    # errors.
    node_list = []
    if plan.nodes:
        node_list += plan.nodes
    if plan.transformations:
        node_list += plan.transformations

    for node in node_list:
        try:
            # Store the node name before hand because `validate` will
            # actually destroy the node is the validation fails.
            node_name = node.name
            node.validate()
        except Exception as exc:
            raise ValueError("Error encountered when validating node {}".format(
                node_name)) from exc

    # Now validate the plan schema itself just in case we missed something
    # from the previous checks.
    plan.validate()

    # Check for non-unique nodes.
    names = set()
    for node in plan.nodes:
        if node.name in names:
            raise ValueError("Nonunique name found: {}".format(node.name))
        names.add(node.name)


def validate_references(model: models.Model) -> None:
    """Validates that all `optplan.ReferenceType` are valid.

    This function recursively checks to see if every `ReferenceType` field
    is actually holding a reference to the appropriate type.

    Args:
        model: The model to check.

    Raises:
        ValueError: If a bad type is encountered.
    """

    def process_field(parent: models.Model,
                      child: Union[str, optplan.ProblemGraphNode],
                      field_type: optplan.ReferenceType) -> None:
        if not child:
            return

        if not isinstance(child, (str, field_type.reference_type)):
            raise ValueError("Expected type {} for node {}, got {}".format(
                field_type.reference_type, child, type(child)))

    _iter_optplan_fields(model, set(), process_field, pass_field_info=True)


def _validate_optplan_version(version: str) -> None:
    """Validates an optplan has a compatible version.

    Raises:
        ValueError: On incompatible version.
    """
    version_parts = [int(part) for part in version.split(".")]
    if version_parts[0] < 0 or version_parts[1] < 2 or version_parts[2] < 3:
        raise ValueError(
            "Optplan must be at least version 0.2.1, got {}".format(version))


def loads(serialized_plan: str) -> optplan.OptimizationPlan:
    """Loads a serialized optimization plan.

    Args:
        serialized_plan: Serialized plan.

    Returns:
        Stored `optplan.OptimizationPlan`.
    """
    plan = optplan.OptimizationPlan(json.loads(serialized_plan))
    validate(plan)
    _validate_optplan_version(plan.version)

    # Unpack the optimization plan.
    # First create a dictionary mapping strings to the node.
    node_map = {}
    for node in plan.nodes:
        node_map[node.name] = node

    # Now, recursively fill any node names with the actual node.
    def process_field(model: models.Model, child_model: str) -> models.Model:
        return node_map[child_model]

    visited = set()
    if plan.transformations:
        for transformation in plan.transformations:
            _iter_optplan_fields(transformation, visited, process_field)
    for node in plan.nodes:
        _iter_optplan_fields(node, visited, process_field)

    validate_references(plan)
    return plan


def dumps(plan: optplan.OptimizationPlan) -> str:
    """Serializes `plan` into a string.

    This is useful for saving `optplan.OptimizationPlan` as a string.

    Args:
        plan: Plan.

    Raises:
        ValueError: If any `optplan.ProblemGraphNode` names are duplicated.
    """
    # Make a deepcopy so we do not touch the original.
    plan = copy.deepcopy(plan)
    validate_references(plan)

    # First make sure every graph node is named.
    autoname(plan)

    # Next extract out every graph node.
    model_list = []
    _extract_nodes(plan, model_list)

    # Now replace all reference types with the correct names.
    _replace_ref_nodes_with_names(plan, model_list)

    # Update the list of nodes.
    plan.nodes = model_list

    validate(plan)
    return json.dumps(plan.to_primitive())


def _extract_nodes(value: Any,
                   model_list: List[optplan.ProblemGraphNode]) -> None:
    """Finds any nodes recursively in `value` and appends them to `model_list`.

    This is used by `dumps` to find all relevant `optplan.ProblemGraphNode`
    instances.

    Args:
        value: Value to recurse on.
        model_list: List of `optplan.ProblemGraphNode` models found so far.
    """
    if isinstance(value, optplan.ProblemGraphNode):
        if value not in model_list:
            model_list.append(value)
        else:
            # This node has been processed before so quit.
            return

    if isinstance(value, models.Model):
        for item in value.keys():
            _extract_nodes(value[item], model_list)
    elif isinstance(value, list):
        for item in value:
            _extract_nodes(item, model_list)
    elif isinstance(value, dict):
        for _, item in value.items():
            _extract_nodes(item, model_list)


def _replace_ref_nodes_with_names(
        model: models.Model,
        model_list: List[optplan.ProblemGraphNode]) -> None:
    """Replaces `optplan.ProblemGraphNode` references with names.

    This is used by `dumps` to replace all `optplan.ProblemGraphNode` instances
    with their names for serialization.

    Args:
        model: Model to recurse on.
        model_list: List of `optplan.ProblemGraphNode` models.
    """

    # Replace model reference with a name.
    def process_field(model: models.Model, child_model: models.Model) -> str:
        if isinstance(child_model, str):
            return child_model

        ind = model_list.index(child_model)
        return model_list[ind].name

    _iter_optplan_fields(model, set(), process_field)
