from enum import Enum
from . import yaml
import minizinc as mzn
from minizinc.helpers import check_result
import pathlib
import re
import subprocess
from json import load, loads, dumps
import jsonschema
import warnings


@yaml.mapping("!Test")
class Test:
    """
    Represents a test case

    Encoded using `!Test` in YAML.
    """

    def __init__(self, **kwargs):
        self.name = yaml.Undefined
        self.solvers = ["gecode", "cbc", "chuffed", "highs"]
        self.check_against = []
        self.expected = []
        self.options = {}
        self.extra_files = []
        self.markers = []
        self.type = "solve"

        for key, value in kwargs.items():
            setattr(self, key, value)

    def run(self, mzn_file, solver, default_options={}):
        """
        Runs this test case given an mzn file path, a solver name and some default options.

        Any options specified in this test case directly will override those provided in
        default options.

        Returns a tuple containing:
        - the model produced
        - the result of running the solver
        - a list of expected results
        - the actual obtained result

        Pass the actual obtained result to passed() to determine if the test case passed.
        """
        options = {k: v for k, v in default_options.items()}
        options.update(self.options)
        file = pathlib.Path(mzn_file)
        solver_name = solver
        extra_files = [file.parent.joinpath(other) for other in self.extra_files]

        try:
            model = mzn.Model([file] + extra_files)
            model.output_type = Solution
            if solver.endswith(".msc"):
                solver = mzn.Solver.load(file.parent.joinpath(solver))
            else:
                solver = mzn.Solver.lookup(solver)
            instance = mzn.Instance(solver, model)
            if self.type == "core":
                # The `core` message is not surfaced by the minizinc Python driver, so run the
                # executable directly and parse the unsatisfiable core(s) from the JSON stream.
                obtained = CoreResult.from_solve(solver_name, file, extra_files, options)
                result = obtained
            elif self.type == "solve":
                result = instance.solve(**options)
                obtained = Result.from_mzn(result)
            elif self.type == "compile":
                with instance.flat(**options) as (fzn, ozn, stats):
                    if solver.inputType == "JSON":
                        obtained = FlatZincJSON.from_mzn(fzn, file.parent)
                        result = obtained
                    else:
                        obtained = FlatZinc.from_mzn(fzn, file.parent)
                        result = obtained
            elif self.type == "output-model":
                with instance.flat(**options) as (fzn, ozn, stats):
                    obtained = OutputModel.from_mzn(ozn, file.parent)
                    result = obtained
            else:
                raise NotImplementedError("Unknown test case type")
        except mzn.MiniZincError as error:
            result = error
            obtained = Error.from_mzn(result)

        required = self.expected if isinstance(self.expected, list) else [self.expected]

        return model, result, required, obtained

    def passed(self, actual):
        """
        Returns whether the given result satisfies this test case
        """
        if isinstance(self.expected, list):
            return any(exp.check(actual) for exp in self.expected)
        return self.expected.check(actual)


@yaml.mapping("!Result")
class Result:
    """
    Represents a result generated by calling `solve()` on an instance.

    Encoded using `!Result` in YAML.
    """

    def __init__(self, **kwargs):
        self.status = yaml.Undefined
        self.solution = yaml.Undefined
        self.core = yaml.Undefined

        for k, v in kwargs.items():
            setattr(self, k, v)

    def check(self, actual):
        """
        Returns whether this result and the given result match.
        """
        if isinstance(actual, CoreResult):
            if self.status is not yaml.Undefined and actual.status != self.status:
                return False
            if self.core is not yaml.Undefined and self.core is not None:
                # The order of entries in a core is solver dependent, so compare as a multiset.
                return any(sorted(self.core) == sorted(c) for c in actual.cores)
            return True

        if not isinstance(actual, mzn.Result):
            return False

        if self.status is not yaml.Undefined and str(actual.status) != self.status:
            return False

        if self.solution is not yaml.Undefined and self.solution is not None:
            required = self.solution
            obtained = actual.solution

            if isinstance(required, list):
                return len(required) == len(obtained) and all(
                    x.is_satisfied(y) for x, y in zip(required, obtained)
                )
            else:
                return required.is_satisfied(obtained)

        return True

    @staticmethod
    def from_mzn(result):
        """
        Constructs a Result from a minizinc Result type produced by running a solver
        as opposed to deserializing from YAML.
        """
        instance = Result()
        instance.status = str(result.status)
        instance.solution = result.solution
        return instance


