"""Class definition of UPPAAL NTA."""
from typing import Callable, List, Optional, Type, Union, cast

import lxml.etree as ET

from uppaalpy.classes import constraint_patcher as cp
from uppaalpy.classes import nodes as n
from uppaalpy.classes import templates as te
from uppaalpy.classes import transitions as tr
from uppaalpy.classes.context import Context
from uppaalpy.classes.expr import ClockConstraintExpression
from uppaalpy.classes.simplethings import Declaration, Query, SystemDeclaration


class NTA:
    """Class that holds information about a network of timed automata.

    Attributes:
        declaration: Declaration object.
        context: Context object generated by parsing the declaration.
        templates: List of template objects.
        system: SystemDeclaration object.
        queries: List of Query objects.
        patch_cache: Constraint cache object. Changes on constraints are stored
            here for fast write operations.
        _associated_file: String for denoting the path of the file read. Used
            by constraint patcher. Set by the from_xml method.
        _doctype: String for xml doctype. Set by from_xml, and used by to_file.
    """

    def __init__(self, **kwargs) -> None:
        """Initialize NTA from keyword arguments.

        Args:
            **kwargs: Keyword arguments for initializing the NTA. See class docstring.
        """
        self.declaration = kwargs.get("declaration")  # type: Optional[Declaration]
        self.context = kwargs["context"]  # type: Context
        self.templates = kwargs.get("templates") or []  # type: List[te.Template]
        self.system = kwargs["system"]  # type: SystemDeclaration
        self.queries = kwargs["queries"]  # type: List[Query]
        self.patch_cache = cp.ConstraintCache(self)
        self._associated_file: str = ""
        self._doctype: str = ""

    @classmethod
    def from_xml(cls: Type["NTA"], path: str) -> "NTA":
        """Given a xml file path, construct an NTA from that xml file."""
        et = ET.parse(path)
        obj = cls.from_element(et.getroot())
        obj._associated_file = path
        obj._doctype = et.docinfo.doctype
        return obj

    @classmethod
    def from_element(cls: Type["NTA"], et) -> "NTA":
        """Construct NTA from Element, and return it."""
        kw = {}
        kw["declaration"] = Declaration.from_element(et.find("declaration"))

        if kw["declaration"] is None or kw["declaration"].text is None:
            kw["context"] = Context.parse_context(None)
        else:
            kw["context"] = Context.parse_context(kw["declaration"].text)
        kw["templates"] = [
            te.Template.from_element(template, kw["context"])
            for template in et.iterchildren("template")
        ]
        kw["system"] = SystemDeclaration.from_element(et.find("system"))
        if et.find("queries") is None:
            kw["queries"] = []
        else:
            kw["queries"] = [
                Query.from_element(query) for query in et.find("queries").iter("query")
            ]
        return cls(**kw)

    def to_element(self):
        """Construct an Element object, and return it."""
        root = ET.Element("nta")

        if self.declaration is not None:
            et_declaration = ET.SubElement(root, "declaration")
            et_declaration.text = self.declaration.text

        for template in self.templates:
            root.append(template.to_element())

        root.append(self.system.to_element())
        queries = ET.SubElement(root, "queries")
        for query in self.queries:
            queries.append(query.to_element())

        return root

    def to_file(self, path: str) -> None:
        r"""Convert the NTA to an element tree and write it into a file.

        File will be printed with pretty printing and '\t' indentation.

        Args:
            path: String denoting the path of the output file.
        """
        elm = self.to_element()
        ET.indent(elm, "\t")
        ET.ElementTree(elm).write(
            path, encoding="utf-8", xml_declaration=True, doctype=self._doctype
        )

    def change_clock_constraint(
        self,
        obj: Union[tr.Transition, n.Location],
        operation: str,
        ccexp: ClockConstraintExpression,
        *,
        threshold_function: Optional[Callable[[str], str]] = None,
    ) -> None:
        """Insert/remove/update a simple constraint expression.

        A constraint patch is created to be cached in the patch cache for the
        changed constraint. Operations can be "insert", "remove", or "update".

        Args:
            obj: Either a Transition or a Location.
            operation: "insert", "remove", or "update".
            ccexp: ClockConstraintExpression object.
                Insert mode: ClockConstraintExpression to be inserted to the
                    guard/invariant label.
                Remove mode: ClockConstraintExpression to be removed from the
                    guard/invariant label.
                Update mode: ClockConstraintExpression whose threshold will be
                    changed.
            threshold_function: Function f: str -> str used in Update mode. The
                treshold x is set to threshold_function(x). Notice that x is
                either an identifier or a value. In both cases its type is str.

        Insert mode:
        A new ClockConstraintExpression is inserted to the guard/invariant
        label. If no guard exists, one is created located between the source and
        target for guards, and on the location for invariants. A
        ConstraintInsert is created and cached in the patch cache.

        Remove mode:
        A ClockConstraintExpression is removed from the guard/invariant. If the
        expression to be removed is the only remaining constraint on the object,
        the label is removed, as well. A ConstraintRemove is created and cached.

        Update mode:
        The ClockConstraintExpression object is updated with the threshold_function.
        A ConstraintUpdate is created and cached.
        """
        label = obj.get_constraint_label()

        if operation == "insert":
            if label is None:  # Create guard or invariant
                obj.create_constraint_label(ccexp, self.context)
                change = cp.ConstraintInsert(
                    ccexp, newly_created=obj.get_constraint_label()
                )
            else:
                label.constraints.append(ccexp)
                change = cp.ConstraintInsert(ccexp)

        elif operation == "remove":
            label.constraints.remove(ccexp)
            if label.constraints == []:
                change = cp.ConstraintRemove(ccexp, True)
                obj.remove_constraint_label()
            else:
                change = cp.ConstraintRemove(ccexp)

        else:  # operation == "update"
            new = threshold_function(ccexp.threshold)
            change = cp.ConstraintUpdate(ccexp, new)
            index = label.constraints.index(ccexp)
            label.constraints.pop(index)
            label.constraints.insert(index, change.generate_new_constraint())

        patch = cp.ConstraintPatch(cast(te.Template, obj.template), change, obj)
        self.patch_cache.cache(patch)

    def flush_constraint_changes(self, out_path):
        """Read the associated_file and write the modified version to new file.

        The xml file the NTA is constructed from is read again and the lines
        are updated according to the patch_cache's update records.
        See: constraint_patcher.py
        """
        with open(self._associated_file) as input_file:
            lines = input_file.readlines()

        self.patch_cache.apply_patches(lines)

        with open(out_path, "w") as output_file:
            output_file.writelines(lines)
