import ast
from type_def import PrimitiveType
from parse_expr import parse_expr

class Lifetime:
    low: int
    original: int

    def __init__(self, scope):
        self.low = self.original = scope
        self.parent = None

    def fold(self):
        while self.parent is not None:
            self.low = self.parent.low
            self.original = self.parent.original
            self.parent = self.parent.parent
        return self

    def ok(self, other):
        self.fold()
        if other == None:
            return False
        other.fold()
        return self.low >= other.original and \
            (other.original != self.low or self.low != 1)

    def __str__(self):
        self.fold()
        return f'({self.low}, {self.original})'


def assign_expr(
        scope: int,
        sym_table: dict[str, Lifetime],
        expr: ast.expr):
    if isinstance(expr, ast.Expression):
        body = expr.body
    else:
        body = expr
    if isinstance(body.type, PrimitiveType):
        body.lifetime = None
    elif isinstance(body, ast.Attribute):
        body.lifetime = assign_expr(scope, sym_table, body.value)
    elif isinstance(body, ast.Subscript):
        body.lifetime = assign_expr(scope, sym_table, body.value)
    elif isinstance(body, ast.Name):
        if body.id in sym_table:
            body.lifetime = sym_table[body.id]
        else:
            body.lifetime = Lifetime(scope)
            sym_table[body.id] = body.lifetime
    else:
        body.lifetime = Lifetime(scope)
    return body.lifetime


def assign_stmt(
        scope: int,
        sym_table: dict[str, Lifetime],
        nodes):
    for node in nodes:
        if isinstance(node, ast.Assign):
            b = assign_expr(scope, sym_table, node.value)
            for target in node.targets:
                a = assign_expr(scope, sym_table, target)
                if a == None and b == None:
                    continue
                if not a.ok(b):
                    print(ast.unparse(node))
                    print(f'{a} <- {b}')
                    assert False
                a.low = min(a.low, b.low)
                a.original = max(a.original, b.original)
                b.parent = a
        elif isinstance(node, ast.If) or isinstance(node, ast.While):
            assign_stmt(scope + 1, sym_table, node.body)
            assign_stmt(scope + 1, sym_table, node.orelse)
        elif isinstance(node, ast.Return):
            a = assign_expr(scope, sym_table, node.value)
            if a != None and a.fold().original > 1:
                print(ast.unparse(node))
                print(a)
                assert False