import t, c
from stdint import *
import memhub
import string
from .tokens import (
    Token, TokenType, TokOp, _init_tables, _kw_lookup, _op_head,
    OpEntry, new_token, token_set_str, token_set_str_literal,
)


# 缩进栈最大深度
LEXER_MAX_INDENT: t.CDefine = 256


# ============================================================
# Lexer 状态结构体
# ============================================================
class Lexer:
    src: str          # 源代码（NUL 结尾）
    pos: t.CSizeT                  # 当前位置
    len: t.CSizeT                  # 源代码长度
    lineno: t.CInt                 # 当前行号（从 1 开始）
    col: t.CInt                    # 当前列号（从 0 开始）
    pool: memhub.MemManager | t.CPtr     # 内存池
    indent_stack: t.CInt | t.CPtr  # 缩进栈（int 数组，从 pool 分配）
    indent_top: t.CInt             # 栈顶索引（-1 表示空栈）
    paren_depth: t.CInt            # 括号深度（>0 时不产生 TokenType.NewLine）
    tokens_head: Token | t.CPtr    # token 链表头
    tokens_tail: Token | t.CPtr    # token 链表尾
    at_line_start: t.CInt          # 1=行首（需处理缩进）
    error_count: t.CInt            # 错误计数

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


def new_lexer(pool: memhub.MemManager | t.CPtr) -> Lexer | t.CPtr:
    """从 pool 分配并返回 Lexer（跨模块安全包装）"""
    return Lexer(pool)


def _lexer_init(lx: Lexer | t.CPtr, src: str,
                pool: memhub.MemManager | t.CPtr):
    """初始化 Lexer"""
    if lx == None: return
    lx.src = src
    lx.pos = 0
    lx.len = string.strlen(src) if src != None else 0
    lx.lineno = 1
    lx.col = 0
    lx.pool = pool
    # 分配缩进栈
    lx.indent_stack = pool.alloc(LEXER_MAX_INDENT * 8)
    if lx.indent_stack != None:
        string.memset(lx.indent_stack, 0, LEXER_MAX_INDENT * 8)
    lx.indent_top = -1
    # 栈底压入 0（基础缩进）
    if lx.indent_stack != None:
        lx.indent_stack[0] = 0
        lx.indent_top = 0
    lx.paren_depth = 0
    lx.tokens_head = None
    lx.tokens_tail = None
    lx.at_line_start = 1
    lx.error_count = 0
    _init_tables(pool)


# ============================================================
# 辅助函数
# ============================================================

def _peek(lx: Lexer | t.CPtr, offset: t.CSizeT) -> t.CChar:
    """查看 pos+offset 处字符，越界返回 '\0'"""
    p: t.CSizeT = lx.pos + offset
    if p >= lx.len:
        return '\0'
    return lx.src[p]


def _peek_cur(lx: Lexer | t.CPtr) -> t.CChar:
    if lx.pos >= lx.len:
        return '\0'
    return lx.src[lx.pos]


def _advance(lx: Lexer | t.CPtr) -> t.CChar:
    """前进一个字符，返回该字符"""
    if lx.pos >= lx.len:
        return '\0'
    ch: t.CChar = lx.src[lx.pos]
    lx.pos += 1
    if ch == '\n':
        lx.lineno += 1
        lx.col = 0
    else:
        lx.col += 1
    return ch


def _is_alpha(ch: t.CChar) -> t.CInt:
    if ch >= 'a' and ch <= 'z': return 1
    if ch >= 'A' and ch <= 'Z': return 1
    if ch == '_': return 1
    return 0


def _is_digit(ch: t.CChar) -> t.CInt:
    if ch >= '0' and ch <= '9': return 1
    return 0


def _is_alnum(ch: t.CChar) -> t.CInt:
    if _is_alpha(ch): return 1
    if _is_digit(ch): return 1
    return 0


def _is_hex(ch: t.CChar) -> t.CInt:
    if ch >= '0' and ch <= '9': return 1
    if ch >= 'a' and ch <= 'f': return 1
    if ch >= 'A' and ch <= 'F': return 1
    return 0


