import typing

import claripy

from . import primitives


# Functions that should be included in claripy, but aren't

def twos_comp_range_constraint(x: claripy.ast.bits, low: int, high: int) -> claripy.ast.Bool:
    """
    Generates a constraint which bounds the input argument in range [low, high), assuming a two's complement representation.

    :param clairpy.ast.bits x: The bits to constrain. Typically, this is a symbolic bitvector.
    :param int low: The lower bound on the range. This number may be negative.
    :param int high: The upper bound on the range.
    """
    num_bits = x.length
    sign_bit = x[(num_bits - 1):(num_bits - 1)]
    if low >= 0 and high >= 0:
        return ((sign_bit == 0) & (low <= x) & (x < high))
    elif low < 0 and high > 0:
        return (((sign_bit == 1) & (primitives.to_twos_comp(low, num_bits) <= x)) | ((sign_bit == 0) & (x < high)))
    elif low < 0 and high < 0:
        return ((sign_bit == 1) & (primitives.to_twos_comp(low, num_bits) <= x) & (x < primitives.to_twos_comp(high, num_bits)))
    else:
        raise ValueError("low was greater than high")

def simplify_kb(expr: claripy.ast.bits, kb: claripy.ast.Bool) -> claripy.ast.bits:
    """
    Simplifies a claripy AST expression, given some knowledge base (kb) of information

    :param claripy.ast.bits expr: The expression to simplify
    :param claripy.ast.Bool kb: The knowledge base which is used to simplify the expr. This is typically a series of equalities conjoined together.
    :return: A simplified version of the input expression, or the original expression if no simplification occurred.
    :rtype: claripy.ast.bits
    """
    if not expr.symbolic:
        return expr
    else:
        v = claripy.BVS("simp_kb_var", expr.length)
        solver = claripy.Solver()
        solver.add(kb)
        solver.add(v == expr)
        solver.simplify()
        for c in solver.constraints:
            if c.op == "__eq__" and len(c.args) == 2:
                (lhs, rhs) = c.args
                if lhs is v:
                    simplified_expr = rhs
                elif rhs is v:
                    simplified_expr = lhs
                else:
                    continue
                return simplified_expr
        return expr

def get_symbol_name(sym):
    return sym.args[0]

def model(constraints,
          extra_symbols: set[typing.Union[claripy.BVS, claripy.FPS]] | frozenset[typing.Union[claripy.BVS, claripy.FPS]]=frozenset(),
          n=1,
          **kwargs) -> list[dict[typing.Union[claripy.BVS, claripy.FPS], typing.Union[claripy.BVV, claripy.FPV]]]:
    # Computes at most n different satisfying assignments for a set of constraints
    solver = claripy.Solver()
    solver.add(constraints)
    extra_constraints = []
    generated_models = []
    while len(generated_models) < n:
        if solver.satisfiable(extra_constraints=extra_constraints, **kwargs):
            models = list(solver._models)

            ret = dict()

            def zero(sym):
                if sym.op == "BVS":
                    return claripy.BVV(0, sym.length)
                elif sym.op == "FPS":
                    return claripy.FPV(0.0, sym.sort)
                else:
                    raise ValueError("Unsupported op")

            def value(sym, v):
                if sym.op == "BVS":
                    return claripy.BVV(v, sym.length)
                elif sym.op == "FPS":
                    return claripy.FPV(v, sym.sort)
                else:
                    raise ValueError("Unsupported op")

            if len(models) > 0:
                # m maps symbol names to either integers or floats. We want symbols to map
                # to BVV or FPS, so we need to do that conversion here
                m = models[0].model
                for c in constraints:
                    for leaf in c.leaf_asts():
                        if leaf.symbolic:
                            leaf_name = get_symbol_name(leaf)
                            if leaf_name in m:
                                ret[leaf] = value(leaf, m[leaf_name])
                            else:
                                ret[leaf] = zero(leaf)

            # Set any symbols not in the model to 0
            for extra in extra_symbols:
                if extra not in ret:
                    ret[extra] = zero(extra)

            generated_models.append(ret)

            if len(generated_models) < n:
                # We need to find a solution that is different from the model we just found in at least one variable
                extra_constraints.append(claripy.Or(*[sym != val for (sym, val) in ret.items()]))
                # We don't want to accumulate models between calls to .satisfiable(), so clear out the model we
                # just found here
                solver._models.clear()
        else:
            return generated_models
    return generated_models