import ast

class TypeErrorInfo:
    def __init__(self, lineno, msg):
        self.lineno = lineno
        self.msg = msg

    def __str__(self):
        return f"Line {self.lineno}: {self.msg}"


class TypeChecker(ast.NodeVisitor):
    def __init__(self):
        self.errors = []
        self.env = {} 

    # ---------- utils ----------

    def error(self, node, msg):
        self.errors.append(TypeErrorInfo(node.lineno, msg))

    def annotation_to_type(self, ann):
        if isinstance(ann, ast.Name):
            return ann.id
        return None

    def literal_type(self, node):
        if isinstance(node, ast.Constant):
            return type(node.value).__name__
        return None

    # ---------- visit ----------

    def visit_AnnAssign(self, node: ast.AnnAssign):
        if not isinstance(node.target, ast.Name):
            return

        var_name = node.target.id
        ann_type = self.annotation_to_type(node.annotation)
        self.env[var_name] = ann_type

        if node.value:
            value_type = self.literal_type(node.value)
            if ann_type and value_type and ann_type != value_type:
                self.error(
                    node,
                    f"Cannot assign {value_type} to {ann_type}"
                )

    def visit_Assign(self, node: ast.Assign):
        value_type = self.literal_type(node.value)
        for target in node.targets:
            if isinstance(target, ast.Name):
                var_name = target.id
                if var_name in self.env:
                    expected = self.env[var_name]
                    if value_type and expected != value_type:
                        self.error(
                            node,
                            f"Cannot assign {value_type} to {expected}"
                        )

    def visit_FunctionDef(self, node: ast.FunctionDef):
        local_env = {}

        for arg in node.args.args:
            if arg.annotation:
                local_env[arg.arg] = self.annotation_to_type(arg.annotation)

        for stmt in node.body:
            if isinstance(stmt, ast.Return):
                if node.returns:
                    expected = self.annotation_to_type(node.returns)
                    actual = self.literal_type(stmt.value)
                    if expected and actual and expected != actual:
                        self.error(
                            stmt,
                            f"Return {actual}, expected {expected}"
                        )

        self.generic_visit(node)


# ---------- runner ----------

def check(code: str):
    tree = ast.parse(code)
    checker = TypeChecker()
    checker.visit(tree)

    for e in checker.errors:
        print(e)


# ---------- example ----------

if __name__ == "__main__":
    code = """
x: int = 1
y: int = "hello"

def foo(a: int) -> int:
    return "nope"
"""
    check(code)