@yaml.mapping("!CoreResult")
class CoreResult:
    """
    Represents the unsatisfiable core(s) reported by a solver via ``%%%mzn-core``.

    The minizinc Python driver does not surface ``core`` messages, so for ``core`` test
    cases the executable is run directly and this object is constructed from its JSON stream.
    It is compared against an expected ``!Result`` with a ``core`` field.
    """

    def __init__(self, **kwargs):
        self.status = kwargs.get("status")
        self.cores = kwargs.get("cores", [])

    @staticmethod
    def from_solve(solver_name, file, extra_files, options):
        executable = str(mzn.default_driver._executable)
        if solver_name.endswith(".msc"):
            solver_arg = str(pathlib.Path(file).parent.joinpath(solver_name))
        else:
            solver_arg = solver_name
        cmd = [executable, "--solver", solver_arg, "--json-stream"]
        for key, value in options.items():
            flag = "--" + str(key).replace("_", "-")
            if value is True:
                cmd.append(flag)
            elif value is False:
                continue
            else:
                cmd += [flag, str(value)]
        cmd.append(str(file))
        cmd += [str(f) for f in extra_files]

        proc = subprocess.run(cmd, capture_output=True, text=True)
        status = None
        cores = []
        for line in proc.stdout.splitlines():
            line = line.strip()
            if not line:
                continue
            try:
                obj = loads(line)
            except ValueError:
                continue
            if obj.get("type") == "status":
                status = obj.get("status")
            elif obj.get("type") == "core":
                cores.append(list(obj.get("core", [])))
        return CoreResult(status=status, cores=cores)


@yaml.sequence("!SolutionSet")
class SolutionSet:
    """
    Represents a set of expected solutions that passes when
    every obtained solution matches an entry in this set.

    Represented with `!SolutionSet` in YAML
    """

    def __init__(self, *args):
        self.items = args

    def __iter__(self):
        return iter(self.items)

    def __len__(self):
        return len(self.items)

    def is_satisfied(self, other):
        """
        Returns whether or not every solution in `other` matches one in this solution set
        """
        return len(self) == len(other) and all(
            any(item.is_satisfied(o) for item in self.items) for o in other
        )


@yaml.mapping("!Solution")
class Solution:
    """
    A solution which is satisfied when every member of the this solution matches a member in the obtained solution.

    That is, the obtained solution can contain extra items that are not compared.

    Represented by `!Solution` in YAML.
    """

    def __init__(self, **kwargs):
        self.inner = {k: v for k, v in kwargs.items()}

    def items(self):
        return self.inner.items()

    def as_dict(self):
        return self.inner

    def is_satisfied(self, other):
        """
        Returns whether or not this solution is satisfied by an actual solution
        """

        def in_other(k, v):
            try:
                return v == other.inner[k]
            except KeyError:
                return False

        return all(in_other(k, v) for k, v in self.inner.items())


@yaml.mapping("!Error")
class Error:
    """
    A MiniZinc error with a type name and message.

    Represented by `!Error` in YAML.
    """

    def __init__(self, **kwargs):
        self.type = yaml.Undefined
        self.message = yaml.Undefined
        self.regex = yaml.Undefined

        for key, value in kwargs.items():
            setattr(self, key, value)

    def check(self, actual):
        """
        Checks if this error matches an actual obtained result/error
        """
        if not isinstance(actual, mzn.MiniZincError):
            return False

        if self.type is not yaml.Undefined:
            kind = type(actual)
            classes = (kind,) + kind.__bases__
            if not any(self.type == c.__name__ for c in classes):
                return False

        if self.message is not yaml.Undefined:
            return self.message == str(actual)

        if self.regex is not yaml.Undefined:
            return re.match(self.regex, str(actual), flags=re.M | re.S) is not None

        return True

    @staticmethod
    def from_mzn(error):
        """
        Creates an Error from a MiniZinc error (as opposed to from deserializing YAML)
        """
        instance = Error()
        instance.type = type(error).__name__
        instance.message = str(error)
        return instance


