from __future__ import annotations
import ast
import random
_JUNK_NAMES = ('_v', '_t', '_acc', '_state', '_buf', '_h', '_z', '_w', '_q', '_n')
_JUNK_STMT_TEMPLATES = ('{v} = ({a} ^ {b}) + ({c} & {d})', '{v} = ((({a} * 0x01000193) ^ {b}) & 0xFFFFFFFF)', '{v} = sum(((({a} >> _k) & 0xFF) for _k in range(0, 32, 8)))', '{v} = ({a} | {b}) - (({c} & {d}) >> 1)', '{v} = pow({a} | 1, 1, ({b} | 0x10001))', '{v} = (({a} + {b}) ^ ({c} - {d})) & 0x7FFFFFFF', '{v} = bytes(((_i * {a}) ^ {b}) & 0xFF for _i in range(8))', '{v} = (~({a}) | ({b} & {c})) & 0xFFFFFF')

def _rint(rng: random.Random) -> int:
    return rng.randrange(1, 1 << 30)

class JunkBranchInjector(ast.NodeTransformer):

    def __init__(self, rng: random.Random, seed_name: str, prob: float=0.4, max_depth: int=0):
        self.rng = rng
        self.seed_name = seed_name
        self.prob = prob
        self.injected = False
        self._depth = 0
        self._max_depth = max_depth

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
        return self._inject_into_body(node)
    visit_AsyncFunctionDef = visit_FunctionDef

    def visit_Module(self, node: ast.Module) -> ast.AST:
        return self._inject_into_body(node)

    def _inject_into_body(self, node):
        for child in ast.iter_child_nodes(node):
            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Module)):
                self.visit(child)
        if not getattr(node, 'body', None):
            return node
        if self.rng.random() > self.prob:
            return node
        block = self._build_block()
        idx = self.rng.randint(0, len(node.body))
        node.body.insert(idx, block)
        self.injected = True
        return node

    def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST:
        for child in node.body:
            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                self.visit(child)
        return node

    def _build_block(self) -> ast.stmt:
        s = self.seed_name
        kind = self.rng.choice(('parity', 'flt3', 'flt5'))
        if kind == 'parity':
            cond_src = f'({s} * ({s} + 1) + 1) % 2 == 0'
        elif kind == 'flt3':
            cond_src = f'({s} * {s} * {s} - {s} + 1) % 3 == 0'
        else:
            cond_src = f'({s} * {s} * {s} * {s} * {s} - {s} + 1) % 5 == 0'
        cond = ast.parse(cond_src, mode='eval').body
        body = [self._random_stmt() for _ in range(self.rng.randint(2, 4))]
        return ast.If(test=cond, body=body, orelse=[])

    def _random_stmt(self) -> ast.stmt:
        v = self.rng.choice(_JUNK_NAMES)
        tmpl = self.rng.choice(_JUNK_STMT_TEMPLATES)
        src = tmpl.format(v=v, a=_rint(self.rng), b=_rint(self.rng), c=_rint(self.rng), d=_rint(self.rng))
        return ast.parse(src).body[0]