"""vast.utils - AST 遍历与序列化工具

与 CPython ``ast`` 模块辅助函数同构：

- ``iter_fields`` / ``iter_child_nodes`` / ``walk``
- ``dump``（含 annotate_fields / include_attributes / indent）
- ``copy_location`` / ``fix_missing_locations`` / ``get_source_segment``
- ``literal_eval``（常量求值）
"""
from __future__ import annotations

from lib.vast.Nodes import AST
from lib.vast.Tokenizer import Token


def iter_fields(node: AST):
    """按 _fields 顺序产出 (字段名, 值)。"""
    for field in node._fields:
        try:
            yield field, getattr(node, field)
        except AttributeError:
            pass


def iter_child_nodes(node: AST):
    """产出全部子节点（递归扁平，跳过 None 与空列表）。"""
    for _name, field in iter_fields(node):
        if isinstance(field, AST):
            yield field
        elif isinstance(field, list):
            for item in field:
                if isinstance(item, AST):
                    yield item


def walk(node: AST):
    """深度优先遍历，产出自身与全部后代节点。"""
    from collections import deque

    todo: deque[AST] = deque([node])
    while todo:
        n: AST = todo.popleft()
        yield n
        todo.extend(iter_child_nodes(n))


def dump(node: AST, annotate_fields: bool = True, include_attributes: bool = False,
         indent: int | str | None = None) -> str:
    """将 AST 节点格式化为字符串（精确复刻 CPython 3.12 ast.dump）。

    - annotate_fields: 是否显示字段名（默认 True）
    - include_attributes: 是否附加 lineno/col_offset 等位置属性
    - indent: 缩进（int 或字符串），None 为单行紧凑输出

    与 CPython 一致的行为：
    - 值为 None 且类默认值为 None 的字段跳过（如 FunctionDef.returns）
    - 节点字段全部为"简单值"且 ≤3 个时单行输出
    """
    if not isinstance(node, AST):
        raise TypeError(f'expected AST, got {type(node).__name__!r}')

    if indent is not None and not isinstance(indent, str):
        indent = ' ' * indent

    def _format(node: object, level: int = 0):
        if indent is not None:
            level += 1
            prefix: str = '\n' + indent * level  # type: ignore[operator]
            sep: str = ',\n' + indent * level  # type: ignore[operator]
        else:
            prefix = ''
            sep = ', '
        if isinstance(node, AST):
            cls: type = type(node)
            args: list[str] = []
            allsimple: bool = True
            keywords: bool = annotate_fields
            for name in node._fields:
                try:
                    value = getattr(node, name)
                except AttributeError:
                    keywords = True
                    continue
                if value is None and getattr(cls, '_node_defaults', {}).get(name, ...) is None:
                    keywords = True
                    continue
                value, simple = _format(value, level)
                allsimple = allsimple and simple
                if keywords:
                    args.append(f'{name}={value}')
                else:
                    args.append(value)
            if include_attributes and node._attributes:
                for name in node._attributes:
                    try:
                        value = getattr(node, name)
                    except AttributeError:
                        continue
                    if value is None and getattr(cls, '_node_defaults', {}).get(name, ...) is None:
                        continue
                    value, simple = _format(value, level)
                    allsimple = allsimple and simple
                    args.append(f'{name}={value}')
            if allsimple and len(args) <= 3:
                return f'{type(node).__name__}({", ".join(args)})', not args
            return f'{type(node).__name__}({prefix}{sep.join(args)})', False
        elif isinstance(node, list):
            if not node:
                return '[]', True
            return f'[{prefix}{sep.join(_format(x, level)[0] for x in node)}]', False
        return repr(node), True

    return _format(node)[0]


def copy_location(new_node: AST, old_node: AST) -> AST:
    """复制 old_node 的位置属性到 new_node（与 ast.copy_location 一致）。"""
    for attr in ('lineno', 'col_offset', 'end_lineno', 'end_col_offset'):
        try:
            setattr(new_node, attr, getattr(old_node, attr))
        except AttributeError:
            pass
    return new_node


def fix_missing_locations(node: AST) -> AST:
    """为缺失位置的节点补全（从最近的祖先复制；递归处理）。"""
    def _fix(n: AST, parent_loc: dict[str, int] | None) -> None:
        has_loc: bool = hasattr(n, 'lineno') and getattr(n, 'lineno', None) is not None
        if hasattr(n, 'lineno') and getattr(n, 'lineno', None) is None and parent_loc:
            for attr, val in parent_loc.items():
                setattr(n, attr, val)
        loc: dict[str, int] | None = None
        if hasattr(n, 'lineno') and getattr(n, 'lineno', None) is not None:
            loc = {a: getattr(n, a) for a in
                   ('lineno', 'col_offset', 'end_lineno', 'end_col_offset')
                   if getattr(n, a, None) is not None}
        for child in iter_child_nodes(n):
            _fix(child, loc)

    _fix(node, None)
    return node