def _is_space(ch: t.CChar) -> t.CInt:
    if ch == ' ': return 1
    if ch == '\t': return 1
    if ch == '\f': return 1
    if ch == '\v': return 1
    if ch == '\r': return 1
    return 0


def _emit(lx: Lexer | t.CPtr, tok: Token | t.CPtr):
    """将 token 追加到链表"""
    if tok == None: return
    if lx.tokens_head == None:
        lx.tokens_head = tok
        lx.tokens_tail = tok
    else:
        lx.tokens_tail.next = tok
        lx.tokens_tail = tok


def _make_op_token(lx: Lexer | t.CPtr, op_id: t.CInt, length: t.CInt,
                   lineno: t.CInt, col: t.CInt) -> Token | t.CPtr:
    tok: Token | t.CPtr = new_token(lx.pool, TokenType.Op, lineno, col)
    if tok == None: return None
    tok.op_subtype = op_id
    tok.end_col_offset = col + length
    return tok


# ============================================================
# 缩进处理
# ============================================================

def _handle_indent(lx: Lexer | t.CPtr):
    """处理行首缩进，生成 TokenType.Indent/TokenType.Dedent。
    跳过空行和注释行（不产生缩进变化）。"""
    # 先跳过空行和注释行
    while lx.pos < lx.len:
        # 记住行首位置，用于非空行回退
        line_start_pos: t.CSizeT = lx.pos
        line_start_col: t.CInt = lx.col
        # 跳过空白（不跨行）
        while lx.pos < lx.len:
            ch: t.CChar = lx.src[lx.pos]
            if ch == ' ' or ch == '\t' or ch == '\f' or ch == '\v':
                lx.pos += 1
                lx.col += 1
            elif ch == '\r':
                # CRLF 行尾符的 \r，跳过（让 \n 触发空行处理）
                lx.pos += 1
            else:
                break
        if lx.pos >= lx.len:
            break
        ch = lx.src[lx.pos]
        if ch == '\n':
            # 空行，跳过
            lx.pos += 1
            lx.lineno += 1
            lx.col = 0
            continue
        if ch == '\r':
            # CRLF 空行（\r 后跟 \n），跳过 \r
            lx.pos += 1
            if lx.pos < lx.len and lx.src[lx.pos] == '\n':
                lx.pos += 1
            lx.lineno += 1
            lx.col = 0
            continue
        if ch == '#':
            # 注释行，跳过到行尾
            while lx.pos < lx.len and lx.src[lx.pos] != '\n':
                lx.pos += 1
            continue
        # 非空行：回退到行首，让后续缩进计算能读到前导空格
        lx.pos = line_start_pos
        lx.col = line_start_col
        break

    if lx.pos >= lx.len:
        # 文件末尾，不再处理缩进
        lx.at_line_start = 0
        return

    # 计算当前行缩进
    indent: t.CInt = 0
    save_pos: t.CSizeT = lx.pos
    save_col: t.CInt = lx.col
    while lx.pos < lx.len:
        ch: t.CChar = lx.src[lx.pos]
        if ch == ' ':
            indent += 1
            lx.pos += 1
            lx.col += 1
        elif ch == '\t':
            indent += 8 - (indent % 8)
            lx.pos += 1
            lx.col += 1
        else:
            break

    # 如果行尾（空行），不处理
    if lx.pos < lx.len:
        ch = lx.src[lx.pos]
        if ch == '\n' or ch == '\r' or ch == '#':
            # 空行或注释行，回退位置让主循环处理
            lx.pos = save_pos
            lx.col = save_col
            lx.at_line_start = 0
            return

    # 与缩进栈比较
    top_indent: t.CInt = lx.indent_stack[lx.indent_top]
    if indent > top_indent:
        # TokenType.Indent
        if lx.indent_top + 1 >= LEXER_MAX_INDENT:
            lx.error_count += 1
            return
        lx.indent_top += 1
        lx.indent_stack[lx.indent_top] = indent
        tok: Token | t.CPtr = new_token(lx.pool, TokenType.Indent, lx.lineno, 0)
        _emit(lx, tok)
    else:
        while indent < lx.indent_stack[lx.indent_top] and lx.indent_top > 0:
            lx.indent_top -= 1
            tok: Token | t.CPtr = new_token(lx.pool, TokenType.Dedent, lx.lineno, 0)
            _emit(lx, tok)
    lx.at_line_start = 0


