"""
transpiler/ir_transformer.py — Stage 2: Python AST to IR conversion.
"""

from __future__ import annotations
import ast
from typing import Any, List, Optional
from transpiler.middleend import ir
from transpiler.infrastructure.errors import UnsupportedFeatureError
from transpiler.frontend.ast_utils import get_line, node_name

class PythonToIRTransformer(ast.NodeVisitor):
    def __init__(self):
        self._temp_count = 0

    def transform(self, node: ast.AST) -> ir.IRNode:
        return self.visit(node)

    def _transform_target(self, target: ast.AST | List[ast.AST]) -> List[ir.IRExpression]:
        if isinstance(target, list):
            res = []
            for t in target:
                res.extend(self._transform_target(t))
            return res
        if isinstance(target, ast.Name):
            return [ir.IRVariable(name=target.id, lineno=target.lineno)]
        if isinstance(target, ast.Attribute):
            return [ir.IRPropertyAccess(
                lineno=target.lineno,
                col_offset=target.col_offset,
                receiver=self.visit(target.value),
                property=target.attr
            )]
        if isinstance(target, ast.Subscript):
            return [ir.IRIndexOperation(
                lineno=target.lineno,
                col_offset=target.col_offset,
                receiver=self.visit(target.value),
                index=self.visit(target.slice)
            )]
        if isinstance(target, (ast.Tuple, ast.List)):
            res = []
            for elt in target.elts:
                res.extend(self._transform_target(elt))
            return res
        raise UnsupportedFeatureError(f"Unsupported assignment target: {type(target).__name__}", get_line(target))

    def _get_type_str(self, node: Optional[ast.AST]) -> Optional[str]:
        if node is None:
            return None
            
        if isinstance(node, ast.Name):
            primitive_map = {
                "int": "int",
                "float": "float",
                "str": "str",
                "bool": "bool",
                "Any": "Any",
                "None": "None"
            }
            return primitive_map.get(node.id, node.id)
            
        if isinstance(node, ast.Constant):
            if node.value is None:
                return "nil"
            if isinstance(node.value, str):
                return f'"{node.value}"'
            return str(node.value)
            
        if isinstance(node, ast.Attribute):
            val = self._get_type_str(node.value)
            if val == "typing":
                return node.attr
            return f"{val}.{node.attr}"
            
        if isinstance(node, ast.Subscript):
            container = self._get_type_str(node.value)
            if container == "List":
                el_type = self._get_type_str(node.slice)
                return f"{{{el_type}}}"
            elif container == "Dict":
                if isinstance(node.slice, ast.Tuple) and len(node.slice.elts) == 2:
                    k_type = self._get_type_str(node.slice.elts[0])
                    v_type = self._get_type_str(node.slice.elts[1])
                    return f"{{[{k_type}]: {v_type}}}"
                return "{[any]: any}"
            elif container == "Optional":
                el_type = self._get_type_str(node.slice)
                return f"{el_type}?"
            elif container == "Union":
                if isinstance(node.slice, ast.Tuple):
                    types = [str(self._get_type_str(e)) for e in node.slice.elts]
                    return " | ".join(types)
                else:
                    return self._get_type_str(node.slice)
            elif container == "Tuple":
                if isinstance(node.slice, ast.Tuple):
                    types = [str(self._get_type_str(e)) for e in node.slice.elts]
                    return f"({', '.join(types)})"
                elif isinstance(node.slice, ast.Name):
                    return f"({self._get_type_str(node.slice)})"
                return "(...any)"
            elif container == "Callable":
                if isinstance(node.slice, ast.Tuple) and len(node.slice.elts) == 2:
                    arg_types_node = node.slice.elts[0]
                    ret_type = self._get_type_str(node.slice.elts[1])
                    
                    if isinstance(arg_types_node, ast.List):
                        args = [str(self._get_type_str(e)) for e in arg_types_node.elts]
                        return f"({', '.join(args)}) -> {ret_type}"
                    elif isinstance(arg_types_node, ast.Constant) and arg_types_node.value is ...:
                        return f"(...any) -> {ret_type}"
                return "(...any) -> any"
            return container

        if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
            left = self._get_type_str(node.left)
            right = self._get_type_str(node.right)
            return f"{left} | {right}"
            
        return "any"

    def _visit_body(self, nodes: List[ast.stmt]) -> List[ir.IRStatement]:
        results = []
        for n in nodes:
            res = self.visit(n)
            if res is None: continue
            if isinstance(res, list):
                results.extend(res)
            else:
                results.append(res)
        return results

    def generic_visit(self, node: ast.AST) -> Any:
        raise UnsupportedFeatureError(
            node_name(node),
            line=getattr(node, "lineno", None),
            col=getattr(node, "col_offset", None)
        )

    def visit_Module(self, node: ast.Module) -> ir.IRModule:
        return ir.IRModule(lineno=1, body=self._visit_body(node.body))

    # --- Expressions ---

    def visit_Constant(self, node: ast.Constant) -> ir.IRLiteral:
        return ir.IRLiteral(
            lineno=getattr(node, "lineno", 0),
            col_offset=getattr(node, "col_offset", 0),
            value=node.value
        )

    def visit_Name(self, node: ast.Name) -> ir.IRVariable:
        return ir.IRVariable(
            lineno=getattr(node, "lineno", 0),
            col_offset=getattr(node, "col_offset", 0),
            name=node.id
        )

    def visit_BinOp(self, node: ast.BinOp) -> ir.IRBinaryOperation:
        return ir.IRBinaryOperation(
            lineno=node.lineno,
            col_offset=node.col_offset,
            left=self.visit(node.left),
            operator=type(node.op).__name__,
            right=self.visit(node.right)
        )

    def visit_Compare(self, node: ast.Compare) -> ir.IRExpression:
        current_left = self.visit(node.left)
        comparisons = []
        for op, comp in zip(node.ops, node.comparators):
            current_right = self.visit(comp)
            comparisons.append(ir.IRBinaryOperation(
                lineno=node.lineno,
                col_offset=node.col_offset,
                left=current_left,
                operator=type(op).__name__,
                right=current_right
            ))
            current_left = current_right
            
        res = comparisons[0]
        for comp in comparisons[1:]:
            res = ir.IRBinaryOperation(
                lineno=node.lineno,
                col_offset=node.col_offset,
                left=res,
                operator="And",
                right=comp
            )
        return res

    def visit_BoolOp(self, node: ast.BoolOp) -> ir.IRExpression:
        # Preserve 'and'/'or' as binary operations so the backend can emit 
        # native Luau 'and'/'or' operators, which follow the same 'return-the-value' 
        # logic as Python.
        values = [self.visit(v) for v in node.values]
        res = values[0]
        op_name = "And" if isinstance(node.op, ast.And) else "Or"
        for next_val in values[1:]:
            res = ir.IRBinaryOperation(
                lineno=node.lineno,
                col_offset=node.col_offset,
                left=res,
                operator=op_name,
                right=next_val
            )
        return res

    def visit_UnaryOp(self, node: ast.UnaryOp) -> ir.IRExpression:
        operand = self.visit(node.operand)
        
        # Constant folding for negative literals (e.g., -1)
        # This ensures the backend IRIndexOperation can detect negative indices.
        if isinstance(node.op, ast.USub) and isinstance(operand, ir.IRLiteral) and isinstance(operand.value, (int, float)):
            return ir.IRLiteral(
                lineno=node.lineno,
                col_offset=node.col_offset,
                value=-operand.value
            )
            
        return ir.IRUnaryOperation(
            lineno=node.lineno,
            col_offset=node.col_offset,
            operator=type(node.op).__name__,
            operand=operand
        )

    def visit_Call(self, node: ast.Call) -> ir.IRExpression:
        # Handle super().method(...)
        if isinstance(node.func, ast.Attribute) and \
           isinstance(node.func.value, ast.Call) and \
           isinstance(node.func.value.func, ast.Name) and \
           node.func.value.func.id == "super":
            return ir.IRSuperCall(
                lineno=node.lineno,
                col_offset=node.col_offset,
                method=node.func.attr,
                args=[self.visit(arg) for arg in node.args]
            )

        if isinstance(node.func, ast.Attribute):
            receiver_node = self.visit(node.func.value)
            method_name = node.func.attr

            # Detect receivers that must use dot-call style in Luau (not colon-call).
            # Python: CFrame.new(...)  ->  Luau: CFrame.new(...)  NOT CFrame:new(...)
            DOT_ONLY_RECEIVERS = {
                "CFrame", "Vector3", "Vector2", "UDim", "UDim2", "Rect", "Color3",
                "TweenInfo", "NumberSequence", "ColorSequence", "Ray", "Region3",
                "math", "table", "string", "os", "task", "Instance",
            }
            receiver_name = None
            if isinstance(node.func.value, ast.Name):
                receiver_name = node.func.value.id
            
            call_style = ir.CallStyle.DOT if receiver_name in DOT_ONLY_RECEIVERS else ir.CallStyle.UNKNOWN

            return ir.IRMethodCall(
                lineno=node.lineno,
                col_offset=node.col_offset,
                receiver=receiver_node,
                method=method_name,
                args=[self.visit(arg) for arg in node.args],
                call_style=call_style,
            )
        return ir.IRFunctionCall(
            lineno=node.lineno,
            col_offset=node.col_offset,
            func=self.visit(node.func),
            args=[self.visit(arg) for arg in node.args]
        )

    def visit_Attribute(self, node: ast.Attribute) -> ir.IRPropertyAccess:
        return ir.IRPropertyAccess(
            lineno=node.lineno,
            col_offset=node.col_offset,
            receiver=self.visit(node.value),
            property=node.attr
        )

    def visit_Subscript(self, node: ast.Subscript) -> ir.IRIndexOperation:
        return ir.IRIndexOperation(
            lineno=node.lineno,
            col_offset=node.col_offset,
            receiver=self.visit(node.value),
            index=self.visit(node.slice)
        )

    def visit_Slice(self, node: ast.Slice) -> ir.IRNode:
        # For now, just visit lower/upper
        # RPy usually doesn't support full slicing yet, but we need to visit it to avoid errors
        if node.lower: self.visit(node.lower)
        if node.upper: self.visit(node.upper)
        if node.step: self.visit(node.step)
        return ir.IRLiteral(value="slice")

    def visit_List(self, node: ast.List) -> ir.IRTableLiteral:
        return ir.IRTableLiteral(
            lineno=node.lineno,
            col_offset=node.col_offset,
            elements=[self.visit(e) for e in node.elts]
        )

    def visit_Tuple(self, node: ast.Tuple) -> ir.IRTableLiteral:
        return ir.IRTableLiteral(
            lineno=node.lineno,
            col_offset=node.col_offset,
            elements=[self.visit(e) for e in node.elts]
        )

    def visit_Dict(self, node: ast.Dict) -> ir.IRTableLiteral:
        fields = []
        for k, v in zip(node.keys, node.values):
            if k is None: continue
            kn = self.visit(k).value if isinstance(k, ast.Constant) else str(self.visit(k))
            fields.append((str(kn), self.visit(v)))
        return ir.IRTableLiteral(lineno=node.lineno, col_offset=node.col_offset, fields=fields)

    def visit_JoinedStr(self, node: ast.JoinedStr) -> ir.IRExpression:
        if not node.values:
            return ir.IRLiteral(value="", lineno=node.lineno, col_offset=node.col_offset)
        
        # Build a concatenation chain: (val1 .. (val2 .. val3))
        res = self.visit(node.values[-1])
        if isinstance(node.values[-1], ast.FormattedValue):
            res = ir.IRFunctionCall(
                func=ir.IRVariable(name="str"),
                args=[res],
                lineno=node.lineno,
                col_offset=node.col_offset
            )

        for i in range(len(node.values) - 2, -1, -1):
            val = self.visit(node.values[i])
            if isinstance(node.values[i], ast.FormattedValue):
                val = ir.IRFunctionCall(
                    func=ir.IRVariable(name="str"),
                    args=[val],
                    lineno=node.lineno,
                    col_offset=node.col_offset
                )
            
            res = ir.IRBinaryOperation(
                lineno=node.lineno,
                left=val,
                operator="Concat",
                right=res
            )
        return res

    def visit_FormattedValue(self, node: ast.FormattedValue) -> ir.IRExpression:
        return self.visit(node.value)

    def visit_Await(self, node: ast.Await) -> ir.IRAwait:
        return ir.IRAwait(
            lineno=node.lineno,
            col_offset=node.col_offset,
            value=self.visit(node.value)
        )

    # --- Statements ---

    def visit_Assign(self, node: ast.Assign) -> List[ir.IRAssignment]:
        # Capture all targets
        ir_targets = [self.visit(t) for t in node.targets]
        ir_value = self.visit(node.value)
        
        # Check if the assignment is wrapped in a persistent call: x = persistent(0)
        # Note: persistent only supports single target in this implementation
        if len(node.targets) == 1 and isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name) and node.value.func.id == "persistent":
            v = self.visit(node.value.args[0]) if node.value.args else ir.IRLiteral(value=None)
            return [ir.IRPersistentAssign(
                lineno=node.lineno,
                col_offset=node.col_offset,
                target=ir_targets[0],
                value=v
            )] # type: ignore
        
        
    def visit_IfExp(self, node: ast.IfExp) -> ir.IRIfElseExpression:
        return ir.IRIfElseExpression(
            lineno=node.lineno,
            condition=self.visit(node.test),
            then_expr=self.visit(node.body),
            else_expr=self.visit(node.orelse)
        )
        
    def visit_Global(self, node: ast.Global) -> None:
        # SSA and Generator handle this via scoped declaration tracking.
        return None

    def visit_Nonlocal(self, node: ast.Nonlocal) -> None:
        # SSA and Generator handle this via scoped declaration tracking.
        return None

    def visit_Assign(self, node: ast.Assign) -> List[ir.IRStatement]:
        # Handle assignments
        if len(node.targets) == 1:
            target_node = node.targets[0]
            ir_targets = self._transform_target(target_node)
            
            # Multi-target / Parallel assignment
            if len(ir_targets) > 1:
                # If RHS is a Tuple/List, it's parallel assignment
                if isinstance(node.value, (ast.Tuple, ast.List)):
                    values = [self.visit(elt) for elt in node.value.elts]
                    return [ir.IRMultiAssign(
                        lineno=node.lineno,
                        targets=ir_targets,
                        rhs_values=values,
                        is_local=all(isinstance(t, ir.IRVariable) for t in ir_targets)
                    )]
                
                # If RHS is a call, it's a multi-return unpack
                v = self.visit(node.value)
                if isinstance(v, (ir.IRFunctionCall, ir.IRMethodCall)):
                    return [ir.IRMultiAssign(
                        lineno=node.lineno,
                        targets=ir_targets,
                        rhs_call=v,
                        is_local=all(isinstance(t, ir.IRVariable) for t in ir_targets)
                    )]

                # Fallback to legacy unpack for other RHS types
                tmp_name = f"_unpack_{node.lineno}_{getattr(node, 'col_offset', 0)}"
                tmp_var = ir.IRVariable(name=tmp_name)
                results: List[ir.IRStatement] = [ir.IRAssignment(
                    lineno=node.lineno,
                    col_offset=getattr(node, 'col_offset', 0),
                    target=tmp_var,
                    value=v,
                    is_local=True
                )]
                for i, target in enumerate(ir_targets):
                    results.append(ir.IRAssignment(
                        lineno=node.lineno,
                        col_offset=getattr(node, 'col_offset', 0),
                        target=target,
                        value=ir.IRIndexOperation(
                            lineno=node.lineno,
                            col_offset=getattr(node, 'col_offset', 0),
                            receiver=tmp_var,
                            index=ir.IRLiteral(value=i + 1)
                        ),
                        is_local=isinstance(target, ir.IRVariable)
                    ))
                return results

            # Single target
            v = self.visit(node.value)
            is_local = isinstance(ir_targets[0], ir.IRVariable)
            return [ir.IRAssignment(
                lineno=node.lineno,
                col_offset=getattr(node, 'col_offset', 0),
                target=ir_targets[0],
                value=v,
                is_local=is_local
            )]
        
        # Multiple targets: x = y = z
        ir_value = self.visit(node.value)
        tmp_name = f"_tmp_assign_{node.lineno}_{getattr(node, 'col_offset', 0)}"
        tmp_var = ir.IRVariable(name=tmp_name)
        results = [ir.IRAssignment(lineno=node.lineno, target=tmp_var, value=ir_value, is_local=True)]
        for t in node.targets:
            for target in self._transform_target(t):
                results.append(ir.IRAssignment(lineno=node.lineno, target=target, value=tmp_var))
        return results
        
    def visit_AnnAssign(self, node: ast.AnnAssign) -> ir.IRStatement:
        target = self.visit(node.target)
        value = self.visit(node.value) if node.value else ir.IRLiteral(value=None)
        
        ann_str = self._get_type_str(node.annotation)
        if ann_str == "persistent":
            return ir.IRPersistentAssign(
                lineno=node.lineno,
                col_offset=node.col_offset,
                target=target,
                value=value
            )
            
        return ir.IRAssignment(
            lineno=node.lineno,
            col_offset=node.col_offset,
            target=target,
            value=value,
            annotation=ann_str,
            is_local=True
        )

    def visit_AugAssign(self, node: ast.AugAssign) -> ir.IRStatement:
        # x += 1 -> x = x + 1
        target = self.visit(node.target)
        value = self.visit(node.value)
        op_name = type(node.op).__name__
        
        return ir.IRAssignment(
            lineno=node.lineno,
            col_offset=node.col_offset,
            target=target,
            value=ir.IRBinaryOperation(
                lineno=node.lineno,
                left=target,
                operator=op_name,
                right=value
            ),
            is_local=False # Usually already declared
        )

    def visit_Expr(self, node: ast.Expr) -> ir.IRStatement:
        # In IR, standalone expressions are often just calls
        val = self.visit(node.value)
        if isinstance(val, (ir.IRMethodCall, ir.IRFunctionCall, ir.IRAwait, ir.IRSuperCall)):
            return val # type: ignore
        # Otherwise wrap in an assignment to _ or discard
        return ir.IRAssignment(
            lineno=node.lineno,
            col_offset=node.col_offset,
            target=ir.IRVariable(name="_", lineno=node.lineno),
            value=val
        )

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ir.IRFunctionDef:
        decorators = [d.id for d in node.decorator_list if isinstance(d, ast.Name)]
        
        # Capture positional defaults
        defaults = [self.visit(d) for d in node.args.defaults]
        # Pad with None for arguments without defaults
        num_args = len(node.args.args)
        num_defaults = len(defaults)
        full_defaults = [None] * (num_args - num_defaults) + defaults

        return ir.IRFunctionDef(
            lineno=node.lineno,
            col_offset=node.col_offset,
            name=node.name,
            args=[a.arg for a in node.args.args],
            arg_types=[self._get_type_str(a.annotation) for a in node.args.args],
            defaults=full_defaults, # type: ignore
            return_type=self._get_type_str(node.returns),
            body=self._visit_body(node.body),
            is_local=True,
            decorators=decorators
        )

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ir.IRAsyncFunctionDef:
        return ir.IRAsyncFunctionDef(
            lineno=node.lineno,
            col_offset=node.col_offset,
            name=node.name,
            args=[a.arg for a in node.args.args],
            arg_types=[self._get_type_str(a.annotation) for a in node.args.args],
            return_type=self._get_type_str(node.returns),
            body=self._visit_body(node.body),
            is_local=True,
            decorators=[d.id for d in node.decorator_list if isinstance(d, ast.Name)]
        )

    def visit_Return(self, node: ast.Return) -> ir.IRStatement:
        if isinstance(node.value, ast.Tuple):
            return ir.IRMultiReturn(
                lineno=node.lineno,
                values=[self.visit(elt) for elt in node.value.elts]
            )
        return ir.IRReturnStatement(
            lineno=node.lineno,
            col_offset=node.col_offset,
            value=self.visit(node.value) if node.value else None
        )

    def visit_If(self, node: ast.If) -> ir.IRIfStatement:
        return ir.IRIfStatement(
            lineno=node.lineno,
            col_offset=node.col_offset,
            condition=self.visit(node.test),
            then_block=self._visit_body(node.body),
            else_block=self._visit_body(node.orelse) if node.orelse else None
        )

    def visit_While(self, node: ast.While) -> ir.IRWhileStatement:
        return ir.IRWhileStatement(
            lineno=node.lineno,
            col_offset=node.col_offset,
            condition=self.visit(node.test),
            body=self._visit_body(node.body)
        )

    def visit_For(self, node: ast.For) -> ir.IRStatement:
        from transpiler.frontend.ast_utils import is_range_call, unpack_range_args
        if is_range_call(node.iter):
            start, stop, step = unpack_range_args(node.iter)
            return ir.IRForStatement(
                lineno=node.lineno,
                col_offset=node.col_offset,
                var=node.target.id if isinstance(node.target, ast.Name) else "_",
                start=self.visit(start),
                end=self.visit(stop),
                step=self.visit(step) if step else None,
                body=self._visit_body(node.body)
            )
        
        # Generic for
        loop_vars = []
        if isinstance(node.target, ast.Name):
            loop_vars = [node.target.id]
        elif isinstance(node.target, (ast.Tuple, ast.List)):
            loop_vars = [elt.id for elt in node.target.elts if isinstance(elt, ast.Name)]

        return ir.IRGenericForStatement(
            lineno=node.lineno,
            col_offset=node.col_offset,
            vars=loop_vars,
            iterator=self.visit(node.iter),
            body=self._visit_body(node.body)
        )

    def visit_Pass(self, node: ast.Pass) -> ir.IRPassStatement:
        return ir.IRPassStatement(lineno=node.lineno)

    def visit_Break(self, node: ast.Break) -> ir.IRBreakStatement:
        return ir.IRBreakStatement(lineno=node.lineno)

    def visit_Continue(self, node: ast.Continue) -> ir.IRContinueStatement:
        return ir.IRContinueStatement(lineno=node.lineno)

    def visit_Lambda(self, node: ast.Lambda) -> ir.IRLambda:
        return ir.IRLambda(
            lineno=node.lineno,
            col_offset=node.col_offset,
            args=[a.arg for a in node.args.args],
            arg_types=[self._get_type_str(a.annotation) for a in node.args.args],
            defaults=[self.visit(d) for d in node.args.defaults],
            body=[ir.IRReturnStatement(value=self.visit(node.body))],
            return_type=None # Lambda return type is harder to extract from AST easily
        )

    def visit_ClassDef(self, node: ast.ClassDef) -> ir.IRClassDef:
        bases = [self.visit(b) for b in node.bases]
        decorators = [d.id for d in node.decorator_list if isinstance(d, ast.Name)]
        return ir.IRClassDef(
            lineno=node.lineno,
            col_offset=node.col_offset,
            name=node.name,
            bases=bases,
            body=self._visit_body(node.body),
            decorators=decorators
        )

    def visit_Import(self, node: ast.Import) -> List[ir.IRImport]:
        return [
            ir.IRImport(
                lineno=node.lineno,
                col_offset=node.col_offset,
                module=alias.name,
                alias=alias.asname
            )
            for alias in node.names
        ]

    def visit_ImportFrom(self, node: ast.ImportFrom) -> ir.IRImportFrom:
        return ir.IRImportFrom(
            lineno=node.lineno,
            col_offset=node.col_offset,
            module=node.module or "",
            names=[alias.name for alias in node.names],
            aliases=[alias.asname for alias in node.names],
            level=node.level
        )
