import t, c
from stdint import *
import memhub
import string
import viperlib
from .tokens import TokOp


# ============================================================
# ASTKind - Node type enum (kind() virtual function return value)
# ============================================================

class ASTKind(t.CEnum):
    # Core 15 nodes
    Module: t.State
    FunctionDef: t.State
    ClassDef: t.State
    Assign: t.State
    If: t.State
    For: t.State
    While: t.State
    Return: t.State
    Expr: t.State
    Name: t.State
    Constant: t.State
    BinOp: t.State
    UnaryOp: t.State
    Call: t.State
    Compare: t.State
    # Module Layer
    Expression: t.State
    Interactive: t.State
    FunctionType: t.State
    # Statement
    Delete: t.State
    AugAssign: t.State
    AnnAssign: t.State
    With: t.State
    Raise: t.State
    Try: t.State
    Assert: t.State
    Global: t.State
    Nonlocal: t.State
    Pass: t.State
    Break: t.State
    Continue: t.State
    Import: t.State
    ImportFrom: t.State
    Match: t.State
    # Expression Layer
    BoolOp: t.State
    Lambda: t.State
    IfExp: t.State
    Dict: t.State
    Set: t.State
    ListComp: t.State
    SetComp: t.State
    DictComp: t.State
    GeneratorExp: t.State
    Await: t.State
    Yield: t.State
    YieldFrom: t.State
    FormattedValue: t.State
    JoinedStr: t.State
    Attribute: t.State
    Subscript: t.State
    Starred: t.State
    List: t.State
    Tuple: t.State
    Slice: t.State
    NamedExpr: t.State
    # Auxiliary Nodes
    ExceptHandler: t.State
    Arguments: t.State
    Arg: t.State
    Keyword: t.State
    Alias: t.State
    WithItem: t.State
    Comprehension: t.State
    OpNode: t.State
    # Pattern Matching
    MatchCase: t.State
    MatchValue: t.State
    MatchSingleton: t.State
    MatchSequence: t.State
    MatchMapping: t.State
    MatchClass: t.State
    MatchStar: t.State
    MatchAs: t.State
    MatchOr: t.State


# ============================================================
# ASTCtx - Name context (Load/Store/Del)
# ============================================================

class ASTCtx(t.CEnum):
    Load: t.State
    Store: t.State
    Del: t.State


# ============================================================
# OpKind - Operator enum (BinOp/UnaryOp/Compare)
# ============================================================

class OpKind(t.CEnum):
    # Binary arithmetic operators
    Add:       t.State
    Sub:       t.State
    Mult:      t.State
    MatMult:   t.State
    Div:       t.State
    Mod:       t.State
    Pow:       t.State
    LShift:    t.State
    RShift:    t.State
    BitOr:     t.State
    BitXor:    t.State
    BitAnd:    t.State
    FloorDiv:  t.State
    # Boolean operators
    And:       t.State
    Or:        t.State
    # Comparison operators
    Eq:        t.State
    Ne:        t.State
    Lt:        t.State
    Le:        t.State
    Gt:        t.State
    Ge:        t.State
    Is:        t.State
    IsNot:     t.State
    In:        t.State
    NotIn:     t.State
    # Unary operators
    Not:       t.State
    UAdd:      t.State
    USub:      t.State
    Invert:    t.State
    # None
    NoneOp:    t.State


# Constant subtype identifiers
CONST_INT: t.CDefine = 1
CONST_FLOAT: t.CDefine = 2
CONST_STR: t.CDefine = 3
CONST_BOOL: t.CDefine = 4
CONST_NONE: t.CDefine = 5
CONST_CHAR: t.CDefine = 6

# Node flags
FLAG_IS_ASYNC: t.CDefine = 1
FLAG_SIMPLE: t.CDefine = 2
FLAG_HAS_STAR: t.CDefine = 4


# ============================================================
# ASTFlag - Node flags enum
# ============================================================

class ASTFlag(t.CEnum):
    IsAsync: t.State
    Simple: t.State
    HasStar: t.State


# ============================================================
# Helper Functions
# ============================================================

def _copy_str(pool: memhub.MemManager | t.CPtr, src: str) -> str:
    """Allocate and copy a C string from pool."""
    if src is None: return None
    slen: t.CSizeT = string.strlen(src)
    buf: str = pool.alloc(slen + 1)
    if buf is None: return None
    string.memcpy(buf, src, slen + 1)
    return buf


