import ast 
from pprint import pprint
from z3 import *
code = """
def f(x,y):
  if x > 5: 
    if y > 1:
      return 2
    else:
      return y
    y = 2
  else:
    if  x > 5:
        return x
    else:
      return 2

"""

tree= ast.parse(code)

class Path:
    def __init__(self, constraints, returns):
        self.constraints = constraints  # list of Z3 conditions
        self.returns = returns          # return value

def eval_expr(expr, symbols):
    "Convert Python AST expression to Z3 expression."
    if isinstance(expr, ast.Compare):
        left = eval_expr(expr.left, symbols)
        right = eval_expr(expr.comparators[0], symbols)
        op = expr.ops[0]
        if isinstance(op, ast.Gt): return left > right
        if isinstance(op, ast.Lt): return left < right
        if isinstance(op, ast.GtE): return left >= right
        if isinstance(op, ast.LtE): return left <= right
        if isinstance(op, ast.Eq): return left == right
        if isinstance(op, ast.NotEq): return left != right
    elif isinstance(expr, ast.Name):
        return symbols[expr.id]
    elif isinstance(expr, ast.Constant):
        return expr.value
    raise NotImplementedError(f"Unsupported: {ast.dump(expr)}")

def walk(node, path_constraints, paths, symbols):
    if isinstance(node, ast.If):
        cond = eval_expr(node.test, symbols)
        # Then branch
        walk(node.body[0], path_constraints + [cond], paths, symbols)

        # Else branch
        walk(node.orelse[0], path_constraints + [Not(cond)], paths, symbols)

    elif isinstance(node, ast.Return):
        paths.append(Path(path_constraints, node.value))


def visit(nodes, counter, path_constraints, paths, symbols):
    node = nodes[counter]
    if isinstance(node, ast.If):
        walk(node, path_constraints, paths, symbols)
    elif isinstance(node, ast.Return):
        walk(node, path_constraints, paths, symbols)
    else:
       visit(nodes, counter+1, path_constraints, paths, symbols)



print(code)
path_constraint = []
paths = []
symbol_x = Int('x')
symbol_y = Int('y')
symbols = {'x': symbol_x, 'y': symbol_y}

func_def = tree.body[0]  # get the function node
visit(func_def.body,0, [], paths, symbols)
#walk(func_def.body[0], [], paths, symbols)

for i, path in enumerate(paths):
    print("Path Constraint", i+1, ":", path.constraints)

# ---------- Step 4: Solve Path Conditions ----------
print("Generated Test Inputs:")
for i, path in enumerate(paths):
    solver = Solver()
    solver.add(path.constraints)
    if solver.check() == sat:
        model = solver.model()
        print(f"✅ Path {i+1}: return '{ast.unparse(path.returns)}' | input: x = {model[symbol_x]}, y={model[symbol_y]}")
    else:
        print(f"❌ Path {i+1} (return {ast.unparse(path.returns)}) is unreachable")