from inspect import getfullargspec

def validate(fun):
  annotations = fun.__annotations__
  def wrapped(*args, **kwargs):

    for arg, arg_name in zip(args, getfullargspec(fun).args):
      if arg_name in annotations:
        t = annotations[arg_name]
        assert isinstance(arg, t)

    for k, v in annotations.items():
      if k != "return" and k in kwargs:
        assert isinstance(kwargs[k], v)

    out = fun(*args, **kwargs)
    if "return" in annotations:
      assert isinstance(out, annotations["return"])
    return out

  return wrapped

class Term:
  pass

class Symbol(Term):
  @validate
  def __init__(self, wrapped):
    self.wrapped = wrapped

  def __repr__(self):
    return repr(self.wrapped)

  def __eq__(self, other):
    return self.wrapped == other.wrapped

class Abstr(Term):
  @validate
  def __init__(self, var : Symbol, body : Term):
    self.var = var
    self.body = body

  def to_tuple(self):
    return ("lambda", self.var, self.body)

  def __repr__(self):
    return repr(self.to_tuple())

class App(Term):
  @validate
  def __init__(self, t : Term, s : Term):
    self.t = t
    self.s = s

  def to_tuple(self):
    return (self.t, self.s)

  def __repr__(self):
    return repr(self.to_tuple())

@validate
def substitute(body : Term, var_name : Symbol, arg : Term):
  "capture avoiding substitution"
  if isinstance(body, Symbol):
    # change the symbol if the symbol is the same
    return arg if body == var_name else body
  elif isinstance(body, Abstr):
    # variable has same name as replaced variable, block substitution.
    if body.var == var_name:
      return body
    # otherwise, substitution recurses down
    else:
      return Abstr(body.var, substitute(body.body, var_name, arg))
  elif isinstance(body, App):
    # substitute into the function and the arguments
    return App(
      substitute(body.t, var_name, arg),
      substitute(body.s, var_name, arg)
    )

@validate
def hit(x : Term):
  "hit expression, perform beta reduction"
  if isinstance(x, App):
    fun = x.t
    body = x.s
    if isinstance(fun, Abstr):
      return substitute(fun.body, fun.var, body)
  return None

a = Symbol("a")
x = Symbol("x")
y = Symbol("y")

expr = Abstr(x, Abstr(y, App(x,y)))
print(expr)

print( hit(App(expr,a)) )