def _set_pos(node: AST | t.CPtr, lineno: t.CInt, col_offset: t.CInt,
             end_lineno: t.CInt, end_col_offset: t.CInt):
    if node is None: return
    node.lineno = lineno
    node.col_offset = col_offset
    node.end_lineno = end_lineno
    node.end_col_offset = end_col_offset


def _inherit_pos(node: AST | t.CPtr, ref: AST | t.CPtr):
    if node is None or ref is None: return
    node.lineno = ref.lineno
    node.col_offset = ref.col_offset
    node.end_lineno = ref.end_lineno
    node.end_col_offset = ref.end_col_offset


def _init_ast(node: AST | t.CPtr, pool: memhub.MemManager | t.CPtr,
              lineno: t.CInt, col: t.CInt):
    """Initialize AST base fields."""
    node.parent = None
    node.pool = pool
    node.lineno = lineno
    node.col_offset = col
    node.end_lineno = lineno
    node.end_col_offset = col
    node.children = None


def _set_parent_list(lst: list[AST | t.CPtr] | t.CPtr, parent: AST | t.CPtr):
    """Set list children parent pointers."""
    if lst is None: return
    n: t.CSizeT = lst.__len__()
    i: t.CSizeT = 0
    while i < n:
        child: AST | t.CPtr = lst.get(i)
        if child is not None:
            child.parent = parent
        i += 1


def _append_child(lst: list[AST | t.CPtr] | t.CPtr, node: AST | t.CPtr,
                  pool: memhub.MemManager | t.CPtr):
    """Append child nodes to a list (module-level function, to avoid conflicts with the AST.append method name).

       Call this function within the AST.append method to bypass the compiler mistaking list.append for a virtual function call to AST.append when doing type inference."""
    if lst is None or node is None:
        return
    lst.append(node)


# ============================================================
# Operator Mapping: TokOp -> OpKind
# ============================================================

def _binop_from_op(op_id: t.CInt) -> t.CInt:
    if op_id == TokOp.Plus:      return OpKind.Add
    if op_id == TokOp.Minus:     return OpKind.Sub
    if op_id == TokOp.Star:      return OpKind.Mult
    if op_id == TokOp.At:        return OpKind.MatMult
    if op_id == TokOp.Slash:     return OpKind.Div
    if op_id == TokOp.Percent:   return OpKind.Mod
    if op_id == TokOp.StarStar:  return OpKind.Pow
    if op_id == TokOp.LtLt:      return OpKind.LShift
    if op_id == TokOp.GtGt:      return OpKind.RShift
    if op_id == TokOp.VBar:      return OpKind.BitOr
    if op_id == TokOp.Caret:     return OpKind.BitXor
    if op_id == TokOp.Amp:       return OpKind.BitAnd
    if op_id == TokOp.DSlash:    return OpKind.FloorDiv
    return OpKind.NoneOp


def _augop_from_op(op_id: t.CInt) -> t.CInt:
    if op_id == TokOp.PlusEq:    return OpKind.Add
    if op_id == TokOp.MinusEq:   return OpKind.Sub
    if op_id == TokOp.StarEq:    return OpKind.Mult
    if op_id == TokOp.AtEq:      return OpKind.MatMult
    if op_id == TokOp.SlashEq:   return OpKind.Div
    if op_id == TokOp.PercentEq: return OpKind.Mod
    if op_id == TokOp.StarEqEq:  return OpKind.Pow
    if op_id == TokOp.LtLtEq:    return OpKind.LShift
    if op_id == TokOp.GtGtEq:    return OpKind.RShift
    if op_id == TokOp.VBarEq:    return OpKind.BitOr
    if op_id == TokOp.CaretEq:   return OpKind.BitXor
    if op_id == TokOp.AmpEq:     return OpKind.BitAnd
    if op_id == TokOp.DSlashEq:  return OpKind.FloorDiv
    return OpKind.NoneOp


def _cmpop_from_op(op_id: t.CInt) -> t.CInt:
    if op_id == TokOp.EqEq:      return OpKind.Eq
    if op_id == TokOp.ExclaimEq: return OpKind.Ne
    if op_id == TokOp.Less:      return OpKind.Lt
    if op_id == TokOp.LessEq:    return OpKind.Le
    if op_id == TokOp.Greater:   return OpKind.Gt
    if op_id == TokOp.GreaterEq: return OpKind.Ge
    return OpKind.NoneOp


# ============================================================
# dump Helper Functions
# ============================================================

