"""
transpiler/analyzer.py — Semantic Analysis Pass for IR (v0.5.0).
"""

from __future__ import annotations
import json
import os
from typing import Dict, Optional, Set, Any
from transpiler.middleend import ir
from transpiler.infrastructure.diagnostics import manager, DiagnosticManager
from transpiler.infrastructure.project_context import ProjectContext, ModuleExport

class SymbolTable:
    def __init__(self, parent: Optional[SymbolTable] = None):
        self.parent = parent
        self.symbols: Dict[str, str] = {} # name -> type_name
        self.definitions: Dict[str, tuple[int, int, Optional[str]]] = {} # name -> (line, col, origin_module)

    def lookup(self, name: str) -> Optional[str]:
        if name in self.symbols:
            return self.symbols[name]
        p = self.parent
        if p is not None:
            return p.lookup(name)
        return None

    def lookup_definition(self, name: str) -> Optional[tuple[int, int, Optional[str]]]:
        if name in self.definitions:
            return self.definitions[name]
        p = self.parent
        if p is not None:
            return p.lookup_definition(name)
        return None

    def define(self, name: str, type_name: str, lineno: int = 0, col_offset: int = 0, origin: Optional[str] = None):
        self.symbols[name] = type_name
        self.definitions[name] = (lineno, col_offset, origin)

SDK_MODULES = {"roblox", "math", "table", "string", "task", "coroutine", "debug", "re", "json", "random"}

