import t, c
from stdint import *
import memhub
import string


# ============================================================
# Token 类型枚举（参考 CPython Lib/token.py）
# 用 t.CEnum 组织为类型安全枚举；旧名 XXX 保留为兼容别名
# ============================================================

class TokenType(t.CEnum):
    EndMarker:      t.State
    NewLine:        t.State
    Indent:         t.State
    Dedent:         t.State
    Name:           t.State
    Number:         t.State
    String:         t.State
    FStringStart:   t.State
    FStringMiddle:  t.State
    FStringEnd:     t.State
    Op:             t.State
    Nl:             t.State
    Comment:        t.State
    Encoding:       t.State


# ============================================================
# 关键字枚举（参考 CPython Lib/keyword.py）
# 用 t.CEnum 组织为类型安全枚举；旧名 KW_XXX 保留为兼容别名
# NotKeyword(0) 表示非关键字（哨兵，避免 False_=0 与 kw_subtype=0 冲突）
# ============================================================

class Keyword(t.CEnum):
    NotKeyword: t.State
    False_:     t.State
    None_:      t.State
    True_:      t.State
    And:        t.State
    As:         t.State
    Assert:     t.State
    Async:      t.State
    Await:      t.State
    Break:      t.State
    Class:      t.State
    Continue:   t.State
    Def:        t.State
    Del:        t.State
    Elif:       t.State
    Else:       t.State
    Except:     t.State
    Finally:    t.State
    For:        t.State
    From:       t.State
    Global:     t.State
    If:         t.State
    Import:     t.State
    In:         t.State
    Is:         t.State
    Lambda:     t.State
    Nonlocal:   t.State
    Not:        t.State
    Or:         t.State
    Pass:       t.State
    Raise:      t.State
    Return:     t.State
    Try:        t.State
    While:      t.State
    With:       t.State
    Yield:      t.State
    Match:      t.State
    Case:       t.State
    Type:       t.State


# ============================================================
# 运算符/分隔符枚举
# 用 t.CEnum 组织为类型安全枚举；旧名 OP_XXX 保留为兼容别名
# 存储在 token.op_subtype 字段
# ============================================================

class TokOp(t.CEnum):
    LPar:        t.State     # (
    RPar:        t.State     # )
    Lsqb:        t.State     # [
    Rsqb:        t.State     # ]
    LBrace:      t.State     # {
    RBrace:      t.State     # }
    Comma:       t.State     # ,
    Colon:       t.State     # :
    Dot:         t.State     # .
    Semi:        t.State    # ;
    At:          t.State    # @
    Equal:       t.State    # =
    RArrow:      t.State    # ->
    PlusEq:      t.State    # +=
    MinusEq:     t.State    # -=
    StarEq:      t.State    # *=
    SlashEq:     t.State    # /=
    DSlashEq:    t.State    # //=
    PercentEq:   t.State    # %=
    AtEq:        t.State    # @=
    AmpEq:       t.State    # &=
    VBarEq:      t.State    # |=
    CaretEq:     t.State    # ^=
    GtGtEq:      t.State    # >>=
    LtLtEq:      t.State    # <<=
    StarEqEq:    t.State    # **=
    DSlash:      t.State    # //
    StarStar:    t.State    # **
    LtLt:        t.State    # <<
    GtGt:        t.State    # >>
    LessEq:      t.State    # <=
    GreaterEq:   t.State    # >=
    EqEq:        t.State    # ==
    ExclaimEq:   t.State    # !=
    Less:        t.State    # <
    Greater:     t.State    # >
    Plus:        t.State    # +
    Minus:       t.State    # -
    Star:        t.State    # *
    Slash:       t.State    # /
    Percent:     t.State    # %
    Amp:         t.State    # &
    VBar:        t.State    # |
    Caret:       t.State    # ^
    Tilde:       t.State    # ~
    ColonEq:     t.State    # :=
    Ellipsis:    t.State    # ...
    Bang:        t.State    # ! (非标准，用于解析容错)


# ============================================================
# Token 结构体
# ============================================================
class Token:
    type: t.CInt               # Token 类型（TokenType.EndMarker/TokenType.NewLine/TokenType.Name/...）
    op_subtype: t.CInt         # 运算符子类型（OP_*），仅 type==TokenType.Op 时有效
    kw_subtype: t.CInt         # 关键字子类型（KW_*），仅 type==TokenType.Name 且是关键字时有效
    str_val: str  # token 文本（NUL 结尾，从 pool 分配）
    int_val: t.CInt64T         # 整数值（TokenType.Number 用）
    float_val: t.CDouble       # 浮点值（TokenType.Number 用）
    is_float: t.CInt           # 1=浮点数, 0=整数
    is_complex: t.CInt         # 1=复数, 0=实数
    lineno: t.CInt             # 起始行号
    col_offset: t.CInt         # 起始列偏移
    end_lineno: t.CInt         # 结束行号
    end_col_offset: t.CInt     # 结束列偏移
    next: Token | t.CPtr       # 链表下一个 token

    def __new__(self, pool: memhub.MemManager | t.CPtr):
        """从 mpool 分配 Token 结构体内存并零初始化"""
        ptr: Token | t.CPtr = pool.alloc(Token.__sizeof__())
        if ptr:
            string.memset(ptr, 0, Token.__sizeof__())
        return ptr


