"""
transpiler/printer.py — Serializes the Luau AST into source code with line tracking.
"""

from __future__ import annotations
from typing import Any, Dict, List
from transpiler.backend import luau_ast as last

# ---------------------------------------------------------------------------
# Operator precedence table (higher number = binds tighter)
# Based on Luau / Lua 5.x operator precedence rules.
# ---------------------------------------------------------------------------
_LUAU_OP_PREC: Dict[str, int] = {
    # Logical
    "or":   1,
    "and":  2,
    # Comparison   (non-associative – always parens across different comparisons)
    "<":    3,  ">": 3,  "<=": 3,  ">=": 3,  "==": 3,  "~=": 3,
    # Concat
    "..":   4,
    # Additive
    "+":    5,  "-":  5,
    # Multiplicative
    "*":    6,  "/":  6,  "//": 6,  "%":  6,
    # Unary ("not", "-", "#", "~") — handled separately
    # Power (right-associative)
    "^":    8,
    # If Expression (lowest priority)
    "if":   0,
}

_UNARY_PREC = 7   # unary ops bind tighter than multiplicative but looser than ^

def _op_prec(op: str) -> int:
    return _LUAU_OP_PREC.get(op, 5)  # default to additive level for unknown ops

class LuauPrinter:
    def __init__(self, indent_size: int = 4) -> None:
        self.indent_size = indent_size
        self.indent_level = 0
        self.output = []
        self.current_line = 1
        self.line_map: Dict[int, int] = {} # Luau Line -> Python Line

    def write(self, text: str):
        if not text:
            return
        lines = text.split("\n")
        for i, line in enumerate(lines):
            if i > 0:
                self.output.append("\n")
                self.current_line += 1
            if line:
                if self.output and self.output[-1] == "\n":
                    self.output.append(" " * (self.indent_size * self.indent_level))
                self.output.append(line)

    def write_indent(self):
        self.output.append(" " * (self.indent_size * self.indent_level))

    def newline(self):
        self.output.append("\n")
        self.current_line += 1

    def visit(self, node: Any) -> None:
        if node is None:
            self.write("nil")
            return
            
        if isinstance(node, last.Node):
            # Phase S4: Use source_lineno if available, otherwise fallback to node.lineno
            src_line = getattr(node, "source_lineno", getattr(node, "lineno", 0))
            if src_line != 0:
                # Record mapping for the first line of this node
                if self.current_line not in self.line_map:
                    self.line_map[self.current_line] = src_line
        
        method_name = f"visit_{node.__class__.__name__}"
        visitor = getattr(self, method_name, self.generic_visit)
        visitor(node)

    def generic_visit(self, node: Any) -> None:
        raise NotImplementedError(f"No visitor for {node.__class__.__name__}")

    def print_node(self, node: last.Node) -> str:
        self.visit(node)
        return "".join(self.output)

    # --- Expressions ---

    def visit_Variable(self, node: last.Variable) -> None:
        self.write(node.name)

    def visit_Literal(self, node: last.Literal) -> None:
        v = node.value
        if v is None: self.write("nil")
        elif v is True: self.write("true")
        elif v is False: self.write("false")
        elif isinstance(v, str):
            escaped = v.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")
            self.write(f'"{escaped}"')
        else:
            self.write(str(v))

    def visit_IfExpression(self, node: last.IfExpression) -> None:
        self.write("if ")
        self.visit(node.condition)
        self.write(" then ")
        self.visit(node.then_expr)
        self.write(" else ")
        self.visit(node.else_expr)

    def _visit_expr(self, node: Any, parent_prec: int = 0, is_right: bool = False) -> None:
        """Emit an expression, adding parentheses only when needed.

        *parent_prec*: the precedence of the enclosing operator (0 = outermost).
        *is_right*: True when this node is the right operand (for left-assoc operators
                    we need parens if child prec == parent prec on the right side).
        """
        if isinstance(node, last.BinaryOperation):
            my_prec = _op_prec(node.operator)
            # Need parens when child binds looser than parent, or same-level on the right
            # of a left-associative operator (handles `a - (b - c)` correctly).
            need_parens = my_prec < parent_prec or (is_right and my_prec == parent_prec)
            if need_parens:
                self.write("(")
            self._visit_expr(node.left, my_prec, is_right=False)
            self.write(f" {node.operator} ")
            self._visit_expr(node.right, my_prec, is_right=True)
            if need_parens:
                self.write(")")
        elif isinstance(node, last.UnaryOperation):
            my_prec = _UNARY_PREC
            need_parens = my_prec < parent_prec
            if need_parens:
                self.write("(")
            
            # Clean up spacing: 'not' needs a space, '-' and '~' do not.
            op = node.operator.strip()
            self.write(op)
            if op == "not":
                self.write(" ")
                
            self._visit_expr(node.operand, my_prec)
            if need_parens:
                self.write(")")
        elif isinstance(node, last.IfExpression):
            my_prec = _op_prec("if") # lowest
            need_parens = my_prec < parent_prec
            if need_parens:
                self.write("(")
            self.visit(node)
            if need_parens:
                self.write(")")
        else:
            self.visit(node)

    def visit_IfElseExpression(self, node: last.IfElseExpression) -> None:
        self.write("if ")
        self.visit(node.condition)
        self.write(" then ")
        self.visit(node.then_expr)
        self.write(" else ")
        self.visit(node.else_expr)

    def visit_BinaryOperation(self, node: last.BinaryOperation) -> None:
        self._visit_expr(node, parent_prec=0)

    def visit_UnaryOperation(self, node: last.UnaryOperation) -> None:
        self._visit_expr(node, parent_prec=0)

    def visit_FunctionCall(self, node: last.FunctionCall) -> None:
        if isinstance(node.func, last.Lambda):
            self.write("(")
            self.visit(node.func)
            self.write(")")
        else:
            self.visit(node.func)
        self.write("(")
        for i, arg in enumerate(node.args):
            if i > 0: self.write(", ")
            self.visit(arg)
        self.write(")")

    def visit_MethodCall(self, node: last.MethodCall) -> None:
        self.visit(node.receiver)
        self.write(f":{node.method}(")
        for i, arg in enumerate(node.args):
            if i > 0: self.write(", ")
            self.visit(arg)
        self.write(")")

    def visit_MemberAccess(self, node: last.MemberAccess) -> None:
        self.visit(node.receiver)
        self.write(f".{node.member}")

    def visit_IndexOperation(self, node: last.IndexOperation) -> None:
        self.visit(node.receiver)
        self.write("[")
        self.visit(node.index)
        self.write("]")

    def visit_TableLiteral(self, node: last.TableLiteral) -> None:
        self.write("{")
        first = True
        for e in node.elements:
            if not first: self.write(", ")
            self.visit(e)
            first = False
        for name, val in node.fields:
            if not first: self.write(", ")
            # Check if name needs square brackets (numeric or non-identifier)
            if name.isidentifier():
                self.write(f"{name} = ")
            else:
                # If it's a number-like string, don't quote it inside brackets for idiomatic Luau
                if name.replace(".", "", 1).isdigit():
                    self.write(f"[{name}] = ")
                else:
                    self.write(f"[\"{name}\"] = ")
            self.visit(val)
            first = False
        self.write("}")

    def _emit_block(self, stmts: list) -> None:
        """Emit an indented block without suppressing returns, as that broke early exit logic."""
        for stmt in stmts:
            self.newline()
            self.write_indent()
            self.visit(stmt)

    def visit_Lambda(self, node: last.Lambda) -> None:
        self.write("function(")
        for i, arg in enumerate(node.args):
            if i > 0: self.write(", ")
            type_ann = node.arg_types[i] if i < len(node.arg_types) else None
            self.write(f"{arg}: {type_ann}" if type_ann else arg)
        self.write(")")
        if node.return_type:
            self.write(f": {node.return_type}")
        
        self.indent_level += 1
        self._emit_block(node.body)
        self.indent_level -= 1
        self.newline()
        self.write_indent()
        self.write("end")

    # --- Statements ---

    def visit_LocalAssign(self, node: last.LocalAssign) -> None:
        self.write(f"local {node.name}")
        if node.type_annotation:
            self.write(f": {node.type_annotation}")
        if node.value:
            self.write(" = ")
            self.visit(node.value)

    def visit_Assignment(self, node: last.Assignment) -> None:
        self.visit(node.target)
        self.write(" = ")
        self.visit(node.value)

    def visit_MultiAssign(self, node: last.MultiAssign) -> None:
        if node.is_local:
            self.write("local ")
        
        for i, target in enumerate(node.targets):
            if i > 0: self.write(", ")
            self.visit(target)
        
        self.write(" = ")
        
        if node.rhs_call:
            self.visit(node.rhs_call)
        else:
            for i, val in enumerate(node.values):
                if i > 0: self.write(", ")
                self.visit(val)

    def visit_CompoundAssignment(self, node: last.CompoundAssignment) -> None:
        self.visit(node.target)
        self.write(f" {node.operator}= ")
        self.visit(node.value)

    def visit_IfStatement(self, node: last.IfStatement) -> None:
        self.write("if ")
        self.visit(node.condition)
        self.write(" then")
        
        self.indent_level += 1
        for stmt in node.then_block:
            self.newline()
            self.write_indent()
            self.visit(stmt)
        self.indent_level -= 1
        
        for e_cond, e_body in node.elif_blocks:
            self.newline()
            self.write_indent()
            self.write("elseif ")
            self.visit(e_cond)
            self.write(" then")
            self.indent_level += 1
            for stmt in e_body:
                self.newline()
                self.write_indent()
                self.visit(stmt)
            self.indent_level -= 1
            
        if node.else_block:
            self.newline()
            self.write_indent()
            self.write("else")
            self.indent_level += 1
            for stmt in node.else_block:
                self.newline()
                self.write_indent()
                self.visit(stmt)
            self.indent_level -= 1
            
        self.newline()
        self.write_indent()
        self.write("end")

    def visit_WhileStatement(self, node: last.WhileStatement) -> None:
        self.write("while ")
        self.visit(node.condition)
        self.write(" do")
        self.indent_level += 1
        for stmt in node.body:
            self.newline()
            self.write_indent()
            self.visit(stmt)
        self.indent_level -= 1
        self.newline()
        self.write_indent()
        self.write("end")

    def visit_ForStatement(self, node: last.ForStatement) -> None:
        self.write(f"for {node.var} = ")
        self.visit(node.start)
        self.write(", ")
        self.visit(node.end)
        if node.step:
            self.write(", ")
            self.visit(node.step)
        self.write(" do")
        self.indent_level += 1
        for stmt in node.body:
            self.newline()
            self.write_indent()
            self.visit(stmt)
        self.indent_level -= 1
        self.newline()
        self.write_indent()
        self.write("end")

    def visit_GenericForStatement(self, node: last.GenericForStatement) -> None:
        self.write("for ")
        for i, v in enumerate(node.vars):
            if i > 0: self.write(", ")
            type_ann = node.var_types[i] if i < len(node.var_types) else None
            self.write(f"{v}: {type_ann}" if type_ann else v)
        self.write(" in ")
        self.visit(node.iterator)
        self.write(" do")
        self.indent_level += 1
        for stmt in node.body:
            self.newline()
            self.write_indent()
            self.visit(stmt)
        self.indent_level -= 1
        self.newline()
        self.write_indent()
        self.write("end")

    def visit_FunctionDef(self, node: last.FunctionDef) -> None:
        if node.is_local: self.write("local ")
        self.write(f"function {node.name}(")
        for i, arg in enumerate(node.args):
            if i > 0: self.write(", ")
            type_ann = node.arg_types[i] if i < len(node.arg_types) else None
            self.write(f"{arg}: {type_ann}" if type_ann else arg)
        self.write(")")
        if node.return_type:
            self.write(f": {node.return_type}")
        
        self.indent_level += 1
        body = node.body
        # Strip trailing bare `return` at function level only — void function trailing return is superfluous in Luau.
        if body and isinstance(body[-1], last.ReturnStatement) and body[-1].is_bare and body[-1].value is None:
            body = body[:-1]
        self._emit_block(body)
        self.indent_level -= 1
        self.newline()
        self.write_indent()
        self.write("end")

    def visit_ReturnStatement(self, node: last.ReturnStatement) -> None:
        # Always emit return — bare returns inside conditionals must not be dropped.
        self.write("return")
        if node.value:
            self.write(" ")
            self.visit(node.value)

    def visit_MultiReturnStatement(self, node: last.MultiReturnStatement) -> None:
        self.write("return ")
        for i, val in enumerate(node.values):
            if i > 0: self.write(", ")
            self.visit(val)

    def visit_BreakStatement(self, node: last.BreakStatement) -> None:
        self.write("break")

    def visit_ContinueStatement(self, node: last.ContinueStatement) -> None:
        self.write("continue")

    def visit_PassStatement(self, node: last.PassStatement) -> None:
        self.write("-- pass")

    def visit_Comment(self, node: last.Comment) -> None:
        if node.text == "RPY_TRACEBACK_MAP_PLACEHOLDER":
             # Inject the traceback registration call
             self.write(f"require(sdk.runtime.traceback_handler).register_map(\"{getattr(self, 'module_name', 'unknown')}\", ")
             # Serialize the line map as a Luau table
             self.write("{")
             first = True
             for luau_line, py_line in self.line_map.items():
                 if not first: self.write(", ")
                 self.write(f"[{luau_line}] = {py_line}")
                 first = False
             self.write("})")
        else:
            self.write(f"-- {node.text}")

    def visit_RawCode(self, node: last.RawCode) -> None:
        # Phase S4: Handle stdlib shims if triggered
        if node.code == "--!strict":
            self.write("--!strict\n")
            # If generator added 'json' or 'base64' to runtime_used, inject shims
            # (In a real system, this would be more structured, but this fits the current snippets model)
            return

        self.write(node.code)

    def _classify_stmt(self, stmt: Any) -> str:
        if isinstance(stmt, (last.LocalAssign, last.Assignment)): return "DECL"
        if isinstance(stmt, last.FunctionDef): return "FUNC"
        if isinstance(stmt, (last.IfStatement, last.WhileStatement, last.ForStatement, last.GenericForStatement)): return "LOGIC"
        if isinstance(stmt, (last.FunctionCall, last.MethodCall, last.ReturnStatement, last.BreakStatement, last.ContinueStatement)): return "EXEC"
        if isinstance(stmt, last.Comment): return "COMMENT"
        if isinstance(stmt, last.RawCode): return "MISC"
        return "MISC"

    def visit_Block(self, node: last.Block) -> None:
        # S5: Native VM Guardrail
        if node.is_function_body and node.can_native:
            self.write("--!native\n")
            self.write_indent()

        body = list(node.body)
        
        # Strip trailing bare `return` at module/block level
        while body and isinstance(body[-1], last.ReturnStatement) and body[-1].is_bare and body[-1].value is None:
            body.pop()
        
        # Strip trailing 'pass' comments too
        while body and isinstance(body[-1], last.PassStatement):
            body.pop()

        # If it's a module, and we stripped a return, make sure we didn't leave an empty block
        # unless it was already empty. 

        prev_kind = None
        for i, stmt in enumerate(body):
            kind = self._classify_stmt(stmt)
            if prev_kind and kind != prev_kind:
                if not (kind == "COMMENT" and prev_kind == "DECL"):
                    self.newline()
            
            if i > 0:
                self.newline()
                self.write_indent()
            
            self.visit(stmt)
            prev_kind = kind