class SemanticAnalyzer:
    def __init__(self, diagnostics: Optional[DiagnosticManager] = None, project: Optional[ProjectContext] = None, module_name: str = "main", uri: Optional[str] = None):
        self.diagnostics = diagnostics or DiagnosticManager()
        self.project = project
        self.module_name = module_name
        self.uri = uri
        self._class_info: Dict[str, Any] = {}
        self._current_class: Optional[ir.IRClassDef] = None
        
        # Load core snippets if available
        self.metadata = self._load_metadata()
        self.symbol_table = SymbolTable()
        self._init_globals()
        self._exports = ModuleExport(module_name, uri)

    def _load_metadata(self) -> Dict[str, Any]:
        path = os.path.join(os.path.dirname(__file__), "api_metadata.json")
        try:
            with open(path, "r") as f:
                return json.load(f)
        except Exception:
            # Silence metadata loading warnings in LSP mode
            return {"classes": {}, "globals": {}}

    def _init_globals(self):
        globals_map = self.metadata.get("globals", {})
        for name, type_name in globals_map.items():
            self.symbol_table.define(name, type_name)

    def analyze(self, node: ir.IRNode):
        method_name = f"visit_{type(node).__name__}"
        visitor = getattr(self, method_name, self.generic_visit)
        visitor(node)
        
        # Ensure all expressions receive an inferred type after visiting
        if isinstance(node, ir.IRExpression) and getattr(node, "inferred_type", None) is None:
            node.inferred_type = self._infer_type(node)

    def generic_visit(self, node: ir.IRNode):
        if hasattr(node, "__dict__"):
            for attr, value in node.__dict__.items():
                if isinstance(value, ir.IRNode):
                    self.analyze(value)
                elif isinstance(value, list):
                    for item in value:
                        if isinstance(item, ir.IRNode):
                            self.analyze(item)
                elif isinstance(value, ir.IRIfElseExpression):
                    self.analyze(value.condition)
                    self.analyze(value.then_expr)
                    self.analyze(value.else_expr)

    def visit_IRModule(self, node: ir.IRModule):
        # We'll store class property info in class_properties map or similar
        self._class_info = {}

        for stmt in node.body:
            self.analyze(stmt)
        
        # After analysis, if we are in a project, update the global project context
        if self._exports:
            self._exports.symbols = self.symbol_table.symbols.copy()
            self._exports.classes = self._class_info.copy()
            
        proj = self.project
        if proj is not None:
            proj.modules[self.module_name] = self._exports

        # Capture metadata for the generator
        node.metadata = {
            "symbols": self.symbol_table.symbols,
            "classes": self._class_info
        }

    def visit_IRAssignment(self, node: ir.IRAssignment):
        self.analyze(node.value)
        inferred = self._infer_type(node.value)
        if inferred and not node.value.inferred_type:
            node.value.inferred_type = inferred
        
        if isinstance(node.target, ir.IRVariable):
            name = getattr(node.target, "name", None)
            if name:
                t = self._infer_type(node.value) or "any"
                self.symbol_table.define(name, t, node.target.lineno, node.target.col_offset)
                
                # Update inferred_type on target node itself
                node.target.inferred_type = t
                
                # Export if we are at module level
                if self.symbol_table.parent is None: # Toplevel
                    self._exports.symbols[node.target.name] = inferred
                    self._exports.definitions[node.target.name] = (node.target.lineno, node.target.col_offset)
        elif isinstance(node.target, ir.IRPropertyAccess):
            # Check for 'self.prop = value'
            if isinstance(node.target.receiver, ir.IRVariable) and node.target.receiver.name == "self":
                curr_class = getattr(self, "_current_class", None)
                if curr_class:
                    curr_class.properties[node.target.property] = inferred or "any"
                    node.target.inferred_type = inferred

        self.analyze(node.target)
        
        # Type check: if target already has a type, ensure value is compatible
        if isinstance(node.target, ir.IRVariable):
            target_type = self.symbol_table.lookup(node.target.name)
            if target_type and target_type != "any" and inferred and inferred != "any":
                if target_type != inferred:
                    self.diagnostics.error(
                        f"Type mismatch: cannot assign '{inferred}' to variable of type '{target_type}'",
                        node.lineno, node.col_offset, self.uri
                    )

    def visit_IRBinaryOperation(self, node: ir.IRBinaryOperation):
        self.analyze(node.left)
        self.analyze(node.right)
        
        lt = node.left.inferred_type or "any"
        rt = node.right.inferred_type or "any"
        op = node.operator
        
        # Mapping AST names to human-readable symbols for error messages
        OP_DISPLAY = {
            "Add": "+", "Sub": "-", "Mult": "*", "Div": "/", "Mod": "%", "Pow": "**", "FloorDiv": "//",
            "Lt": "<", "LtE": "<=", "Gt": ">", "GtE": ">=", "Eq": "==", "NotEq": "!=",
            "And": "and", "Or": "or"
        }
        display_op = OP_DISPLAY.get(op, op)
        
        # Arithmetic operators: both must be numbers
        if op in ("Add", "Sub", "Mult", "Div", "Mod", "Pow", "FloorDiv"):
            # Attempt to find the operator position (e.g., after the left operand)
            # This is a heuristic: left_end + 1
            op_col = node.col_offset
            if hasattr(node.left, "name"):
                op_col = node.left.col_offset + len(node.left.name) + 1 # +1 for space
            elif isinstance(node.left, ir.IRLiteral):
                op_col = node.left.col_offset + len(str(node.left.value)) + 1
            
            if lt != "any" and lt != "number":
                self.diagnostics.error(f"Operator '{display_op}' is not supported for type '{lt}'", node.lineno, op_col, self.uri)
            if rt != "any" and rt != "number":
                self.diagnostics.error(f"Operator '{display_op}' is not supported for type '{rt}'", node.lineno, op_col, self.uri)
        
        # Comparisons: usually need same types
        elif op in ("Lt", "LtE", "Gt", "GtE", "Eq", "NotEq"):
            if lt != "any" and rt != "any" and lt != rt:
                self.diagnostics.warning(f"Comparison between mismatched types '{lt}' and '{rt}'", node.lineno, node.col_offset, self.uri)

    def _track_usage(self, name: str, node: ir.IRNode):
        # sys.stderr.write(f"DEBUG: _track_usage called for '{name}' in '{self.module_name}'\n")
        if self.project and self.module_name:
            defn = self.symbol_table.lookup_definition(name)
            target_mod = self.module_name
            if defn and defn[2]: # origin is third in the tuple
                target_mod = defn[2]
            
            mod = self.project.get_module(target_mod)
            if name not in mod.usages:
                mod.usages[name] = []
            mod.usages[name].append((self.uri or "unknown", node.lineno, node.col_offset))

    def visit_IRVariable(self, node: ir.IRVariable):
        # Resolve symbol in symbol table
        symbol = self.symbol_table.lookup(node.name)
        if symbol:
            node.inferred_type = symbol
        # Record usage
        self._track_usage(node.name, node)
            
        # Record definition location for LSP
        defn = self.symbol_table.lookup_definition(node.name)
        if defn:
            node.definition_location = defn # (line, col, origin)

    def visit_IRFunctionCall(self, node: ir.IRFunctionCall):
        self.analyze(node.func)
        for arg in node.args:
            self.analyze(arg)
        
        # Handle SDK member calls (e.g., from string import upper -> upper("..."))
        if isinstance(node.func, ir.IRVariable) and node.func.inferred_type:
            if node.func.inferred_type.startswith("SDK_MEMBER:"):
                qualified = node.func.inferred_type[11:]
                if "." in qualified:
                    mod, member = qualified.split(".", 1)
                    # Transform this function call into a method call (qualified library call)
                    # We can't easily replace 'node' itself in a visitor, but we can 
                    # change its 'func' to an Attribute if we had that in IR.
                    # In RPy IR, we have IRPropertyAccess.
                    # Let's change node.func to an IRPropertyAccess.
                    node.func = ir.IRPropertyAccess(
                        receiver=ir.IRVariable(name=mod, lineno=node.lineno, col_offset=node.col_offset),
                        property=member,
                        lineno=node.lineno,
                        col_offset=node.col_offset
                    )

        if not node.inferred_type:
            node.inferred_type = self._infer_type(node)

    def visit_IRMethodCall(self, node: ir.IRMethodCall):
        self.analyze(node.receiver)
        for arg in node.args:
            self.analyze(arg)

        receiver_type = self._infer_type(node.receiver)
        node.call_style = ir.CallStyle.COLON # Default

        # 1. Aggressively map common Roblox patterns regardless of type info
        # as these are almost always intended to be the CamelCase Roblox versions.
        common_rbx_patterns = {
            "connect": "Connect",
            "wait": "Wait",
            "destroy": "Destroy",
            "clone": "Clone",
            "is_a": "IsA",
            "get_service": "GetService",
            "wait_for_child": "WaitForChild",
            "get_children": "GetChildren",
            "find_first_child": "FindFirstChild",
            "set_core_gui_enabled": "SetCoreGuiEnabled"
        }
        if node.method in common_rbx_patterns:
            node.method = common_rbx_patterns[node.method]
        elif "_" in node.method:
            # General snake_case to CamelCase fallback for any colon call
            # This is idiomatic for Python-to-Roblox transpilation.
            node.method = "".join(p.capitalize() for p in node.method.split("_"))

        if receiver_type:
            if receiver_type.startswith("Module:"):
                node.call_style = ir.CallStyle.DOT # Modules use dots
                # Infer return type from module exports
                mod_name = receiver_type[7:]
                proj = self.project
                if proj is not None:
                    mod_export = proj.modules.get(mod_name)
                    if mod_export:
                        t = mod_export.symbols.get(node.method)
                        if t and t.startswith("Function:"):
                            node.inferred_type = t[9:]
                return

            is_optional = receiver_type.endswith("?")
            base_type = receiver_type[:-1] if is_optional else receiver_type
            class_name = base_type[6:] if base_type.startswith("Class:") else base_type

            class_info = self._class_info.get(class_name)
            if class_info:
                # 1. Direct match
                method_type = class_info.get("methods", {}).get(node.method)
                
                # 2. Case-insensitive match for Roblox APIs (if metadata available)
                if not method_type:
                    search_name = node.method.lower()
                    for m_name in class_info.get("methods", {}):
                        if m_name.lower() == search_name:
                            node.method = m_name # Remap to CamelCase
                            method_type = class_info["methods"][m_name]
                            break
                
                if method_type in ("classmethod", "staticmethod"):
                    node.call_style = ir.CallStyle.DOT
                else:
                    node.call_style = ir.CallStyle.COLON
                return

    def visit_IRImport(self, node: ir.IRImport):
        if node.module in SDK_MODULES:
            # e.g., import math -> math is a module
            self.symbol_table.define(node.alias or node.module, f"Module:{node.module}")
            return

        if self.project:
            # We don't perform deep analysis here, just define the module name in symbol table
            self.symbol_table.define(node.alias or node.module, f"Module:{node.module}")

    def visit_IRImportFrom(self, node: ir.IRImportFrom):
        if node.module == "roblox":
            for i, name in enumerate(node.names):
                alias = node.aliases[i] or name
                # e.g., from roblox import task -> task is a module/global
                if name in SDK_MODULES:
                    self.symbol_table.define(alias, f"Module:{name}")
                else:
                    self.symbol_table.define(alias, "any")
            return

        if node.module in SDK_MODULES:
            for i, name in enumerate(node.names):
                alias = node.aliases[i] or name
                # e.g., from math import floor -> floor is a member of math
                self.symbol_table.define(alias, f"SDK_MEMBER:{node.module}.{name}")
            return

        if self.project is not None:
            mod_export = self.project.modules.get(node.module)
            if mod_export:
                for i, name in enumerate(node.names):
                    alias = node.aliases[i] or name
                    t = mod_export.symbols.get(name, "any")
                    self.symbol_table.define(alias, t, origin=node.module)
                    # Track usage of the imported symbol in its source module
                    self._track_usage(alias, node)
            else:
                proj = self.project
                if proj is not None:
                    self.diagnostics.hint(
                        f"Module '{node.module}' not found. (Root: {proj.workspace_root or 'None'})",
                        line=node.lineno, col=node.col_offset
                    )

    def visit_IRPropertyAccess(self, node: ir.IRPropertyAccess):
        self.analyze(node.receiver)
        receiver_type = self._infer_type(node.receiver)
        
        if not receiver_type or receiver_type == "any":
            return

        # --- Module Access Fallback ---
        if receiver_type.startswith("Module:"):
            mod_name = receiver_type[7:]
            proj = self.project
            if proj is not None:
                mod_export = proj.modules.get(mod_name)
                if mod_export:
                    t = mod_export.symbols.get(node.property)
                    if t:
                        node.inferred_type = t
                        return
            
            if mod_name in SDK_MODULES:
                # For SDK modules, we assume the property exists for now
                node.inferred_type = "any"
                return

            self.diagnostics.error(
                f"Module '{mod_name}' has no attribute '{node.property}'.",
                line=node.lineno, col=node.col_offset
            )
            return

        is_optional = receiver_type.endswith("?")
        base_type = receiver_type[:-1] if is_optional else receiver_type
        
        if is_optional:
            receiver_name = getattr(node.receiver, "name", "object") if hasattr(node.receiver, "name") else "object"
            self.diagnostics.warning(
                f"Property access '{node.property}' on optional type '{receiver_type}'.",
                line=node.lineno, col=node.col_offset,
                hint=f"Consider checking 'if {receiver_name} then' first."
            )

        class_info = self._get_class_info(base_type)
        if class_info:
            props = class_info.get("properties", {})
            prop_type = props.get(node.property)
            if prop_type:
                node.inferred_type = prop_type
                return
            
            # --- Service / Child Fallback for Instances ---
            is_instance = base_type in ("Instance", "DataModel")
            if not is_instance:
                # Check inheritance
                curr = base_type
                while curr:
                    info = self._get_class_info(curr)
                    if info and info.get("inherits") in ("Instance", "DataModel"):
                        is_instance = True
                        break
                    curr = info.get("inherits") if info else None

            if is_instance:
                # 1. Try case-insensitive match for properties
                search_name = node.property.replace("_", "").lower()
                for p_name in props:
                    if p_name.lower() == search_name:
                        node.property = p_name
                        node.inferred_type = props[p_name]
                        return
                
                # 2. Assume it's a child or dynamic property
                node.inferred_type = "any"
                return

            # Check signals
            signals = class_info.get("signals", {})
            signal_type = signals.get(node.property)
            if signal_type:
                node.inferred_type = signal_type
                # Track usage of the signal
                self._track_usage(node.property, node)
                return
            
            # Record usage of the property/method
            self._track_usage(node.property, node)
            
            # Not found: Fuzzy match
            from transpiler.frontend.ast_utils import get_closest_match
            all_possibilities = list(props.keys()) + list(signals.keys()) + list(class_info.get("methods", {}).keys())
            closest = get_closest_match(node.property, all_possibilities)
            
            self.diagnostics.error(
                f"Property '{node.property}' does not exist on type '{base_type}'.",
                line=node.lineno, col=node.col_offset,
                hint=f"Did you mean '{closest}'?" if closest else f"Verify the member exists on {base_type}."
            )

    def visit_IRFunctionDef(self, node: ir.IRFunctionDef):
        # Define the function itself in the parent symbol table first
        if node.name:
            # We use a special prefix for function return types in symbols or just store the type
            # For simplicity let's store it as "ReturnType(className)" or similar
            # Actually, let's just use the function name to store its return type for now.
            pass

        old_table = self.symbol_table
        self.symbol_table = SymbolTable(parent=old_table)
        
        # Define arguments with types if available
        for i, arg_name in enumerate(node.args):
            t_ann = node.arg_types[i] if i < len(node.arg_types) else None
            # Arguments are defined at the function start
            # For now we use the function line/col, in a better IR we'd have arg locs
            self.symbol_table.define(arg_name, t_ann or "any", node.lineno, node.col_offset)
        
        for stmt in node.body:
            self.analyze(stmt)
            
        # Optional: Attempt to infer return type if not explicitly set
        if not node.return_type:
            # Simple heuristic: find all return statements
            returns = []
            def find_returns(n):
                if isinstance(n, ir.IRReturnStatement):
                    returns.append(n)
                elif hasattr(n, "__dict__"):
                    for attr, v in n.__dict__.items():
                        if isinstance(v, list):
                            for item in v: find_returns(item)
                        elif isinstance(v, ir.IRNode): find_returns(v)
            
            find_returns(node)
            if returns:
                # Just take the first one for now
                ret_type = self._infer_type(returns[0].value) if returns[0].value else "nil"
                node.return_type = ret_type
            else:
                node.return_type = "any"

        self.symbol_table = old_table
        name = getattr(node, "name", None)
        if name:
            self.symbol_table.define(name, f"Function:{node.return_type}", node.lineno, node.col_offset)
            # Export if we are at module level
            if self.symbol_table.parent is None:
                self._exports.symbols[name] = f"Function:{node.return_type}"
                self._exports.definitions[name] = (node.lineno, node.col_offset)

    def visit_IRGenericForStatement(self, node: ir.IRGenericForStatement):
        self.analyze(node.iterator)
        
        # Heuristic for iterator types
        if not node.var_types:
            it_type = self._infer_type(node.iterator)
            if it_type and it_type.startswith("Class:"):
                class_name = it_type[6:]
                class_info = self._get_class_info(class_name)
                if class_info and "__iter__" in class_info.get("methods", {}):
                    # For now, assume it returns any. In a full impl, we'd check __iter__'s return type.
                    node.var_types = ["any"]
            
            if isinstance(node.iterator, ir.IRFunctionCall) and isinstance(node.iterator.func, ir.IRVariable):
                it_name = node.iterator.func.name
                if it_name == "ipairs":
                    node.var_types = ["number", "any"]
                elif it_name == "pairs":
                    node.var_types = ["any", "any"]
        
        old_table = self.symbol_table
        self.symbol_table = SymbolTable(parent=old_table)
        for i, v in enumerate(node.vars):
            v_type = node.var_types[i] if i < len(node.var_types) else "any"
            self.symbol_table.define(v, v_type, node.lineno, node.col_offset)
        
        for stmt in node.body:
            self.analyze(stmt)
        self.symbol_table = old_table

    def visit_IRWhileStatement(self, node: ir.IRWhileStatement):
        self.analyze(node.condition)
        for stmt in node.body:
            self.analyze(stmt)

    def visit_IRForStatement(self, node: ir.IRForStatement):
        self.analyze(node.start)
        self.analyze(node.end)
        if node.step: self.analyze(node.step)
        
        old_table = self.symbol_table
        self.symbol_table = SymbolTable(parent=old_table)
        self.symbol_table.define(node.var, "number", node.lineno, node.col_offset)
        # S5: Mark loop variable as certain 'number'
        node.var_type_info = ir.TypeMetadata(luau_type="number", confidence="certain")
        
        for stmt in node.body:
            self.analyze(stmt)
        self.symbol_table = old_table

    def _extract_narrowings(self, condition: ir.IRExpression) -> Dict[str, str]:
        """
        Return a dict of {var_name: narrowed_type} that holds in the then-branch
        of an if-statement given `condition`.

        Supported patterns:
        - `x`                     → x: non-optional (strip `?`)
        - `x is None`             → x: "nil"
        - `x is not None`         → x: base type (strip `?`)
        - `isinstance(x, T)`      → x: T  (where T is an IRVariable naming a type)
        - `cond1 and cond2`       → merge both dicts
        - `not cond`              → invert single-var narrowings
        """
        narrowings: Dict[str, str] = {}

        # Pattern 1: bare variable  →  strip optional
        if isinstance(condition, ir.IRVariable):
            name = getattr(condition, "name", None)
            if name:
                existing = self.symbol_table.lookup(name)
                if existing and existing.endswith("?"):
                    narrowings[name] = existing[:-1]
        
        # Pattern 2: `not x`  →  variable is nil/None
        elif isinstance(condition, ir.IRUnaryOperation) and condition.operator in ("not ", "not"):
            if isinstance(condition.operand, ir.IRVariable):
                name = getattr(condition.operand, "name", None)
                if name:
                    existing = self.symbol_table.lookup(name)
                    if existing:
                        narrowings[name] = "nil"

        # Pattern 3: binary comparisons and logical combinators
        elif isinstance(condition, ir.IRBinaryOperation):
            op = condition.operator

            # `x is None` or `x == None`
            if op in ("Is", "Eq", "is", "=="):
                lhs, rhs = condition.left, condition.right
                # normalise order
                if isinstance(rhs, ir.IRVariable) and isinstance(lhs, ir.IRLiteral) and lhs.value is None:
                    lhs, rhs = rhs, lhs
                if isinstance(lhs, ir.IRVariable) and isinstance(rhs, ir.IRLiteral) and rhs.value is None:
                    name = getattr(lhs, "name", None)
                    if name:
                        narrowings[name] = "nil"

            # `x is not None` / `x != None`
            elif op in ("IsNot", "NotEq", "is not", "!="):
                lhs, rhs = condition.left, condition.right
                if isinstance(rhs, ir.IRVariable) and isinstance(lhs, ir.IRLiteral) and lhs.value is None:
                    lhs, rhs = rhs, lhs
                if isinstance(lhs, ir.IRVariable) and isinstance(rhs, ir.IRLiteral) and rhs.value is None:
                    name = getattr(lhs, "name", None)
                    if name:
                        existing = self.symbol_table.lookup(name)
                        if existing and existing.endswith("?"):
                            narrowings[name] = existing[:-1]  # x: T? → T

            # `cond and cond`  →  merge
            elif op in ("And", "and"):
                narrowings.update(self._extract_narrowings(condition.left))
                narrowings.update(self._extract_narrowings(condition.right))

        # Pattern 4: isinstance(x, T)
        elif isinstance(condition, ir.IRFunctionCall):
            if isinstance(condition.func, ir.IRVariable) and getattr(condition.func, "name", None) == "isinstance":
                if len(condition.args) == 2:
                    target = condition.args[0]
                    type_arg = condition.args[1]
                    if isinstance(target, ir.IRVariable) and isinstance(type_arg, ir.IRVariable):
                        # Map Python primitive names to Luau type strings
                        TYPE_MAP = {
                            "int": "number", "float": "number", "str": "string",
                            "bool": "boolean", "list": "{any}", "dict": "{[any]: any}",
                        }
                        t_name = getattr(target, "name", None)
                        ta_name = getattr(type_arg, "name", None)
                        if t_name and ta_name:
                            narrowings[t_name] = TYPE_MAP.get(ta_name, ta_name)

        return narrowings

    def _invert_narrowings(self, then_narrowings: Dict[Any, str]) -> Dict[str, str]:
        """
        Compute else-branch narrowings by inverting then-branch ones.
        If then-branch: x → T      else-branch: x → T? (re-add optional)
        If then-branch: x → nil    else-branch: x → existing base type
        """
        inverted: Dict[str, str] = {}
        for name, narrowed_type in then_narrowings.items():
            if not isinstance(name, str):
                continue
            existing = self.symbol_table.lookup(name)
            if narrowed_type == "nil":
                # then: x is nil  →  else: x is not nil → base type
                if existing and existing.endswith("?"):
                    inverted[name] = existing[:-1]
            elif existing and existing.endswith("?"):
                # then: x is not nil  →  else: x might still be nil
                inverted[name] = existing  # keep optional
            # isinstance narrowings don't invert cleanly → skip (keep outer type)
        return inverted

    def visit_IRIfStatement(self, node: ir.IRIfStatement):
        self.analyze(node.condition)

        # --- Flow-sensitive type narrowing ---
        then_narrowings = self._extract_narrowings(node.condition)
        else_narrowings = self._invert_narrowings(then_narrowings)

        # Collect original types for restoration
        saved_types: Dict[str, str] = {}
        for name in then_narrowings:
            t = self.symbol_table.lookup(name)
            if t is not None:
                saved_types[name] = t

        # Apply then-branch narrowings
        for name, narrowed in then_narrowings.items():
            self.symbol_table.define(name, narrowed)

        for stmt in node.then_block:
            self.analyze(stmt)

        # Restore original types before the else-branch
        for name, orig in saved_types.items():
            self.symbol_table.define(name, orig)

        # Apply else-branch narrowings (if any)
        for name, narrowed in else_narrowings.items():
            self.symbol_table.define(name, narrowed)

        if node.else_block:
            for stmt in node.else_block:
                self.analyze(stmt)

        # Restore after the full if/else block
        for name, orig in saved_types.items():
            self.symbol_table.define(name, orig)


    def visit_IRClassDef(self, node: ir.IRClassDef):
        self.symbol_table.define(node.name, f"Class:{node.name}", node.lineno, node.col_offset)
        
        # Track class info for metadata
        props = []
        meths = {}
        for m in node.body:
            if isinstance(m, ir.IRFunctionDef):
                func_node: ir.IRFunctionDef = m
                if "property" in func_node.decorators:
                    props.append(func_node.name)
                if func_node.name != "__init__":
                    mtype = "method"
                    if "classmethod" in func_node.decorators: mtype = "classmethod"
                    elif "staticmethod" in func_node.decorators: mtype = "staticmethod"
                    meths[func_node.name] = mtype
                else:
                    # Scan __init__ for self.xxx = ...
                    for stmt in func_node.body:
                        if isinstance(stmt, ir.IRAssignment) and isinstance(stmt.target, ir.IRPropertyAccess):
                            if isinstance(stmt.target.receiver, ir.IRVariable) and stmt.target.receiver.name == "self":
                                if stmt.target.property not in props:
                                    props.append(stmt.target.property)

        self._class_info[node.name] = {
            "bases": [str(b) for b in node.bases],
            "properties": props,
            "methods": meths,
            "lineno": node.lineno,
            "col_offset": node.col_offset
        }
        
        if self.symbol_table.parent is None:
            self._exports.symbols[node.name] = f"Class:{node.name}"
            self._exports.definitions[node.name] = (node.lineno, node.col_offset)

        old_class = self._current_class
        self._current_class = node
        
        # We might want a dedicated metadata entry for user classes
        # ... (removed redundancy)
        
        for stmt in node.body:
            self.analyze(stmt)
            
        self._current_class = old_class

    def _infer_type_metadata(self, node: ir.IRExpression) -> Optional[ir.TypeMetadata]:
        if node.type_info: return node.type_info
        
        if isinstance(node, ir.IRLiteral):
            if isinstance(node.value, str): return ir.TypeMetadata("string", "certain")
            if isinstance(node.value, bool): return ir.TypeMetadata("boolean", "certain")
            if isinstance(node.value, (int, float)): return ir.TypeMetadata("number", "certain")
            if node.value is None: return ir.TypeMetadata("nil", "certain")

        if isinstance(node, ir.IRBinaryOperation):
            left_meta = self._infer_type_metadata(node.left)
            right_meta = self._infer_type_metadata(node.right)
            if left_meta and right_meta and left_meta.luau_type == "number" and right_meta.luau_type == "number":
                if node.operator in ("+", "-", "*", "/", "%", "**", "//"):
                    conf = "certain" if (left_meta.confidence == "certain" and right_meta.confidence == "certain") else "inferred"
                    return ir.TypeMetadata("number", conf)
            return ir.TypeMetadata("any", "unknown")

        if isinstance(node, ir.IRIfElseExpression):
            cond_meta = self._infer_type_metadata(node.condition)
            then_meta = self._infer_type_metadata(node.then_expr)
            else_meta = self._infer_type_metadata(node.else_expr)
            
            if then_meta and else_meta and then_meta.luau_type == else_meta.luau_type:
                conf = "certain" if (then_meta.confidence == "certain" and else_meta.confidence == "certain") else "inferred"
                return ir.TypeMetadata(then_meta.luau_type, conf)
            return ir.TypeMetadata("any", "inferred")

        if isinstance(node, ir.IRVariable):
            t = self.symbol_table.lookup(node.name)
            if t: return ir.TypeMetadata(t, "inferred")

        return ir.TypeMetadata("any", "unknown")

    def _infer_type(self, node: ir.IRExpression) -> Optional[str]:
        meta = self._infer_type_metadata(node)
        if meta and meta.luau_type != "any":
            return meta.luau_type
        
        # Analyze specific function/method/property signatures to lift "any"
        if isinstance(node, ir.IRMethodCall):
            if node.method in ("new", "Angles", "fromEulerAnglesYXZ", "fromEulerAnglesXYZ", "fromMatrix", "fromRGB", "fromHSV"):
                if isinstance(node.receiver, ir.IRVariable):
                    constructor_name = node.receiver.name
                    if constructor_name in ("Vector3", "Vector2", "CFrame", "UDim2", "UDim", "Color3", "Rect"):
                        return constructor_name
        elif isinstance(node, ir.IRFunctionCall):
                name = getattr(node.func, "name", None)
                if name:
                    if name in ("Vector3", "Vector2", "CFrame", "UDim2", "UDim", "Color3", "Rect"):
                        return name
                    
                    # Check for class constructor
                    t = self.symbol_table.lookup(name)
                    if t and t.startswith("Class:"):
                        return t
                    
        elif isinstance(node, ir.IRPropertyAccess):
            if node.property in ("Position", "Size", "LookVector", "RightVector", "UpVector"):
                return "Vector3"
            if node.property == "CFrame":
                return "CFrame"
            if node.property == "Color":
                return "Color3"
            if node.property in ("X", "Y", "Z"):
                return "number"

        return meta.luau_type if meta else None

    def _get_class_info(self, class_name: str) -> Optional[Dict[str, Any]]:
        # 1. Check user-defined classes first
        if class_name.startswith("Class:"):
            target = class_name[6:]
            return self._class_info.get(target)

        # 2. Check SDK/Built-in classes
        classes = self.metadata.get("classes", {})
        info = classes.get(class_name)
        if not info:
            return self._class_info.get(class_name) # Fallback
        
        # Merge with base classes if 'inherits' is present
        base_name = info.get("inherits")
        if base_name:
            base_info = self._get_class_info(base_name)
            if base_info:
                # Proper deep merge for properties, methods, signals
                merged = base_info.copy()
                for key in ["properties", "methods", "signals"]:
                    # Handle both list and dict styles
                    if isinstance(info.get(key), list) and isinstance(base_info.get(key), list):
                        merged[key] = list(set(base_info.get(key, []) + info.get(key, [])))
                    else:
                        merged[key] = base_info.get(key, {}).copy()
                        merged[key].update(info.get(key, {}))
                return merged
        return info

def analyze(module: ir.IRModule, diagnostics: Optional[DiagnosticManager] = None, project: Optional[ProjectContext] = None, module_name: str = "main"):
    analyzer = SemanticAnalyzer(diagnostics, project, module_name)
    analyzer.analyze(module)