# ============================================================
# 标识符/关键字
# ============================================================

def _read_name(lx: Lexer | t.CPtr):
    start: t.CSizeT = lx.pos
    start_lineno: t.CInt = lx.lineno
    start_col: t.CInt = lx.col
    while lx.pos < lx.len:
        ch: t.CChar = lx.src[lx.pos]
        if _is_alnum(ch):
            lx.pos += 1
            lx.col += 1
        else:
            break
    length: t.CSizeT = lx.pos - start
    tok: Token | t.CPtr = new_token(lx.pool, TokenType.Name, start_lineno, start_col)
    if tok == None: return
    token_set_str(lx.pool, tok, lx.src, start, length)
    tok.end_lineno = lx.lineno
    tok.end_col_offset = start_col + t.CInt(length)
    # 查找关键字
    kw_id: t.CInt = _kw_lookup(tok.str_val)
    if kw_id:
        tok.kw_subtype = kw_id
    _emit(lx, tok)


# ============================================================
# 数字
# ============================================================

def _read_number(lx: Lexer | t.CPtr):
    start: t.CSizeT = lx.pos
    start_lineno: t.CInt = lx.lineno
    start_col: t.CInt = lx.col
    is_float: t.CInt = 0
    is_complex: t.CInt = 0

    ch0: t.CChar = lx.src[lx.pos]
    # 处理 0x/0o/0b 前缀
    if ch0 == '0' and lx.pos + 1 < lx.len:
        ch1: t.CChar = lx.src[lx.pos + 1]
        if ch1 == 'x' or ch1 == 'X':
            lx.pos += 2
            lx.col += 2
            while lx.pos < lx.len:
                ch: t.CChar = lx.src[lx.pos]
                if _is_hex(ch) or ch == '_':
                    lx.pos += 1
                    lx.col += 1
                else:
                    break
            _emit_number(lx, start, start_lineno, start_col, 0, 0, 0)
            return
        if ch1 == 'o' or ch1 == 'O':
            lx.pos += 2
            lx.col += 2
            while lx.pos < lx.len:
                ch: t.CChar = lx.src[lx.pos]
                if ch >= '0' and ch <= '7':
                    lx.pos += 1
                    lx.col += 1
                elif ch == '_':
                    lx.pos += 1
                    lx.col += 1
                else:
                    break
            _emit_number(lx, start, start_lineno, start_col, 0, 0, 0)
            return
        if ch1 == 'b' or ch1 == 'B':
            lx.pos += 2
            lx.col += 2
            while lx.pos < lx.len:
                ch: t.CChar = lx.src[lx.pos]
                if ch == '0' or ch == '1' or ch == '_':
                    lx.pos += 1
                    lx.col += 1
                else:
                    break
            _emit_number(lx, start, start_lineno, start_col, 0, 0, 0)
            return

    # 十进制整数/浮点
    while lx.pos < lx.len:
        ch: t.CChar = lx.src[lx.pos]
        if _is_digit(ch) or ch == '_':
            lx.pos += 1
            lx.col += 1
        else:
            break

    # 小数部分
    if lx.pos < lx.len and lx.src[lx.pos] == '.':
        is_float = 1
        lx.pos += 1
        lx.col += 1
        while lx.pos < lx.len:
            ch: t.CChar = lx.src[lx.pos]
            if _is_digit(ch) or ch == '_':
                lx.pos += 1
                lx.col += 1
            else:
                break

    # 指数部分
    if lx.pos < lx.len:
        ch: t.CChar = lx.src[lx.pos]
        if ch == 'e' or ch == 'E':
            is_float = 1
            lx.pos += 1
            lx.col += 1
            if lx.pos < lx.len:
                ch = lx.src[lx.pos]
                if ch == '+' or ch == '-':
                    lx.pos += 1
                    lx.col += 1
            while lx.pos < lx.len:
                ch: t.CChar = lx.src[lx.pos]
                if _is_digit(ch) or ch == '_':
                    lx.pos += 1
                    lx.col += 1
                else:
                    break

    # 虚数后缀
    if lx.pos < lx.len and lx.src[lx.pos] == 'j':
        is_complex = 1
        is_float = 1
        lx.pos += 1
        lx.col += 1

    _emit_number(lx, start, start_lineno, start_col, is_float, is_complex, 0)