# ============================================================
# 关键字查找表
# 由于 TransPyC 无 dict 字面量，用链表实现
# ============================================================
class KwEntry:
    name: str
    kw_id: t.CInt
    next: KwEntry | t.CPtr

    def __new__(self, pool: memhub.MemManager | t.CPtr):
        ptr: KwEntry | t.CPtr = pool.alloc(KwEntry.__sizeof__())
        if ptr:
            string.memset(ptr, 0, KwEntry.__sizeof__())
        return ptr


# 全局关键字链表头（懒初始化）
_kw_head: KwEntry | t.CPtr = None
_kw_pool: memhub.MemManager | t.CPtr = None


def _kw_intern(pool: memhub.MemManager | t.CPtr, name: str, kw_id: t.CInt):
    """将关键字名注册到链表"""
    global _kw_head
    entry: KwEntry | t.CPtr = KwEntry(pool)
    if entry == None: return
    slen: t.CSizeT = string.strlen(name)
    buf: str = pool.alloc(slen + 1)
    if buf == None: return
    string.memcpy(buf, name, slen + 1)
    entry.name = buf
    entry.kw_id = kw_id
    entry.next = _kw_head
    _kw_head = entry


def _kw_lookup(name: str) -> t.CInt:
    """查找关键字，返回 kw_id（0 表示非关键字）"""
    cur: KwEntry | t.CPtr = _kw_head
    while cur != None:
        if string.strcmp(cur.name, name) == 0:
            return cur.kw_id
        cur = cur.next
    return 0


def _init_keywords(pool: memhub.MemManager | t.CPtr):
    """初始化关键字表（仅一次）"""
    global _kw_head, _kw_pool
    if _kw_head != None: return
    _kw_pool = pool
    _kw_intern(pool, "False", Keyword.False_)
    _kw_intern(pool, "None", Keyword.None_)
    _kw_intern(pool, "True", Keyword.True_)
    _kw_intern(pool, "and", Keyword.And)
    _kw_intern(pool, "as", Keyword.As)
    _kw_intern(pool, "assert", Keyword.Assert)
    _kw_intern(pool, "async", Keyword.Async)
    _kw_intern(pool, "await", Keyword.Await)
    _kw_intern(pool, "break", Keyword.Break)
    _kw_intern(pool, "class", Keyword.Class)
    _kw_intern(pool, "continue", Keyword.Continue)
    _kw_intern(pool, "def", Keyword.Def)
    _kw_intern(pool, "del", Keyword.Del)
    _kw_intern(pool, "elif", Keyword.Elif)
    _kw_intern(pool, "else", Keyword.Else)
    _kw_intern(pool, "except", Keyword.Except)
    _kw_intern(pool, "finally", Keyword.Finally)
    _kw_intern(pool, "for", Keyword.For)
    _kw_intern(pool, "from", Keyword.From)
    _kw_intern(pool, "global", Keyword.Global)
    _kw_intern(pool, "if", Keyword.If)
    _kw_intern(pool, "import", Keyword.Import)
    _kw_intern(pool, "in", Keyword.In)
    _kw_intern(pool, "is", Keyword.Is)
    _kw_intern(pool, "lambda", Keyword.Lambda)
    _kw_intern(pool, "nonlocal", Keyword.Nonlocal)
    _kw_intern(pool, "not", Keyword.Not)
    _kw_intern(pool, "or", Keyword.Or)
    _kw_intern(pool, "pass", Keyword.Pass)
    _kw_intern(pool, "raise", Keyword.Raise)
    _kw_intern(pool, "return", Keyword.Return)
    _kw_intern(pool, "try", Keyword.Try)
    _kw_intern(pool, "while", Keyword.While)
    _kw_intern(pool, "with", Keyword.With)
    _kw_intern(pool, "yield", Keyword.Yield)
    _kw_intern(pool, "match", Keyword.Match)
    _kw_intern(pool, "case", Keyword.Case)
    _kw_intern(pool, "type", Keyword.Type)


# ============================================================
# 运算符查找表
# 由于 TransPyC 无 dict，用数组+线性查找
# 运算符按长度降序排列，以便贪婪匹配
# ============================================================
class OpEntry:
    name: str   # 运算符文本（如 "==", "+="）
    op_id: t.CInt
    length: t.CInt
    next: OpEntry | t.CPtr

    def __new__(self, pool: memhub.MemManager | t.CPtr):
        ptr: OpEntry | t.CPtr = pool.alloc(OpEntry.__sizeof__())
        if ptr:
            string.memset(ptr, 0, OpEntry.__sizeof__())
        return ptr


_op_head: OpEntry | t.CPtr = None


def _op_intern(pool: memhub.MemManager | t.CPtr, name: str, op_id: t.CInt):
    global _op_head
    slen: t.CSizeT = string.strlen(name)
    entry: OpEntry | t.CPtr = OpEntry(pool)
    if entry == None: return
    buf: str = pool.alloc(slen + 1)
    if buf == None: return
    string.memcpy(buf, name, slen + 1)
    entry.name = buf
    entry.op_id = op_id
    entry.length = t.CInt(slen)
    entry.next = _op_head
    _op_head = entry