def _emit(buf: t.CChar | t.CPtr, size: t.CSizeT, pos: t.CSizeT,
          text: str) -> t.CSizeT:
    """Append literal to buf+pos."""
    if buf is None or pos >= size:
        return pos
    cur: t.CChar | t.CPtr = (t.CVoid | t.CPtr)(t.CUInt64T(buf) + pos)
    rem: t.CSizeT = size - pos
    n: t.CInt = viperlib.snprintf(cur, rem, "%s", text)
    if n < 0:
        return pos
    return pos + t.CSizeT(n)


def _emit_str(buf: t.CChar | t.CPtr, size: t.CSizeT, pos: t.CSizeT,
              text: str) -> t.CSizeT:
    """Append a string value to buf pos (using %s format)"""
    if buf is None or pos >= size:
        return pos
    cur: t.CChar | t.CPtr = (t.CVoid | t.CPtr)(t.CUInt64T(buf) + pos)
    rem: t.CSizeT = size - pos
    n: t.CInt = viperlib.snprintf(cur, rem, "%s", text)
    if n < 0:
        return pos
    return pos + t.CSizeT(n)


def _emit_int(buf: t.CChar | t.CPtr, size: t.CSizeT, pos: t.CSizeT,
              val: t.CInt64T) -> t.CSizeT:
    """Append an integer value to buf pos (using %lld format)."""
    if buf is None or pos >= size:
        return pos
    cur: t.CChar | t.CPtr = (t.CVoid | t.CPtr)(t.CUInt64T(buf) + pos)
    rem: t.CSizeT = size - pos
    n: t.CInt = viperlib.snprintf(cur, rem, "%lld", val)
    if n < 0:
        return pos
    return pos + t.CSizeT(n)


def _op_name(op: t.CInt) -> str:
    """Return the string name of an OpKind."""
    if op == OpKind.Add:       return "Add"
    if op == OpKind.Sub:       return "Sub"
    if op == OpKind.Mult:      return "Mult"
    if op == OpKind.MatMult:   return "MatMult"
    if op == OpKind.Div:       return "Div"
    if op == OpKind.Mod:       return "Mod"
    if op == OpKind.Pow:       return "Pow"
    if op == OpKind.LShift:    return "LShift"
    if op == OpKind.RShift:    return "RShift"
    if op == OpKind.BitOr:     return "BitOr"
    if op == OpKind.BitXor:    return "BitXor"
    if op == OpKind.BitAnd:    return "BitAnd"
    if op == OpKind.FloorDiv:  return "FloorDiv"
    if op == OpKind.And:       return "And"
    if op == OpKind.Or:        return "Or"
    if op == OpKind.Eq:        return "Eq"
    if op == OpKind.Ne:        return "Ne"
    if op == OpKind.Lt:        return "Lt"
    if op == OpKind.Le:        return "Le"
    if op == OpKind.Gt:        return "Gt"
    if op == OpKind.Ge:        return "Ge"
    if op == OpKind.Is:        return "Is"
    if op == OpKind.IsNot:     return "IsNot"
    if op == OpKind.In:        return "In"
    if op == OpKind.NotIn:     return "NotIn"
    if op == OpKind.Not:       return "Not"
    if op == OpKind.UAdd:      return "UAdd"
    if op == OpKind.USub:      return "USub"
    if op == OpKind.Invert:    return "Invert"
    return "NoneOp"


def _dump_list(lst: list[AST | t.CPtr] | t.CPtr, buf: t.CChar | t.CPtr,
               size: t.CSizeT, pos: t.CSizeT) -> t.CSizeT:
    """Traverse a list[AST | CPtr] container, polymorphically dump each element."""
    if lst is None:
        pos = _emit(buf, size, pos, "[]")
        return pos
    pos = _emit(buf, size, pos, "[")
    n: t.CSizeT = lst.__len__()
    i: t.CSizeT = 0
    while i < n:
        if i > 0:
            pos = _emit(buf, size, pos, ", ")
        child: AST | t.CPtr = lst.get(i)
        if child is not None:
            pos = child.dump(buf, size, pos)
        i += 1
    pos = _emit(buf, size, pos, "]")
    return pos