def _emit_number(lx: Lexer | t.CPtr, start: t.CSizeT,
                 start_lineno: t.CInt, start_col: t.CInt,
                 is_float: t.CInt, is_complex: t.CInt, is_hex: t.CInt):
    length: t.CSizeT = lx.pos - start
    tok: Token | t.CPtr = new_token(lx.pool, TokenType.Number, start_lineno, start_col)
    if tok == None: return
    token_set_str(lx.pool, tok, lx.src, start, length)
    tok.is_float = is_float
    tok.is_complex = is_complex
    tok.end_lineno = lx.lineno
    tok.end_col_offset = start_col + t.CInt(length)
    # 解析数值
    if is_float:
        tok.float_val = string.atof(tok.str_val)
    else:
        tok.int_val = string.atoll(tok.str_val)
    _emit(lx, tok)


# ============================================================
# 字符串
# ============================================================

def _read_string(lx: Lexer | t.CPtr):
    """读取字符串字面量。
    支持前缀: r, b, u, f, rb, br, rf, fr（大小写不敏感）
    支持引号: ', ", ''', """
    start_lineno: t.CInt = lx.lineno
    start_col: t.CInt = lx.col
    prefix_start: t.CSizeT = lx.pos

    # 读取前缀（r/b/u/f）
    is_raw: t.CInt = 0
    is_bytes: t.CInt = 0
    is_fstring: t.CInt = 0
    while lx.pos < lx.len:
        ch: t.CChar = lx.src[lx.pos]
        if ch == 'r' or ch == 'R':
            is_raw = 1
            lx.pos += 1
            lx.col += 1
        elif ch == 'b' or ch == 'B':
            is_bytes = 1
            lx.pos += 1
            lx.col += 1
        elif ch == 'u' or ch == 'U':
            lx.pos += 1
            lx.col += 1
        elif ch == 'f' or ch == 'F':
            is_fstring = 1
            lx.pos += 1
            lx.col += 1
        else:
            break

    if lx.pos >= lx.len:
        lx.error_count += 1
        return

    # 确定引号
    quote: t.CChar = lx.src[lx.pos]
    if quote != '\'' and quote != '"':
        lx.error_count += 1
        return

    # 检查三引号
    is_triple: t.CInt = 0
    if lx.pos + 2 < lx.len:
        if lx.src[lx.pos] == quote and lx.src[lx.pos + 1] == quote and lx.src[lx.pos + 2] == quote:
            is_triple = 1

    content_start: t.CSizeT = lx.pos
    if is_triple:
        lx.pos += 3
        lx.col += 3
    else:
        lx.pos += 1
        lx.col += 1

    # 读取字符串内容
    content_buf_start: t.CSizeT = lx.pos
    while lx.pos < lx.len:
        ch = lx.src[lx.pos]
        if is_triple:
            if ch == quote:
                if lx.pos + 2 < lx.len:
                    if lx.src[lx.pos + 1] == quote and lx.src[lx.pos + 2] == quote:
                        break
                elif lx.pos + 2 == lx.len:
                    break
            if ch == '\\':
                lx.pos += 1
                lx.col += 1
                if lx.pos < lx.len:
                    if lx.src[lx.pos] == '\n':
                        lx.lineno += 1
                        lx.col = 0
                    lx.pos += 1
                    lx.col += 1
                continue
            if ch == '\n':
                lx.lineno += 1
                lx.col = 0
            lx.pos += 1
            lx.col += 1
        else:
            if ch == quote:
                break
            if ch == '\n':
                # 单引号字符串不能跨行（除非续行）
                lx.error_count += 1
                return
            if ch == '\\' and not is_raw:
                lx.pos += 1
                lx.col += 1
                if lx.pos < lx.len:
                    lx.pos += 1
                    lx.col += 1
                continue
            lx.pos += 1
            lx.col += 1

    content_end: t.CSizeT = lx.pos
    # 跳过结束引号
    if is_triple:
        if lx.pos + 2 < lx.len:
            lx.pos += 3
            lx.col += 3
        else:
            lx.pos = lx.len
    else:
        if lx.pos < lx.len:
            lx.pos += 1
            lx.col += 1

    # 创建 token
    tok: Token | t.CPtr = new_token(lx.pool, TokenType.String, start_lineno, start_col)
    if tok == None: return
    content_len: t.CSizeT = content_end - content_buf_start

    # 处理转义序列（非 raw 字符串）
    if is_raw == 0 and content_len > 0:
        # 分配转换缓冲区（最大情况：每个字符都是普通字符，无需扩展）
        conv_buf: str = lx.pool.alloc(content_len + 1)
        if conv_buf == None:
            token_set_str(lx.pool, tok, lx.src, content_buf_start, content_len)
        else:
            wpos: t.CSizeT = 0
            rpos: t.CSizeT = 0
            while rpos < content_len:
                ch: t.CChar = lx.src[content_buf_start + rpos]
                if ch == '\\' and rpos + 1 < content_len:
                    rpos += 1
                    next_ch: t.CChar = lx.src[content_buf_start + rpos]
                    if next_ch == 'n':
                        conv_buf[wpos] = '\n'
                    elif next_ch == 't':
                        conv_buf[wpos] = '\t'
                    elif next_ch == 'r':
                        conv_buf[wpos] = '\r'
                    elif next_ch == '\\':
                        conv_buf[wpos] = '\\'
                    elif next_ch == '\'':
                        conv_buf[wpos] = '\''
                    elif next_ch == '"':
                        conv_buf[wpos] = '"'
                    elif next_ch == '0':
                        conv_buf[wpos] = '\0'
                    else:
                        # 未知转义，保留原样（包括 \xNN 等以后再处理）
                        conv_buf[wpos] = '\\'
                        wpos += 1
                        conv_buf[wpos] = next_ch
                    wpos += 1
                    rpos += 1
                else:
                    conv_buf[wpos] = ch
                    wpos += 1
                    rpos += 1
            conv_buf[wpos] = '\0'
            token_set_str_literal(lx.pool, tok, conv_buf)
    else:
        token_set_str(lx.pool, tok, lx.src, content_buf_start, content_len)
    tok.end_lineno = lx.lineno
    tok.end_col_offset = lx.col
    # flags 字段复用：bit0=raw, bit1=bytes, bit2=fstring, bit3=triple
    if is_raw: tok.kw_subtype = tok.kw_subtype | 1
    if is_bytes: tok.kw_subtype = tok.kw_subtype | 2
    if is_fstring: tok.kw_subtype = tok.kw_subtype | 4
    if is_triple: tok.kw_subtype = tok.kw_subtype | 8
    _emit(lx, tok)


