"""vast.unparser - AST 还原为源码字符串（与 ast.unparse 语义一致）

TPC 编译器多处用 ``ast.unparse`` 提取类型注解字符串（如 ``t.CArray[...]``、
``list[i32, N]``、数组维度、默认值文本等）。本模块实现全部节点类型的
源码还原，保证 ``import vast as ast`` 后这些调用零修改。
"""
from __future__ import annotations

from lib.vast import Nodes as N


class Unparser:
    """递归将 vast 节点还原为源码文本。"""

    def __init__(self) -> None:
        self._Depth: int = 0

    # ------------------------------------------------------------------
    # 顶层
    # ------------------------------------------------------------------
    def Unparse(self, node: N.AST) -> str:
        if isinstance(node, N.Module):
            return '\n'.join(self._stmt(s) for s in node.body)
        if isinstance(node, N.Expression):
            return self._expr(node.body)
        return self._expr(node) if isinstance(node, N.expr) else self._stmt(node)  # type: ignore[arg-type]

    # ------------------------------------------------------------------
    # 语句
    # ------------------------------------------------------------------
    def _stmt(self, node: N.AST) -> str:
        name: str = type(node).__name__
        m: dict[str, object] = {
            'FunctionDef': self._FunctionDef,
            'AsyncFunctionDef': self._FunctionDef,
            'ClassDef': self._ClassDef,
            'Return': self._Return,
            'Delete': self._Delete,
            'Assign': self._Assign,
            'AugAssign': self._AugAssign,
            'AnnAssign': self._AnnAssign,
            'For': self._For,
            'AsyncFor': self._For,
            'While': self._While,
            'If': self._If,
            'With': self._With,
            'AsyncWith': self._With,
            'Raise': self._Raise,
            'Try': self._Try,
            'TryStar': self._Try,
            'Assert': self._Assert,
            'Import': self._Import,
            'ImportFrom': self._ImportFrom,
            'Global': self._Global,
            'Nonlocal': self._Global,
            'Expr': self._ExprStmt,
            'Pass': self._Pass,
            'Break': lambda n: 'break',
            'Continue': lambda n: 'continue',
            'Match': self._Match,
        }
        fn = m.get(name)
        if fn is not None:
            return fn(node)  # type: ignore[operator]
        raise NotImplementedError(f'unparse stmt: {name}')

    def _Block(self, body: list[N.stmt]) -> str:
        self._Depth += 1
        pad: str = '    ' * self._Depth
        lines: list[str] = [f'{pad}{self._stmt(s)}' for s in body]
        self._Depth -= 1
        return '\n'.join(lines)

    def _FunctionDef(self, node: N.FunctionDef) -> str:
        is_extern: bool = bool(getattr(node, 'Extern', False))
        kw: str = 'async def' if isinstance(node, N.AsyncFunctionDef) else ('extern def' if is_extern else 'def')
        decos: str = ''.join(f'@{self._expr(d)}\n' for d in node.decorator_list)
        params: str = self._Arguments(node.args)
        ret: str = ''
        returns: N.expr | None = node.returns
        if is_extern and isinstance(returns, N.BinOp) and isinstance(returns.op, N.BitOr):
            # 还原 extern 原始注解：去掉改写时追加的 | t.State
            right: N.expr = returns.right
            if (isinstance(right, N.Attribute) and right.attr == 'State'
                    and isinstance(right.value, N.Name) and right.value.id == 't'):
                returns = returns.left
        if returns is not None:
            ret = f' -> {self._expr(returns)}'
        body: str = self._Block(node.body)
        return f'{decos}{kw} {node.name}({params}){ret}:\n{body}'

    def _ClassDef(self, node: N.ClassDef) -> str:
        decos: str = ''.join(f'@{self._expr(d)}\n' for d in node.decorator_list)
        args: list[str] = [self._expr(b) for b in node.bases]
        for kw in node.keywords:
            args.append(f'{kw.arg}={self._expr(kw.value)}' if kw.arg else f'**{self._expr(kw.value)}')
        header: str = f'class {node.name}' + (f'({", ".join(args)})' if args else '') + ':'
        return f'{decos}{header}\n{self._Block(node.body)}'

    def _Return(self, node: N.Return) -> str:
        return 'return' if node.value is None else f'return {self._expr(node.value)}'

    def _Delete(self, node: N.Delete) -> str:
        return 'del ' + ', '.join(self._expr(t) for t in node.targets)

    def _Assign(self, node: N.Assign) -> str:
        return ', '.join(self._expr(t) for t in node.targets) + f' = {self._expr(node.value)}'

    def _AugAssign(self, node: N.AugAssign) -> str:
        ops: dict[type, str] = {
            N.Add: '+', N.Sub: '-', N.Mult: '*', N.Div: '/', N.FloorDiv: '//',
            N.Mod: '%', N.Pow: '**', N.LShift: '<<', N.RShift: '>>',
            N.BitOr: '|', N.BitXor: '^', N.BitAnd: '&', N.MatMult: '@',
        }
        op: str = ops.get(type(node.op), '?')
        return f'{self._expr(node.target)} {op}= {self._expr(node.value)}'

    def _AnnAssign(self, node: N.AnnAssign) -> str:
        val: str = f' = {self._expr(node.value)}' if node.value is not None else ''
        return f'{self._expr(node.target)}: {self._expr(node.annotation)}{val}'

    def _For(self, node: N.For) -> str:
        kw: str = 'async for' if isinstance(node, N.AsyncFor) else 'for'
        orelse: str = f'\n{"    " * self._Depth}else:\n{self._Block(node.orelse)}' if node.orelse else ''
        return f'{kw} {self._expr(node.target)} in {self._expr(node.iter)}:\n{self._Block(node.body)}{orelse}'

    def _While(self, node: N.While) -> str:
        orelse: str = f'\n{"    " * self._Depth}else:\n{self._Block(node.orelse)}' if node.orelse else ''
        return f'while {self._expr(node.test)}:\n{self._Block(node.body)}{orelse}'

    def _If(self, node: N.If) -> str:
        out: str = f'if {self._expr(node.test)}:\n{self._Block(node.body)}'
        if node.orelse:
            if len(node.orelse) == 1 and isinstance(node.orelse[0], N.If):
                out += f'\n{"    " * self._Depth}elif {self._expr(node.orelse[0].test)}:\n{self._Block(node.orelse[0].body)}'
                out += self._IfTail(node.orelse[0])
            else:
                out += f'\n{"    " * self._Depth}else:\n{self._Block(node.orelse)}'
        return out

    def _IfTail(self, node: N.If) -> str:
        out: str = ''
        if node.orelse:
            if len(node.orelse) == 1 and isinstance(node.orelse[0], N.If):
                out += f'\n{"    " * self._Depth}elif {self._expr(node.orelse[0].test)}:\n{self._Block(node.orelse[0].body)}'
                out += self._IfTail(node.orelse[0])
            else:
                out += f'\n{"    " * self._Depth}else:\n{self._Block(node.orelse)}'
        return out

    def _With(self, node: N.With) -> str:
        kw: str = 'async with' if isinstance(node, N.AsyncWith) else 'with'
        items: list[str] = []
        for item in node.items:
            s: str = self._expr(item.context_expr)
            if item.optional_vars is not None:
                s += f' as {self._expr(item.optional_vars)}'
            items.append(s)
        return f'{kw} {", ".join(items)}:\n{self._Block(node.body)}'

    def _Raise(self, node: N.Raise) -> str:
        if node.exc is None:
            return 'raise'
        out: str = f'raise {self._expr(node.exc)}'
        if node.cause is not None:
            out += f' from {self._expr(node.cause)}'
        return out

    def _Try(self, node: N.Try) -> str:
        out: str = f'try:\n{self._Block(node.body)}'
        for h in node.handlers:
            hdr: str = 'except*' if isinstance(node, N.TryStar) else 'except'
            if h.type is not None:
                hdr += f' {self._expr(h.type)}'
            if h.name:
                hdr += f' as {h.name}'
            out += f'\n{"    " * self._Depth}{hdr}:\n{self._Block(h.body)}'
        if node.orelse:
            out += f'\n{"    " * self._Depth}else:\n{self._Block(node.orelse)}'
        if node.finalbody:
            out += f'\n{"    " * self._Depth}finally:\n{self._Block(node.finalbody)}'
        return out

    def _Assert(self, node: N.Assert) -> str:
        out: str = f'assert {self._expr(node.test)}'
        if node.msg is not None:
            out += f', {self._expr(node.msg)}'
        return out

    def _Import(self, node: N.Import) -> str:
        names: list[str] = []
        for a in node.names:
            s: str = a.name
            if a.asname:
                s += f' as {a.asname}'
            names.append(s)
        return 'import ' + ', '.join(names)

    def _ImportFrom(self, node: N.ImportFrom) -> str:
        dots: str = '.' * node.level
        mod: str = node.module or ''
        names: list[str] = []
        for a in node.names:
            s: str = a.name
            if a.asname:
                s += f' as {a.asname}'
            names.append(s)
        return f'from {dots}{mod} import {", ".join(names)}'

    def _Global(self, node: N.AST) -> str:
        kw: str = 'global' if isinstance(node, N.Global) else 'nonlocal'
        return f'{kw} {", ".join(node.names)}'

    def _ExprStmt(self, node: N.Expr) -> str:
        return self._expr(node.value)

    def _Pass(self, node: N.Pass) -> str:
        return 'pass'

    def _Match(self, node: N.Match) -> str:
        out: str = f'match {self._expr(node.subject)}:\n'
        self._Depth += 1
        pad: str = '    ' * self._Depth
        for case in node.cases:
            guard: str = f' if {self._expr(case.guard)}' if case.guard else ''
            out += f'{pad}case {self._Pattern(case.pattern)}{guard}:\n'
            self._Depth += 1
            out += self._Block(case.body)
            self._Depth -= 1
        self._Depth -= 1
        return out

    def _Pattern(self, node: N.pattern) -> str:
        name: str = type(node).__name__
        if isinstance(node, N.MatchValue):
            return self._expr(node.value)
        if isinstance(node, N.MatchSingleton):
            return repr(node.value)
        if isinstance(node, N.MatchSequence):
            return '[' + ', '.join(self._Pattern(p) for p in node.patterns) + ']'
        if isinstance(node, N.MatchMapping):
            parts: list[str] = []
            for k, p in zip(node.keys, node.patterns):
                parts.append(f'{self._expr(k)}: {self._Pattern(p)}')
            if node.rest:
                parts.append(f'**{node.rest}')
            return '{' + ', '.join(parts) + '}'
        if isinstance(node, N.MatchClass):
            pos: list[str] = [self._Pattern(p) for p in node.patterns]
            kw: list[str] = [f'{a}={self._Pattern(p)}' for a, p in zip(node.kwd_attrs, node.kwd_patterns)]
            return f'{self._expr(node.cls)}({", ".join(pos + kw)})'
        if isinstance(node, N.MatchStar):
            return f'*{node.name}' if node.name else '*_'
        if isinstance(node, N.MatchAs):
            if node.pattern is not None:
                return f'{self._Pattern(node.pattern)} as {node.name}'
            return node.name or '_'
        if isinstance(node, N.MatchOr):
            return ' | '.join(self._Pattern(p) for p in node.patterns)
        raise NotImplementedError(f'unparse pattern: {name}')

    # ------------------------------------------------------------------
    # 表达式
    # ------------------------------------------------------------------
    def _expr(self, node: N.AST) -> str:
        name: str = type(node).__name__
        m: dict[str, object] = {
            'Name': self._Name,
            'Constant': self._Constant,
            'Attribute': self._Attribute,
            'Call': self._Call,
            'BinOp': self._BinOp,
            'UnaryOp': self._UnaryOp,
            'BoolOp': self._BoolOp,
            'Compare': self._Compare,
            'Subscript': self._Subscript,
            'Slice': self._Slice,
            'Tuple': self._Tuple,
            'List': self._List,
            'Set': self._Set,
            'Dict': self._Dict,
            'JoinedStr': self._JoinedStr,
            'FormattedValue': self._FormattedValue,
            'Lambda': self._Lambda,
            'IfExp': self._IfExp,
            'Starred': self._Starred,
            'ListComp': self._ListComp,
            'SetComp': self._ListComp,
            'GeneratorExp': self._ListComp,
            'DictComp': self._DictComp,
            'NamedExpr': self._NamedExpr,
            'Await': self._Await,
            'Yield': self._Yield,
            'YieldFrom': self._YieldFrom,
        }
        fn = m.get(name)
        if fn is not None:
            return fn(node)  # type: ignore[operator]
        raise NotImplementedError(f'unparse expr: {name}')

    def _Name(self, node: N.Name) -> str:
        return node.id

    def _Constant(self, node: N.Constant) -> str:
        v = node.value
        if v is Ellipsis:
            return '...'
        if v is None:
            return 'None'
        if v is True:
            return 'True'
        if v is False:
            return 'False'
        if isinstance(v, str):
            return repr(v)
        if isinstance(v, bytes):
            return repr(v)
        if isinstance(v, complex):
            # 复数文本
            return repr(v)
        return repr(v)

    def _Attribute(self, node: N.Attribute) -> str:
        return f'{self._expr(node.value)}.{node.attr}'

    def _Call(self, node: N.Call) -> str:
        args: list[str] = [self._expr(a) for a in node.args]
        for kw in node.keywords:
            if kw.arg is None:
                args.append(f'**{self._expr(kw.value)}')
            else:
                args.append(f'{kw.arg}={self._expr(kw.value)}')
        return f'{self._expr(node.func)}({", ".join(args)})'

    _BINOP_OPS: dict[type, str] = {
        N.Add: '+', N.Sub: '-', N.Mult: '*', N.Div: '/', N.FloorDiv: '//',
        N.Mod: '%', N.Pow: '**', N.LShift: '<<', N.RShift: '>>',
        N.BitOr: '|', N.BitXor: '^', N.BitAnd: '&', N.MatMult: '@',
    }

    def _BinOp(self, node: N.BinOp) -> str:
        op: str = self._BINOP_OPS.get(type(node.op), '?')
        return f'{self._expr(node.left)} {op} {self._expr(node.right)}'

    def _UnaryOp(self, node: N.UnaryOp) -> str:
        ops: dict[type, str] = {N.USub: '-', N.UAdd: '+', N.Not: 'not ', N.Invert: '~'}
        op: str = ops.get(type(node.op), '?')
        return f'{op}{self._expr(node.operand)}'

    def _BoolOp(self, node: N.BoolOp) -> str:
        kw: str = 'or' if isinstance(node.op, N.Or) else 'and'
        return f' {kw} '.join(self._expr(v) for v in node.values)

    _CMPOPS: dict[type, str] = {
        N.Eq: '==', N.NotEq: '!=', N.Lt: '<', N.LtE: '<=', N.Gt: '>', N.GtE: '>=',
        N.Is: 'is', N.IsNot: 'is not', N.In: 'in', N.NotIn: 'not in',
    }

    def _Compare(self, node: N.Compare) -> str:
        parts: list[str] = [self._expr(node.left)]
        for op, comp in zip(node.ops, node.comparators):
            parts.append(self._CMPOPS.get(type(op), '?'))
            parts.append(self._expr(comp))
        return ' '.join(parts)

    def _Subscript(self, node: N.Subscript) -> str:
        return f'{self._expr(node.value)}[{self._SliceLike(node.slice)}]'

    def _SliceLike(self, node: N.AST) -> str:
        if isinstance(node, N.Slice):
            return self._Slice(node)
        if isinstance(node, N.Tuple):
            # 切片元组：a[1, 2] / a[i, j:k] 元素按需渲染
            return ', '.join(self._SliceLike(e) for e in node.elts)
        return self._expr(node)

    def _Slice(self, node: N.Slice) -> str:
        """渲染切片：a[1:2] / a[::2] / a[:2] / a[1:2:3]。"""
        parts: list[str] = []
        if node.lower is not None:
            parts.append(self._expr(node.lower))
        # 冒号：upper 或 step 存在时才需要第一个冒号
        if node.upper is not None or node.step is not None:
            if len(parts) < 1:
                parts.append('')
            parts.append(self._expr(node.upper) if node.upper is not None else '')
        if node.step is not None:
            if len(parts) < 2:
                parts.append('')
                parts.append('')
            parts.append(self._expr(node.step))
        return ':'.join(parts)

    def _Tuple(self, node: N.Tuple) -> str:
        if len(node.elts) == 1:
            return f'({self._expr(node.elts[0])},)'
        return '(' + ', '.join(self._expr(e) for e in node.elts) + ')'

    def _List(self, node: N.List) -> str:
        return '[' + ', '.join(self._expr(e) for e in node.elts) + ']'

    def _Set(self, node: N.Set) -> str:
        return '{' + ', '.join(self._expr(e) for e in node.elts) + '}'

    def _Dict(self, node: N.Dict) -> str:
        parts: list[str] = []
        for k, v in zip(node.keys, node.values):
            if k is None:
                parts.append(f'**{self._expr(v)}')
            else:
                parts.append(f'{self._expr(k)}: {self._expr(v)}')
        return '{' + ', '.join(parts) + '}'

    def _JoinedStr(self, node: N.JoinedStr) -> str:
        out: list[str] = []
        for v in node.values:
            if isinstance(v, N.Constant):
                # 字面量段：转义内部引号与花括号
                lit: str = str(v.value)
                lit = lit.replace('{', '{{').replace('}', '}}')
                lit = lit.replace("'", "\\'")
                out.append(lit)
            else:
                out.append('{' + self._FormattedValueRaw(v) + '}')
        inner: str = ''.join(out)
        # 优先单引号（与 CPython unparse 一致）
        return "f'" + inner + "'"

    def _FormattedValue(self, node: N.FormattedValue) -> str:
        return "f'{" + self._FormattedValueRaw(node) + "}'"

    def _FormattedValueRaw(self, node: N.FormattedValue) -> str:
        s: str = self._expr(node.value)
        if node.conversion >= 0:
            s += f'!{chr(node.conversion)}'
        if node.format_spec is not None:
            if isinstance(node.format_spec, N.JoinedStr):
                spec: str = ''
                for v in node.format_spec.values:
                    if isinstance(v, N.Constant):
                        spec += str(v.value)
                    else:
                        spec += '{' + self._FormattedValueRaw(v) + '}'
            else:
                spec = self._expr(node.format_spec)
            s += f':{spec}'
        return s

    def _Lambda(self, node: N.Lambda) -> str:
        return f'lambda {self._Arguments(node.args)}: {self._expr(node.body)}'

    def _IfExp(self, node: N.IfExp) -> str:
        return f'{self._expr(node.body)} if {self._expr(node.test)} else {self._expr(node.orelse)}'

    def _Starred(self, node: N.Starred) -> str:
        return f'*{self._expr(node.value)}'

    def _ComprehensionClauses(self, gens: list[N.comprehension]) -> str:
        out: list[str] = []
        for g in gens:
            kw: str = 'async for' if g.is_async else 'for'
            s: str = f' {kw} {self._expr(g.target)} in {self._expr(g.iter)}'
            for cond in g.ifs:
                s += f' if {self._expr(cond)}'
            out.append(s)
        return ''.join(out)

    def _ListComp(self, node: N.AST) -> str:
        cls_name: str = type(node).__name__
        if isinstance(node, N.SetComp):
            fmt: str = '{'
            close: str = '}'
            elt: str = self._expr(node.elt)
        elif isinstance(node, N.GeneratorExp):
            fmt, close, elt = '(', ')', self._expr(node.elt)
        elif isinstance(node, N.DictComp):
            return '{' + f'{self._expr(node.key)}: {self._expr(node.value)}' + self._ComprehensionClauses(node.generators) + '}'
        else:
            fmt, close, elt = '[', ']', self._expr(node.elt)
        return fmt + elt + self._ComprehensionClauses(node.generators) + close

    def _DictComp(self, node: N.DictComp) -> str:
        return '{' + f'{self._expr(node.key)}: {self._expr(node.value)}' + \
            self._ComprehensionClauses(node.generators) + '}'

    def _NamedExpr(self, node: N.NamedExpr) -> str:
        return f'{self._expr(node.target)} := {self._expr(node.value)}'

    def _Await(self, node: N.Await) -> str:
        return f'await {self._expr(node.value)}'

    def _Yield(self, node: N.Yield) -> str:
        return 'yield' if node.value is None else f'yield {self._expr(node.value)}'

    def _YieldFrom(self, node: N.YieldFrom) -> str:
        return f'yield from {self._expr(node.value)}'

    # ------------------------------------------------------------------
    # 参数列表
    # ------------------------------------------------------------------
    def _Arguments(self, args: N.arguments) -> str:
        parts: list[str] = []
        n_args: int = len(args.posonlyargs) + len(args.args)
        n_defaults: int = len(args.defaults)
        # 合并 posonly + 普通参数，defaults 尾部对齐
        all_args: list[N.arg] = list(args.posonlyargs) + list(args.args)
        all_defaults: list[N.expr | None] = [None] * (n_args - n_defaults) + list(args.defaults)
        for i, a in enumerate(all_args):
            s: str = a.arg
            if a.annotation is not None:
                s += f': {self._expr(a.annotation)}'
            if all_defaults[i] is not None:
                s += f'={self._expr(all_defaults[i])}'
            parts.append(s)
            if args.posonlyargs and i == len(args.posonlyargs) - 1:
                parts.append('/')
        if args.vararg is not None:
            s = f'*{args.vararg.arg}'
            if args.vararg.annotation is not None:
                s += f': {self._expr(args.vararg.annotation)}'
            parts.append(s)
        elif args.kwonlyargs:
            parts.append('*')
        for a, d in zip(args.kwonlyargs, args.kw_defaults):
            s = a.arg
            if a.annotation is not None:
                s += f': {self._expr(a.annotation)}'
            if d is not None:
                s += f'={self._expr(d)}'
            parts.append(s)
        if args.kwarg is not None:
            s = f'**{args.kwarg.arg}'
            if args.kwarg.annotation is not None:
                s += f': {self._expr(args.kwarg.annotation)}'
            parts.append(s)
        return ', '.join(parts)


def Unparse(node: N.AST) -> str:
    """便捷入口。"""
    return Unparser().Unparse(node)