def _dump_op_list(lst: t.CPtr, buf: t.CChar | t.CPtr,
                  size: t.CSizeT, pos: t.CSizeT) -> t.CSizeT:
    """Traverse a list[CInt] (OpKind values), output a list of op names."""
    if lst is None:
        pos = _emit(buf, size, pos, "[]")
        return pos
    ops_list: list[t.CInt] | t.CPtr = (list[t.CInt] | t.CPtr)(lst)
    pos = _emit(buf, size, pos, "[")
    n: t.CSizeT = ops_list.__len__()
    i: t.CSizeT = 0
    while i < n:
        if i > 0:
            pos = _emit(buf, size, pos, ", ")
        op: t.CInt = ops_list.get(i)
        pos = _emit(buf, size, pos, _op_name(op))
        i += 1
    pos = _emit(buf, size, pos, "]")
    return pos


# ============================================================
# AST - Polymorphic Base Class

# @t.CVTable enables vtable, supporting virtual functions like kind()/type_name()/dump()/accept()
# Child nodes are stored in a list[AST | CPtr] container (O(1) random access, more efficient than linked lists)
# The parent field maintains a pointer to the parent (used for upward traversal)
# ============================================================

@t.CVTable
class AST:
    """AST node base class. All specific node classes inherit from this.

    Fields:
        parent: Pointer to the parent node (for upward traversal)
        pool: Allocator (used for allocating child nodes/strings)
        lineno/col_offset: Starting position
        end_lineno/end_col_offset: Ending position
        children: List of child nodes (body block, added via append)
    """
    parent:          AST | t.CPtr
    pool:            memhub.MemManager | t.CPtr
    lineno:          t.CInt
    col_offset:      t.CInt
    end_lineno:      t.CInt
    end_col_offset:  t.CInt
    children:        list[AST | t.CPtr] | t.CPtr

    def kind(self) -> t.CInt:
        """Returns the node type (ASTKind value), overridden by subclasses"""
        return 0

    def type_name(self) -> str:
        """Return the node type name, subclasses override"""
        return "AST"

    def dump(self, buf: t.CChar | t.CPtr, size: t.CSizeT,
             pos: t.CSizeT) -> t.CSizeT:
        """Polymorphic dump: subclasses override to define custom format. Returns the written pos"""
        return pos

    def accept(self, visitor: t.CPtr):
        """Visitor pattern entry point: subclasses override to dispatch to ASTVisitor in the visitor pattern.
        Current implementation is empty. override this method to dispatch to visitor.visit_X(self)"""
        pass

    def append(self, node: AST | t.CPtr):
        """Append node to the children list.

        Used for parser to build statement blocks: module.append(stmt), if_node.append(stmt), etc.  
        The children list is lazily initialized on first call.  

        Note: list[AST|CPtr] is currently stored by value (compiler limitation), so you must set  
        node.parent before appending, so the parent in the copy is correct."""
        if self is None or node is None:
            return
        if self.children is None:
            self.children = list[AST | t.CPtr](self.pool, 8)
        node.parent = self
        _append_child(self.children, node, self.pool)


# ============================================================
# Stmt - Statement Base Class
#
# 所有语句节点（Assign/If/For/While/Return/...）继承 Stmt。
# isinstance(node, ast.Stmt) 可在运行时区分"语句"与"表达式"
# （依赖 TPV isinstance 子类集合判定 + 三层 vtable 继承）。
# 不 override kind()（继承 AST 的 0），具体节点各自返回 ASTKind 值。
# ============================================================

@t.CVTable
class Stmt(AST):
    """语句节点基类（Assign/If/For/While/Return/...）"""


# ============================================================
# ExprNode - Expression Base Class
#
# 所有表达式节点（Name/Constant/Call/BinOp/...）继承 ExprNode。
# isinstance(node, ast.ExprNode) 可在运行时区分"表达式"与"语句"。
#
# 命名说明：不能叫 Expr——ASTKind.Expr（base.vp L23）与 stmts.vp 的
# Expr 类已是"表达式语句"（Expression Statement）的既有语义。
# ============================================================

@t.CVTable
class ExprNode(AST):
    """表达式节点基类（Name/Constant/Call/BinOp/...）"""


# ============================================================
# Module dump function (for ast.dump(node, buf, size))
# ============================================================

def dump(node: AST | t.CPtr, buf: t.CChar | t.CPtr, size: t.CSizeT):
    """Serialize AST node to buf (with capacity size)"""
    if node is None or buf is None or size == 0:
        return
    buf[0] = '\0'
    final_pos: t.CSizeT = node.dump(buf, size, 0)
    # NUL-terminate to prevent printf overflow
    if final_pos < size:
        buf[final_pos] = '\0'
    else:
        buf[size - 1] = '\0'