# ============================================================
# 运算符
# ============================================================

def _read_operator(lx: Lexer | t.CPtr):
    """贪婪匹配运算符"""
    start_lineno: t.CInt = lx.lineno
    start_col: t.CInt = lx.col
    remaining: t.CSizeT = lx.len - lx.pos

    # 遍历运算符表，找最长匹配
    best_entry: OpEntry | t.CPtr = None
    best_len: t.CInt = 0
    cur: OpEntry | t.CPtr = _op_head
    while cur != None:
        if cur.length <= t.CInt(remaining):
            # 比较 cur.name 与 src[pos:pos+cur.length]
            match: t.CInt = 1
            i: t.CSizeT = 0
            while i < t.CSizeT(cur.length):
                if lx.src[lx.pos + i] != cur.name[i]:
                    match = 0
                    break
                i += 1
            if match:
                if cur.length > best_len:
                    best_len = cur.length
                    best_entry = cur
        cur = cur.next

    if best_entry == None:
        # 未知字符，报错并跳过
        lx.error_count += 1
        lx.pos += 1
        lx.col += 1
        return

    # 更新括号深度
    if best_entry.op_id == TokOp.LPar or best_entry.op_id == TokOp.Lsqb or best_entry.op_id == TokOp.LBrace:
        lx.paren_depth += 1
    elif best_entry.op_id == TokOp.RPar or best_entry.op_id == TokOp.Rsqb or best_entry.op_id == TokOp.RBrace:
        if lx.paren_depth > 0:
            lx.paren_depth -= 1

    tok: Token | t.CPtr = _make_op_token(lx, best_entry.op_id, best_entry.length,
                                          start_lineno, start_col)
    if tok == None: return
    token_set_str(lx.pool, tok, lx.src, lx.pos, t.CSizeT(best_len))
    lx.pos += t.CSizeT(best_len)
    lx.col += best_len
    _emit(lx, tok)