class CachedResult:
    """
    An object used to store a model and result from solving an instance

    This is so that `CheckItem`s do not have to re-solve an instance when
    verifying against another solver.
    """

    def __init__(self):
        self.model = None
        self.result = None
        self.obtained = None

    def test(self, solver):
        """
        Checks that the result stored in this object, when given as data for the
        stored model with the given solver, is satisfied.
        """
        solver_instance = mzn.Solver.lookup(solver)

        solution = (
            [s.as_dict() for s in self.result.solution]
            if isinstance(self.result.solution, list)
            else [self.result.solution.as_dict()]
        )
        passed = check_result(
            model=self.model,
            result=mzn.Result(self.result.status, solution, {}),
            solver=solver_instance,
        )
        return passed


@yaml.scalar("!FlatZinc")
class FlatZinc:
    """
    A FlatZinc result, encoded by !FlatZinc in YAML.
    """

    def __init__(self, path):
        self.path = path
        self.base = None
        self.fzn = None

    def check(self, actual):
        if not isinstance(actual, FlatZinc):
            return False

        if self.fzn is None:
            with open(actual.base.joinpath(self.path), encoding="utf-8") as f:
                self.fzn = f.read()

        return self.normalized() == actual.normalized()

    def normalized(self):
        """
        Quick and dirty normalization of FlatZinc

        Used to compare two `.fzn` files a bit more robustly than naive string comparison
        """
        fzn = self.fzn
        var_nums = re.findall(r"X_INTRODUCED_(\d+)_", fzn)
        for new, old in enumerate(sorted(set(int(x) for x in var_nums))):
            # Compress the introduced variable numbers, removing any holes and starting from 0
            needle = "X_INTRODUCED_{}_".format(old)
            replacement = "X_INTRODUCED_{}_".format(new)
            fzn = fzn.replace(needle, replacement)
        # Remove empty lines, truncate multiple spaces
        lines = [
            re.sub(
                r"::[^;]*",
                ann_sort,
                re.sub(r"  +", " ", line),
                count=1,
            )
            for line in fzn.split("\n")
            if len(line.strip()) > 0 and not line.strip().startswith("%")
        ]

        # Sort array elements in specific constraints
        for constraint_name in ["array_bool_and", "array_bool_or", "bool_clause"]:
            lines = [
                (
                    re.sub(
                        rf"(\[)([^\]]+)(\])",
                        lambda m: m.group(1)
                        + ", ".join(sorted([x.strip() for x in m.group(2).split(",")]))
                        + m.group(3),
                        line,
                        count=0,
                    )
                    if line.startswith(f"constraint {constraint_name}(")
                    else line
                )
                for line in lines
            ]

        return "\n".join(sorted(lines))  # Sort the FlatZinc lines

    def get_value(self):
        if self.fzn is None:
            return self.path
        return self.normalized()

    @staticmethod
    def from_mzn(fzn, base):
        """
        Creates a `FlatZinc` object from a `File` returned by `flat()` in the minizinc interface.

        Also takes the base path the mzn file was from, so that when loading the expected fzn file
        it can be done relative to the mzn path.
        """
        with open(fzn.name, encoding="utf-8") as file:
            instance = FlatZinc(None)
            instance.fzn = file.read()
            instance.base = base
            return instance


def ann_sort(match):
    s = match.group(0)
    anns = sorted([ann for ann in s.split("::") if ann != ""])
    return "::" + "::".join(anns)