def get_source_segment(source: str, node: AST, padded: bool = False) -> str | None:
    """提取节点对应的源码片段（与 ast.get_source_segment 一致）。"""
    try:
        lineno: int | None = getattr(node, 'lineno', None)
        end_lineno: int | None = getattr(node, 'end_lineno', None)
        col_offset: int | None = getattr(node, 'col_offset', None)
        end_col_offset: int | None = getattr(node, 'end_col_offset', None)
        if lineno is None or end_lineno is None or col_offset is None or end_col_offset is None:
            return None
        lines: list[str] = source.splitlines(keepends=True)
        if lineno < 1 or lineno > len(lines):
            return None
        first: str = lines[lineno - 1]
        # col_offset 是 UTF-8 字节偏移 → 转字符偏移
        start_char: int = _ByteToChar(first, col_offset)
        if lineno == end_lineno:
            end_char: int = _ByteToChar(first, end_col_offset)
            segment: str = first[start_char:end_char]
        else:
            last: str = lines[end_lineno - 1]
            end_char = _ByteToChar(last, end_col_offset)
            middle: list[str] = lines[lineno:end_lineno - 1]
            segment = first[start_char:] + ''.join(middle) + last[:end_char]
        if padded:
            return segment
        return segment
    except Exception:
        return None


def _ByteToChar(line: str, byte_offset: int) -> int:
    """UTF-8 字节偏移 → 字符偏移。"""
    b: bytes = line.encode('utf-8')
    return len(b[:byte_offset].decode('utf-8', errors='replace')) if byte_offset > 0 else 0


def literal_eval(node_or_string: AST | str):
    """安全求值常量表达式（与 ast.literal_eval 语义一致）。

    支持常量、元组/列表/集合/字典、一元 +/-、BinOp 算术、切片。
    """
    if isinstance(node_or_string, str):
        from lib.vast.Tokenizer import Tokenizer

        toks: list[Token] = Tokenizer(node_or_string).Tokenize()
        if toks and toks[0].Type == 'NEWLINE':
            raise ValueError('malformed node or string')
        from lib.vast.Parser import _ComputeLineStarts

        from lib.vast.ExprParser import ExprParserMixin

        class _EV(ExprParserMixin):
            def __init__(self) -> None:
                super().__init__(toks, node_or_string, _ComputeLineStarts(node_or_string))

        parser = _EV()
        node: AST = parser.ParseExpr()
    elif isinstance(node_or_string, AST):
        node = node_or_string
    else:
        raise TypeError('bad argument type')

    return _EvalLiteral(node)


def _EvalLiteral(node: AST):
    """递归常量求值核心。"""
    if isinstance(node, AST):
        pass
    name: str = type(node).__name__
    # 常量
    if name == 'Constant':
        value = node.value
        if isinstance(value, (int, float, complex, str, bool, bytes, type(None), Ellipsis)):  # type: ignore[comparison-overlap]
            return value
        raise ValueError('malformed node or string: unsupported constant')
    if name == 'Tuple':
        return tuple(_EvalLiteral(e) for e in node.elts)
    if name == 'List':
        return [_EvalLiteral(e) for e in node.elts]
    if name == 'Set':
        return {_EvalLiteral(e) for e in node.elts}
    if name == 'Dict':
        if len(node.keys) != len(node.values):
            raise ValueError('malformed node or string: Dict')
        out: dict = {}
        for k, v in zip(node.keys, node.values):
            if k is None:
                raise ValueError('malformed node or string: ** in Dict')
            out[_EvalLiteral(k)] = _EvalLiteral(v)
        return out
    if name == 'UnaryOp':
        operand = _EvalLiteral(node.operand)
        if isinstance(node.op, (__import__('lib.vast.Nodes', fromlist=['USub']).USub,)):
            return -operand
        if isinstance(node.op, (__import__('lib.vast.Nodes', fromlist=['UAdd']).UAdd,)):
            return +operand
        if isinstance(node.op, (__import__('lib.vast.Nodes', fromlist=['Not']).Not,)):
            return not operand
        raise ValueError('malformed node or string: unary op')
    if name == 'BinOp':
        from lib.vast.Nodes import Add, Sub, Mult, Div, FloorDiv, Mod, Pow, LShift, RShift, BitOr, BitXor, BitAnd

        left = _EvalLiteral(node.left)
        right = _EvalLiteral(node.right)
        op = node.op
        if isinstance(op, Add):
            return left + right
        if isinstance(op, Sub):
            return left - right
        if isinstance(op, Mult):
            return left * right
        if isinstance(op, Div):
            return left / right
        if isinstance(op, FloorDiv):
            return left // right
        if isinstance(op, Mod):
            return left % right
        if isinstance(op, Pow):
            return left ** right
        if isinstance(op, LShift):
            return left << right
        if isinstance(op, RShift):
            return left >> right
        if isinstance(op, BitOr):
            return left | right
        if isinstance(op, BitXor):
            return left ^ right
        if isinstance(op, BitAnd):
            return left & right
        raise ValueError('malformed node or string: bin op')
    if name == 'Slice':
        lower = _EvalLiteral(node.lower) if node.lower is not None else None
        upper = _EvalLiteral(node.upper) if node.upper is not None else None
        step = _EvalLiteral(node.step) if node.step is not None else None
        return slice(lower, upper, step)
    raise ValueError(f'malformed node or string: {name}')