from __future__ import annotations
from collections import deque
from typing import TYPE_CHECKING

import cfg_build
import syntax
from cfg.CFG_Node import CFG_START

if TYPE_CHECKING:
    from cfg.CFG import CFG

GLOBAL_SCOPE = ""

# A scoped variable: the function it belongs to, and its name.
# The scope is GLOBAL_SCOPE ("") for variables outside any function.
# e.g. ("f", "x")  →  variable "x" defined in function "f"
#      ("",  "x")  →  variable "x" at global scope
Var = tuple[str, str]

class BackwardAnalysis:
    def __init__(self, cfg: "CFG") -> None:
        self.cfg = cfg
        self.uses: dict[int, set[Var]] = {}
        self.defs: dict[int, set[Var]] = {}

        self.__funcs: dict[str, tuple] = dict(cfg_build.FUNCTIONS)
        self.__func_parent, self._func_params = self.__collect_function_metadata()
        self.func_scope: dict[int, str] = self.__compute_function_scope()
        self.__extract_uses_and_defs()

    # Walk the AST and collect function-parent and parameter information.
    def __collect_function_metadata(self) -> tuple[dict[str, str | None], dict[str, tuple[str, ...]]]:
        func_parent: dict[str, str | None] = {}
        func_params: dict[str, tuple[str, ...]] = {}

        def visit(expr: syntax.EXPRESSION | None, current_func: str | None) -> None:
            if expr is None:
                return
            if isinstance(expr, syntax.LET):
                decls = expr.decl if isinstance(expr.decl, list) else [expr.decl]
                # Register metadata for each declared function.
                for d in decls:
                    if isinstance(d, syntax.DECL):
                        func_parent[d.f_name] = current_func
                        func_params[d.f_name] = tuple(d.params)
                # Recurse into function bodies and the in-expression.
                for d in decls:
                    if isinstance(d, syntax.DECL):
                        visit(d.body, d.f_name)
                    else:
                        visit(d, current_func)
                visit(expr.body, current_func)
                return
            for _, child in expr.children():
                visit(child, current_func)

        visit(self.cfg.ast, None)
        return func_parent, func_params

    # Calculates the scope (in which function is it?) of each node in the CFG.
    def __compute_function_scope(self) -> dict[int, str]:
        # The first function whose BFS claims a node wins.
        functions = self.__funcs
        func_scope: dict[int, str] = {}
        all_f_start_ids: set[int] = {fs.id for _, (fs, _) in functions.items()}

        for f_name, (f_start, f_end) in functions.items():
            queue: deque = deque([f_start])
            while queue:
                node = queue.popleft()
                if node.id in func_scope:
                    continue  # already claimed by an earlier function
                func_scope[node.id] = f_name
                # Stop here — do not follow into the caller context.
                if node.id == f_end.id:
                    continue
                for child in node.children:
                    # Do not follow into a different function's START.
                    if (
                        isinstance(child, CFG_START)
                        and child.id in all_f_start_ids
                        and child.id != f_start.id
                    ):
                        continue
                    queue.append(child)

        return func_scope

    # Populate uses and defs for every node in the CFG.
    def __extract_uses_and_defs(self) -> None:
        for node in self.cfg.nodes():
            nid = node.id
            func = self.func_scope.get(nid)
            ast = node.ast_node

            uses: set[Var] = set()
            defs: set[Var] = set()

            if isinstance(node, CFG_START) and isinstance(ast, syntax.DECL):
                # Function entry defines each formal parameter.
                for param in ast.params:
                    defs.add((ast.f_name, param))
            elif ast is not None:
                if isinstance(ast, syntax.ID):
                    resolved = self.resolve_var(func, ast.name)
                    uses.add(resolved)
                elif isinstance(ast, syntax.ASSIGN):
                    resolved = self.resolve_var(func, ast.var.name)
                    defs.add(resolved)

            self.uses[nid] = uses
            self.defs[nid] = defs

    # Resolve a variables name and scope by walking up the hierarchy
    def resolve_var(self, func: str | None, name: str) -> Var:
        if func is None:
            return GLOBAL_SCOPE, name

        cur: str | None = func
        seen: set[str] = set()
        while cur is not None and cur not in seen:
            seen.add(cur)
            if name in self._func_params.get(cur, ()):
                return cur, name
            cur = self.__func_parent.get(cur)

        # Fallback: local variable in the current function scope
        return func, name