@yaml.scalar("!FlatZincJSON")
class FlatZincJSON:
    """
    A FlatZinc JSON result, encoded by !FlatZincJSON in YAML.
    """

    def __init__(self, path):
        self.path = path
        self.base = None
        self.fzn = None

    def check(self, actual):
        if not isinstance(actual, FlatZincJSON):
            return False

        if self.fzn is None:
            with open(actual.base.joinpath(self.path), encoding="utf-8") as f:
                self.fzn = load(f)

        obtained = actual.normalized()

        schema_path = (
            pathlib.Path(__file__).parent.parent.parent / "docs/en/fznjson.json"
        )
        if schema_path.exists():
            with open(schema_path, encoding="utf-8") as fp:
                schema = load(fp)
            jsonschema.validate(instance=obtained, schema=schema)
        else:
            warnings.warn(
                "Unable to find FlatZinc JSON schema. Cannot verify validity."
            )

        return self.normalized() == obtained

    def normalized(self):
        var_nums = set()
        for v in list(self.fzn["variables"].keys()) + list(self.fzn["arrays"].keys()):
            m = re.match(r"X_INTRODUCED_(\d+)_", v)
            if m is not None:
                var_nums.add(int(m.group(1)))
        var_map = {
            f"X_INTRODUCED_{src}_": f"X_INTRODUCED_{dst}_"
            for dst, src in enumerate(sorted(var_nums))
        }

        def get_ident(obj):
            if isinstance(obj, dict):
                return obj.get("id")
            return obj

        def walk(obj, key=None):
            if isinstance(obj, dict):
                return {walk(k): walk(v, k) for k, v in obj.items()}
            if isinstance(obj, list):
                items = [walk(v) for v in obj]
                if key == "ann":
                    return sorted(items, key=get_ident)
                return items
            if isinstance(obj, str):
                return var_map.get(obj, obj)
            return obj

        return walk(self.fzn)

    def get_value(self):
        if self.fzn is None:
            return self.path
        return dumps(self.normalized())

    @staticmethod
    def from_mzn(fzn, base):
        """
        Creates a `FlatZinc` object from a `File` returned by `flat()` in the minizinc interface.

        Also takes the base path the mzn file was from, so that when loading the expected fzn file
        it can be done relative to the mzn path.
        """
        with open(fzn.name, encoding="utf-8") as file:
            instance = FlatZincJSON(None)
            instance.fzn = load(file)
            instance.base = base
            return instance


def ann_sort(match):
    s = match.group(0)
    anns = sorted([ann for ann in s.split("::") if ann != ""])
    return "::" + "::".join(anns)


@yaml.scalar("!OutputModel")
class OutputModel:
    """
    An OZN result, encoded by !OutputModel in YAML.
    """

    def __init__(self, path):
        self.path = path
        self.base = None
        self.ozn = None

    def check(self, actual):
        if not isinstance(actual, OutputModel):
            return False

        if self.ozn is None:
            with open(actual.base.joinpath(self.path), encoding="utf-8") as f:
                self.ozn = f.read()

        return self.ozn == actual.ozn

    def get_value(self):
        if self.ozn is None:
            return self.path
        return self.ozn

    @staticmethod
    def from_mzn(ozn, base):
        """
        Creates a `OutputModel` object from a `File` returned by `flat()` in the minizinc interface.

        Also takes the base path the mzn file was from, so that when loading the expected ozn file
        it can be done relative to the mzn path.
        """
        with open(ozn.name, encoding="utf-8") as file:
            instance = OutputModel(None)
            instance.ozn = file.read()
            instance.base = base
            return instance


@yaml.mapping("!Suite")
class Suite:
    """
    Represents a test suite configuration

    Encoded using `!Suite` in YAML.
    """

    def __init__(self, **kwargs):
        self.solvers = None
        self.options = {}
        self.strict = True
        self.regex = r"."
        self.includes = ["*"]
        self.excludes = []

        for key, value in kwargs.items():
            setattr(self, key, value)