# ============================================================
# 主词法分析入口
# ============================================================

def tokenize(lx: Lexer | t.CPtr) -> Token | t.CPtr:
    """执行词法分析，返回 token 链表头"""
    if lx == None: return None
    while lx.pos < lx.len:
        # 行首缩进处理
        if lx.at_line_start and lx.paren_depth == 0:
            _handle_indent(lx)
            if lx.pos >= lx.len:
                break

        ch: t.CChar = _peek_cur(lx)
        if ch == '\0':
            break

        # 换行
        if ch == '\n':
            if lx.paren_depth > 0:
                # 括号内，产生 TokenType.Nl
                tok: Token | t.CPtr = new_token(lx.pool, TokenType.Nl, lx.lineno, lx.col)
                _emit(lx, tok)
            else:
                # 逻辑换行
                tok: Token | t.CPtr = new_token(lx.pool, TokenType.NewLine, lx.lineno, lx.col)
                _emit(lx, tok)
                lx.at_line_start = 1
            lx.pos += 1
            lx.lineno += 1
            lx.col = 0
            continue

        # 反斜线续行
        if ch == '\\' and lx.pos + 1 < lx.len and lx.src[lx.pos + 1] == '\n':
            lx.pos += 2
            lx.lineno += 1
            lx.col = 0
            continue

        # 空白
        if _is_space(ch):
            lx.pos += 1
            lx.col += 1
            continue

        # 注释
        if ch == '#':
            while lx.pos < lx.len and lx.src[lx.pos] != '\n':
                lx.pos += 1
            continue

        # 字符串前缀（如 r"..."）或字符串
        if ch == '\'' or ch == '"':
            _read_string(lx)
            continue
        # 字符串前缀字母后跟引号（必须在 _is_alpha 之前检测，否则 f/r/b/u 会被当作 TokenType.Name）
        if (ch == 'r' or ch == 'R' or ch == 'b' or ch == 'B' or
            ch == 'u' or ch == 'U' or ch == 'f' or ch == 'F'):
            if lx.pos + 1 < lx.len:
                ch1: t.CChar = lx.src[lx.pos + 1]
                if ch1 == '\'' or ch1 == '"':
                    _read_string(lx)
                    continue
                # 双前缀如 rb, br, rf, fr
                if lx.pos + 2 < lx.len:
                    ch2: t.CChar = lx.src[lx.pos + 2]
                    if ((ch1 == 'r' or ch1 == 'R' or ch1 == 'b' or ch1 == 'B' or
                         ch1 == 'f' or ch1 == 'F' or ch1 == 'u' or ch1 == 'U') and
                        (ch2 == '\'' or ch2 == '"')):
                        _read_string(lx)
                        continue

        # 标识符/关键字
        if _is_alpha(ch):
            _read_name(lx)
            continue

        # 数字
        if _is_digit(ch):
            _read_number(lx)
            continue

        # 运算符
        _read_operator(lx)

    # 文件结尾：产生 TokenType.NewLine（如果有内容）
    if lx.tokens_tail != None:
        if lx.tokens_tail.type != TokenType.NewLine and lx.tokens_tail.type != TokenType.Nl:
            tok: Token | t.CPtr = new_token(lx.pool, TokenType.NewLine, lx.lineno, lx.col)
            _emit(lx, tok)

    # 产生剩余 TokenType.Dedent
    while lx.indent_top > 0:
        lx.indent_top -= 1
        tok: Token | t.CPtr = new_token(lx.pool, TokenType.Dedent, lx.lineno, 0)
        _emit(lx, tok)

    # TokenType.EndMarker
    tok: Token | t.CPtr = new_token(lx.pool, TokenType.EndMarker, lx.lineno, 0)
    _emit(lx, tok)

    return lx.tokens_head