def _init_operators(pool: memhub.MemManager | t.CPtr):
    """初始化运算符表（按长度降序注册，便于贪婪匹配）"""
    global _op_head
    if _op_head != None: return
    # 3 字符运算符
    _op_intern(pool, "**=", TokOp.StarEqEq)
    _op_intern(pool, "//=", TokOp.DSlashEq)
    _op_intern(pool, ">>=", TokOp.GtGtEq)
    _op_intern(pool, "<<=", TokOp.LtLtEq)
    _op_intern(pool, "...", TokOp.Ellipsis)
    # 2 字符运算符
    _op_intern(pool, "->", TokOp.RArrow)
    _op_intern(pool, "+=", TokOp.PlusEq)
    _op_intern(pool, "-=", TokOp.MinusEq)
    _op_intern(pool, "*=", TokOp.StarEq)
    _op_intern(pool, "/=", TokOp.SlashEq)
    _op_intern(pool, "%=", TokOp.PercentEq)
    _op_intern(pool, "@=", TokOp.AtEq)
    _op_intern(pool, "&=", TokOp.AmpEq)
    _op_intern(pool, "|=", TokOp.VBarEq)
    _op_intern(pool, "^=", TokOp.CaretEq)
    _op_intern(pool, "//", TokOp.DSlash)
    _op_intern(pool, "**", TokOp.StarStar)
    _op_intern(pool, "<<", TokOp.LtLt)
    _op_intern(pool, ">>", TokOp.GtGt)
    _op_intern(pool, "<=", TokOp.LessEq)
    _op_intern(pool, ">=", TokOp.GreaterEq)
    _op_intern(pool, "==", TokOp.EqEq)
    _op_intern(pool, "!=", TokOp.ExclaimEq)
    _op_intern(pool, ":=", TokOp.ColonEq)
    # 1 字符运算符
    _op_intern(pool, "(", TokOp.LPar)
    _op_intern(pool, ")", TokOp.RPar)
    _op_intern(pool, "[", TokOp.Lsqb)
    _op_intern(pool, "]", TokOp.Rsqb)
    _op_intern(pool, "{", TokOp.LBrace)
    _op_intern(pool, "}", TokOp.RBrace)
    _op_intern(pool, ",", TokOp.Comma)
    _op_intern(pool, ":", TokOp.Colon)
    _op_intern(pool, ".", TokOp.Dot)
    _op_intern(pool, ";", TokOp.Semi)
    _op_intern(pool, "@", TokOp.At)
    _op_intern(pool, "=", TokOp.Equal)
    _op_intern(pool, "<", TokOp.Less)
    _op_intern(pool, ">", TokOp.Greater)
    _op_intern(pool, "+", TokOp.Plus)
    _op_intern(pool, "-", TokOp.Minus)
    _op_intern(pool, "*", TokOp.Star)
    _op_intern(pool, "/", TokOp.Slash)
    _op_intern(pool, "%", TokOp.Percent)
    _op_intern(pool, "&", TokOp.Amp)
    _op_intern(pool, "|", TokOp.VBar)
    _op_intern(pool, "^", TokOp.Caret)
    _op_intern(pool, "~", TokOp.Tilde)


def _init_tables(pool: memhub.MemManager | t.CPtr):
    """初始化所有查找表（幂等）"""
    _init_keywords(pool)
    _init_operators(pool)


# ============================================================
# Token 创建辅助函数
# ============================================================

def new_token(pool: memhub.MemManager | t.CPtr, ttype: t.CInt, lineno: t.CInt,
              col_offset: t.CInt) -> Token | t.CPtr:
    """创建一个空 Token，调用方填充字段"""
    tok: Token | t.CPtr = Token(pool)
    if tok == None: return None
    tok.type = ttype
    tok.lineno = lineno
    tok.col_offset = col_offset
    tok.end_lineno = lineno
    tok.end_col_offset = col_offset
    return tok


def token_set_str(pool: memhub.MemManager | t.CPtr, tok: Token | t.CPtr,
                  src: str, start: t.CSizeT, length: t.CSizeT):
    """从 src[start:start+length] 复制文本到 token.str_val

    注意：length == 0 时也分配 1 字节存 NUL，让空字符串与 NULL 指针区分开。
    """
    if tok == None: return
    buf: str = pool.alloc(length + 1)
    if buf == None: return
    i: t.CSizeT = 0
    while i < length:
        buf[i] = src[start + i]
        i += 1
    buf[length] = '\0'
    tok.str_val = buf


def token_set_str_literal(pool: memhub.MemManager | t.CPtr, tok: Token | t.CPtr,
                          text: str):
    """直接复制一个 C 字符串到 token.str_val"""
    if tok == None or text == None: return
    slen: t.CSizeT = string.strlen(text)
    buf: str = pool.alloc(slen + 1)
    if buf == None: return
    string.memcpy(buf, text, slen + 1)
    tok.str_val = buf