#!/usr/bin/env python3

# This program implements a variant of the F1 language extended with primitive recursion over
# natural numbers. The syntax rougly is as follows:
#
#   EXPR ::= [0-9]+                  ; numbers
#          | [a-z]+                  ; variables
#          | true                    ; true constant
#          | false                   ; false constant
#          | (+ EXPR EXPR)           ; addition
#          | (if EXPR EXPR EXPR)     ; branching
#          | (rec EXPR EXPR EXPR)    ; primitive recursion
#          | (fn VAR TEXPR EXPR)     ; lambda abstraction
#
#   TEXPR ::= int                    ; integer type
#          |  bool                   ; boolean type
#          |  (tfn VAR TEXPR TEXPR)  ; function type

# FIXME Decide on how to report type checking errors. Currently it is done via a mix of returning
#  None and throwing exceptions. Same for eval.

class AstNum(int):
    """A (natural) number."""
    def typeof(self, context):
        return int
    def eval(self, c):
        return self
    def __repr__(self):
        return f"AstNum({self.real})" # Ugly hack

class AstVar(str):
    """A variable."""
    def typeof(self, context):
        return context[self]
    def eval(self, c):
        return c[self]
    def __repr__(self):
        return f"AstVar({self})"

class AstAdd(object):
    """An addition expression."""
    def __init__(self, e0, e1):
        self.e0, self.e1 = e0, e1
    def typeof(self, context):
        t0 = self.e0.typeof(context)
        t1 = self.e1.typeof(context)
        if t0 == int and t1 == int:
            return int
        else:
            return None
    def eval(self, c):
        v0 = self.e0.eval(c)
        v1 = self.e1.eval(c)
        return v0 + v1
    def __repr__(self):
        return f"AstAdd({repr(self.e0)},{repr(self.e1)})"

class AstRec(object):
    """A primitive recursion expression."""
    def __init__(self, e0, e1, e2):
        self.e0, self.e1, self.e2 = e0, e1, e2
    def typeof(self, c):
        target_type = self.e0.typeof(c)
        base_type = self.e1.typeof(c)
        step_type = self.e2.typeof(c)
        if target_type != int:
            return None
        if isinstance(step_type, TypeFn) and isinstance(step_type.tr, TypeFn) and \
                step_type.at == int and step_type.tr.at == base_type and \
                step_type.tr.tr == base_type:
            return base_type
        else:
            return None
    def eval(self, c):
        target = self.e0.eval(c)
        base = self.e1.eval(c)
        step = self.e2.eval(c)
        # FIXME This is broken. There should only be one case. Evaluation of an expression with
        # the type int should only evaluate to python values that are one of AstNum or int.
        if not (isinstance(target, AstNum) or isinstance(target, int)):
            return None
        if base == None:
            return None
        elif step == None:
            return None
        if target > 0:
            target0 = AstNum(target - 1)
            return AstApp(AstApp(step, target0), AstRec(target0, self.e1, step)).eval(c)
        else:
            return base
    def __repr__(self):
        return f"AstRec({self.e0},{self.e1},{self.e2})"

class AstIf(object):
    """If expression"""
    def __init__(self, e0, e1, e2):
        self.e0, self.e1, self.e2 = e0, e1, e2
    def typeof(self, context):
        t0 = self.e0.typeof(context)
        t1 = self.e1.typeof(context)
        t2 = self.e2.typeof(context)
        if t0 == bool and t1 == t2:
            return t1
        else:
            return None
    def eval(self, c):
        v0 = self.e0.eval(c)
        if v0 == AstTrue():
            return self.e1.eval(c)
        elif v0 == AstFalse():
            return self.e2.eval(c)
        else:
            return None
    def __repr__(self):
        return f"AstIf({repr(self.e0)},{repr(self.e1)},{repr(self.e2)})"

class AstLet(object):
    """Let expression."""
    def __init__(self, x, e0, e1):
        self.x, self.e0, self.e1 = x, e0, e1
    def typeof(self, c0):
        t0 = self.e0.typeof(c0)
        c1 = c0.copy()
        c1[self.x] = t0
        return self.e1.typeof(c1)
    def eval(self, c0):
        v0 = self.e0.eval(c0)
        c1 = c0.copy()
        c1[self.x] = v0
        return self.e1.eval(c1)
    def __repr__(self):
        return f"AstLet({repr(self.x)},{repr(self.e0)},{repr(self.e1)})"

class AstTrue(object):
    """The true constant expression.""" 
    def typeof(self, context):
        return bool
    def eval(self, c):
        return self
    def __repr__(self):
        return f"AstTrue()"
    def __eq__(self, other):
        return isinstance(other, AstTrue)

class AstFalse(object):
    """The false constant expression.""" 
    def typeof(self, context):
        return bool
    def eval(self, c):
        return self
    def __repr__(self):
        return f"AstFalse()"
    def __eq__(self, other):
        return isinstance(other, AstFalse)

