"""
transpiler/ir.py — Intermediate Representation (IR) for RPy.
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import List, Optional, Union, Dict, Any
from enum import Enum, auto

class CallStyle(Enum):
    DOT = auto()
    COLON = auto()
    UNKNOWN = auto()

class EffectLevel(Enum):
    PURE = auto()
    READS_MEMORY = auto()
    WRITES_MEMORY = auto()
    HAS_SIDE_EFFECT = auto()
    UNKNOWN = auto()

@dataclass(kw_only=True)
class IRNode:
    """Base class for all IR nodes."""
    lineno: int = 0
    col_offset: int = 0

@dataclass
class TypeMetadata:
    luau_type: str  # 'number', 'string', 'boolean', 'any', 'nil'
    confidence: str # 'certain' | 'inferred' | 'unknown'

# --- Expressions ---

@dataclass(kw_only=True)
class IRExpression(IRNode):
    inferred_type: Optional[str] = None
    type_info: Optional[TypeMetadata] = None
    effect_level: EffectLevel = EffectLevel.UNKNOWN

@dataclass(kw_only=True)
class IRLiteral(IRExpression):
    value: Any = None
    effect_level: EffectLevel = EffectLevel.PURE
    
    def __repr__(self):
        return f"Lit({repr(self.value)})"

@dataclass(kw_only=True)
class IRVariable(IRExpression):
    name: str = ""
    original_name: Optional[str] = None # Used for SSA renaming
    definition_location: Optional[tuple[int, int, Optional[str]]] = None # (line, col, origin_module)
    effect_level: EffectLevel = EffectLevel.PURE
    
    def __repr__(self):
        return f"Var({self.name})"

@dataclass(kw_only=True)
class IRBinaryOperation(IRExpression):
    left: IRExpression = field(default_factory=lambda: IRLiteral(value=0))
    operator: str = "+"
    right: IRExpression = field(default_factory=lambda: IRLiteral(value=0))
    
    def __repr__(self):
        return f"({self.left} {self.operator} {self.right})"

@dataclass(kw_only=True)
class IRUnaryOperation(IRExpression):
    operator: str = "-"
    operand: IRExpression = field(default_factory=lambda: IRLiteral(value=0))
    
    def __repr__(self):
        return f"{self.operator}({self.operand})"

@dataclass(kw_only=True)
class IRFunctionCall(IRExpression):
    func: IRExpression = field(default_factory=lambda: IRVariable(name=""))
    args: List[IRExpression] = field(default_factory=list)
    
    def __repr__(self):
        return f"{self.func}({', '.join(map(str, self.args))})"

@dataclass(kw_only=True)
class IRMethodCall(IRExpression):
    receiver: IRExpression = field(default_factory=lambda: IRVariable())
    method: str = ""
    args: List[IRExpression] = field(default_factory=list)
    call_style: CallStyle = CallStyle.UNKNOWN
    
    def __repr__(self):
        sep = ":" if self.call_style == CallStyle.COLON else "."
        return f"{self.receiver}{sep}{self.method}({', '.join(map(str, self.args))})"

@dataclass(kw_only=True)
class IRSuperCall(IRExpression):
    method: str = ""
    args: List[IRExpression] = field(default_factory=list)
    
    def __repr__(self):
        return f"super().{self.method}({', '.join(map(str, self.args))})"

@dataclass(kw_only=True)
class IRPropertyAccess(IRExpression):
    receiver: IRExpression = field(default_factory=lambda: IRVariable())
    property: str = ""
    
    def __repr__(self):
        return f"{self.receiver}.{self.property}"

@dataclass(kw_only=True)
class IRTableLiteral(IRExpression):
    elements: List[IRExpression] = field(default_factory=list)
    fields: List[tuple[str, IRExpression]] = field(default_factory=list)
    
    def __repr__(self):
        parts = [str(e) for e in self.elements] + [f"{k}={v}" for k, v in self.fields]
        return "{" + ", ".join(parts) + "}"

@dataclass(kw_only=True)
class IRIndexOperation(IRExpression):
    receiver: IRExpression = field(default_factory=lambda: IRVariable())
    index: IRExpression = field(default_factory=lambda: IRLiteral(value=0))
    
    def __repr__(self):
        return f"{self.receiver}[{self.index}]"

@dataclass(kw_only=True)
class IRPhi(IRExpression):
    # block_id -> versioned variable or expression
    options: Dict[int, IRExpression] = field(default_factory=dict)
    
    def __repr__(self):
        val_str = ", ".join(f"B{bid}:{opt}" for bid, opt in self.options.items())
        return f"Phi({val_str})"

@dataclass(kw_only=True)
class IRIfElseExpression(IRExpression):
    condition: IRExpression = field(default_factory=lambda: IRLiteral(value=True))
    then_expr: IRExpression = field(default_factory=lambda: IRLiteral())
    else_expr: IRExpression = field(default_factory=lambda: IRLiteral())

    def __repr__(self):
        return f"({self.then_expr} if {self.condition} else {self.else_expr})"

@dataclass(kw_only=True)
class IRAwait(IRExpression):
    value: IRExpression = field(default_factory=lambda: IRVariable())
    
    def __repr__(self):
        return f"await {self.value}"

# --- Statements ---

@dataclass(kw_only=True)
class IRStatement(IRNode):
    pass

@dataclass(kw_only=True)
class IRModule(IRStatement):
    body: List[IRStatement] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)

@dataclass(kw_only=True)
class IRAssignment(IRStatement):
    target: IRExpression = field(default_factory=lambda: IRVariable())
    value: IRExpression = field(default_factory=lambda: IRLiteral())
    annotation: Optional[str] = None
    is_local: bool = False

    
    def __repr__(self):
        return f"{self.target} = {self.value}"

@dataclass(kw_only=True)
class IRPersistentAssign(IRStatement):
    target: IRExpression = field(default_factory=lambda: IRVariable())
    value: IRExpression = field(default_factory=lambda: IRLiteral())
    
    def __repr__(self):
        return f"PERSISTENT {self.target} = {self.value}"

@dataclass(kw_only=True)
class IRFunctionDef(IRStatement):
    name: str = ""
    args: List[str] = field(default_factory=list)
    body: List[IRStatement] = field(default_factory=list)
    is_local: bool = True
    is_method: bool = False
    arg_types: List[Optional[str]] = field(default_factory=list)
    defaults: List[Optional[IRExpression]] = field(default_factory=list)
    return_type: Optional[str] = None
    decorators: List[str] = field(default_factory=list)
    
    def __repr__(self):
        return f"def {self.name}({', '.join(self.args)})"

@dataclass(kw_only=True)
class IRAsyncFunctionDef(IRFunctionDef):
    def __repr__(self):
        return f"async def {self.name}({', '.join(self.args)})"

@dataclass(kw_only=True)
class IRLambda(IRExpression):
    args: List[str] = field(default_factory=list)
    body: List[IRStatement] = field(default_factory=list)
    arg_types: List[Optional[str]] = field(default_factory=list)
    defaults: List[IRExpression] = field(default_factory=list)
    return_type: Optional[str] = None

@dataclass(kw_only=True)
class IRIfStatement(IRStatement):
    condition: IRExpression = field(default_factory=lambda: IRLiteral(value=True))
    then_block: List[IRStatement] = field(default_factory=list)
    elif_blocks: List[tuple[IRExpression, List[IRStatement]]] = field(default_factory=list)
    else_block: Optional[List[IRStatement]] = None
    
    def __repr__(self):
        return f"if {self.condition}"

@dataclass(kw_only=True)
class IRWhileStatement(IRStatement):
    condition: IRExpression = field(default_factory=lambda: IRLiteral(value=True))
    body: List[IRStatement] = field(default_factory=list)
    
    def __repr__(self):
        return f"while {self.condition}"

@dataclass(kw_only=True)
class IRForStatement(IRStatement):
    var: str = ""
    start: IRExpression = field(default_factory=lambda: IRLiteral(value=0))
    end: IRExpression = field(default_factory=lambda: IRLiteral(value=10))
    body: List[IRStatement] = field(default_factory=list)
    step: Optional[IRExpression] = None
    
    def __repr__(self):
        body_str = "\n".join(f"      {s}" for s in self.body)
        return f"for {self.var} in {self.start}..{self.end}:\n{body_str}"

@dataclass(kw_only=True)
class IRGenericForStatement(IRStatement):
    vars: List[str] = field(default_factory=list)
    var_types: List[Optional[str]] = field(default_factory=list)
    iterator: IRExpression = field(default_factory=lambda: IRVariable(name=""))
    body: List[IRStatement] = field(default_factory=list)

@dataclass(kw_only=True)
class IRReturnStatement(IRStatement):
    value: Optional[IRExpression] = None
    
    def __repr__(self):
        return f"return {self.value}"

@dataclass(kw_only=True)
class IRMultiReturn(IRStatement):
    """return a, b, c  — native Luau multi-return.
    Emitted when Python 'return <tuple>' is detected.
    """
    values: List[IRExpression] = field(default_factory=list)

    def __repr__(self):
        return f"return {', '.join(map(str, self.values))}"


@dataclass(kw_only=True)
class IRMultiAssign(IRStatement):
    """local x, y = f()  OR  local x, y = a, b

    Shapes:
      rhs_call is set   → local <targets> = <call>      (unpack multi-return)
      rhs_values        → local <targets> = <v1>, <v2>  (parallel assign)
    """
    targets: List[IRExpression] = field(default_factory=list)
    rhs_call: Optional[IRExpression] = None
    rhs_values: List[IRExpression] = field(default_factory=list)
    is_local: bool = True
    is_loop_assignment: bool = False # Skip emission if it's just mirroring loop variables

    def __repr__(self):
        if self.rhs_call:
            return f"local {', '.join(map(str, self.targets))} = {self.rhs_call}"
        return f"local {', '.join(map(str, self.targets))} = {', '.join(map(str, self.rhs_values))}"

@dataclass(kw_only=True)
class IRBreakStatement(IRStatement):
    def __repr__(self): return "break"

@dataclass(kw_only=True)
class IRContinueStatement(IRStatement):
    def __repr__(self): return "continue"

@dataclass(kw_only=True)
class IRPassStatement(IRStatement):
    def __repr__(self): return "pass"

@dataclass(kw_only=True)
class IRComment(IRStatement):
    text: str = ""
    def __repr__(self): return f"# {self.text}"

@dataclass(kw_only=True)
class IRClassDef(IRStatement):
    name: str = ""
    bases: List[IRExpression] = field(default_factory=list)
    body: List[IRStatement] = field(default_factory=list)
    properties: Dict[str, str] = field(default_factory=dict)
    decorators: List[str] = field(default_factory=list)

@dataclass(kw_only=True)
class IRImport(IRStatement):
    module: str = ""
    alias: Optional[str] = None

@dataclass(kw_only=True)
class IRImportFrom(IRStatement):
    module: str = ""
    names: List[str] = field(default_factory=list)
    aliases: List[Optional[str]] = field(default_factory=list)
    level: int = 0

# --- CFG Terminators ---

@dataclass(kw_only=True)
class IRJump(IRStatement):
    target_block_id: int = 0
    def __repr__(self): return f"jump B{self.target_block_id}"

@dataclass(kw_only=True)
class IRBranch(IRStatement):
    condition: IRExpression = field(default_factory=lambda: IRLiteral(value=True))
    true_block_id: int = 0
    false_block_id: int = 0
    
    def __repr__(self):
        return f"branch {self.condition} then B{self.true_block_id} else B{self.false_block_id}"

@dataclass(kw_only=True)
class IRGenericForTerminator(IRStatement):
    vars: List[str] = field(default_factory=list)
    iterator: IRExpression = field(default_factory=lambda: IRVariable(name=""))
    true_block_id: int = 0  # Body
    false_block_id: int = 0 # Exit
    
    def __repr__(self):
        return f"for {', '.join(self.vars)} in {self.iterator} then B{self.true_block_id} else B{self.false_block_id}"