class AstFn(object):
    """Lambda abstraction expression.""" 
    def __init__(self, x, t, e):
        self.x, self.t, self.e = x, t, e
    def typeof(self, c0):
        c1 = c0.copy()
        c1[self.x] = self.t
        r = self.e.typeof(c1)
        if r:
            return TypeFn(self.x, self.t, r)
        else:
            return None
    def eval(self, c):
        return self
    def __repr__(self):
        return f"AstFn({repr(self.x)},{repr(self.t)},{repr(self.e)})"

class AstApp(object):
    """Lambda application expression.""" 
    def __init__(self, e0, e1):
        self.e0, self.e1 = e0, e1
    def typeof(self, c):
        t0 = self.e0.typeof(c)
        t1 = self.e1.typeof(c)
        if t0.at == t1:
            return t0.tr
    def eval(self, c0):
        v0 = self.e0.eval(c0)
        v1 = self.e1.eval(c0)
        c1 = c0.copy()
        c1[v0.x] = v1
        return v0.e.eval(c1)
    def __repr__(self):
        return f"AstApp({repr(self.e0)},{repr(self.e1)})"

class TypeFn(object):
    """Function type expression.""" 
    def __init__(self, x, at, tr):
        self.x, self.at, self.tr = x, at, tr
    def __repr__(self):
        return f"AstFn({repr(self.x)},{repr(self.at)},{repr(self.tr)})"
    def __eq__(self, other):
        return isinstance(other, type(self)) and self.x == other.x and \
                self.at == other.at and self.tr == other.tr

def readsexpr(txt):
    """Parse text containing balanced parenthesises into symbolic expressions."""
    wc = ['\n', '\t']
    pc = ['(', ')']
    for c in wc: 
        txt = txt.replace(c, " ")
    for c in pc: 
        txt = txt.replace(c, f" {c} ")
    tokens = [s for s in txt.split(" ") if len(s) > 0]
    stack = [[]]
    for t in tokens:
        if t == "(":
            stack.append([])
        elif t == ")":
            l = stack.pop()
            stack[-1].append(l)
        else:
            try:
                stack[-1].append(int(t))
            except ValueError:
                stack[-1].append(t)
    return stack[0]

def parse(sexpr):
    """Parse symbolic expressions into a more structured abstract syntax tree."""
    # FIXME This grammar is not correct. We are likely to return an AST
    # that is invalid and does not make sense. This will then be caught by
    # the type checker or evaluator. But it may be better to report the problem
    # here.
    if isinstance(sexpr, list):
        if sexpr[0] == "+":
            return AstAdd(parse(sexpr[1]), parse(sexpr[2]))
        elif sexpr[0] == "if":
            return AstIf(parse(sexpr[1]), parse(sexpr[2]), parse(sexpr[3]))
        elif sexpr[0] == "fn":
            return AstFn(parse(sexpr[1]), parse(sexpr[2]), parse(sexpr[3]))
        elif sexpr[0] == "tfn":
            return TypeFn(parse(sexpr[1]), parse(sexpr[2]), parse(sexpr[3]))
        elif sexpr[0] == "let":
            return AstLet(parse(sexpr[1]), parse(sexpr[2]), parse(sexpr[3]))
        elif sexpr[0] == "rec":
            return AstRec(parse(sexpr[1]), parse(sexpr[2]), parse(sexpr[3]))
        else:
            return AstApp(parse(sexpr[0]), parse(sexpr[1]))
    elif isinstance(sexpr, str):
        if sexpr == "true":
            return AstTrue()
        elif sexpr == "false":
            return AstFalse()
        elif sexpr == "int":
            return int
        elif sexpr == "bool":
            return bool
        else:
            return AstVar(sexpr)
    elif isinstance(sexpr, int):
        return AstNum(sexpr)
    else:
        raise TypeError(sexpr)

def readexprs(txt):
    """Parse text into a structured abstract syntax tree via conversion to
    symbolic expressions."""
    return [parse(sexpr) for sexpr in readsexpr(txt)]

def selftest():
    """Run a basic self test using provided testcases."""
    a, b, n, x = AstVar('a'), AstVar('b'), AstVar('n'), AstVar('x')
    testcases = [(AstNum(32), int),
            (AstAdd(AstNum(1), AstNum(1)), int),
            (AstIf(AstTrue(), AstNum(1), AstNum(2)), int),
            (AstFn(a, int, a), TypeFn(a, int, int)),
            (AstApp(AstFn(a, int, a),AstNum(42)), int),
            (AstLet(a, AstNum(1), a), int),
            (AstRec(AstNum(0),AstTrue(),AstFn(n, int, AstFn(b, bool, AstFalse()))), bool),
            (AstTrue(), bool),
            (AstFalse(), bool)]
    for (expression, t) in testcases:
        rt = expression.typeof({})
        rv = expression.eval({})
        if rt == None or rv == None:
            print(f"Failed: {e, rt, rv}")
            break
        elif rt != t:
            print(f"Type failed: {e, rt, t}")
            break

if __name__ == "__main__":
    import sys
    selftest()
    for e in readexprs(sys.stdin.read()):
        t = e.typeof({})
        print(f"# expr: {repr(e)}")
        print(f"# type: {repr(t)}")
        if t:
            print(repr(e.eval({})))