from __future__ import annotations
import ast
import random

class NumberObfuscator(ast.NodeTransformer):

    def __init__(self, rng: random.Random, prob: float=0.85):
        self.rng = rng
        self.prob = prob

    def visit_Constant(self, node: ast.Constant) -> ast.AST:
        v = node.value
        if isinstance(v, bool):
            return node
        if isinstance(v, int):
            if self.rng.random() <= self.prob:
                return self._obfuscate(v)
        return node

    def visit_match_case(self, node: ast.match_case) -> ast.AST:
        if node.guard is not None:
            node.guard = self.visit(node.guard)
        node.body = [self.visit(stmt) for stmt in node.body]
        return node

    def _obfuscate(self, n: int) -> ast.AST:
        choices = ['xor', 'sum', 'shift', 'frombytes', 'affine', 'split']
        kind = self.rng.choice(choices)
        try:
            return getattr(self, f'_t_{kind}')(n)
        except Exception:
            return ast.Constant(value=n)

    def _t_xor(self, n: int) -> ast.AST:
        k = self.rng.randrange(1, 1 << 32)
        return self._parse_expr(f'(({n ^ k}) ^ {k})')

    def _t_sum(self, n: int) -> ast.AST:
        a = self.rng.randint(-(1 << 16), 1 << 16)
        return self._parse_expr(f'(({n - a}) + ({a}))')

    def _t_shift(self, n: int) -> ast.AST:
        if abs(n) >= 1 << 60:
            return self._t_xor(n)
        s = self.rng.randint(1, 8)
        return self._parse_expr(f'(({n} << {s}) >> {s})')

    def _t_frombytes(self, n: int) -> ast.AST:
        signed = n < 0
        bit_len = max(1, n.bit_length() + (1 if signed else 0))
        byte_len = (bit_len + 7) // 8
        try:
            data = n.to_bytes(byte_len, 'big', signed=signed)
        except OverflowError:
            return self._t_xor(n)
        signed_arg = 'True' if signed else 'False'
        return self._parse_expr(f"int.from_bytes({data!r}, 'big', signed={signed_arg})")

    def _t_affine(self, n: int) -> ast.AST:
        a = self.rng.randint(2, 64)
        b = self.rng.randint(-(1 << 12), 1 << 12)
        return self._parse_expr(f'((({n * a + b}) - ({b})) // {a})')

    def _t_split(self, n: int) -> ast.AST:
        p = self.rng.randint(-(1 << 16), 1 << 16)
        q = n - p
        return self._parse_expr(f'({p} + ({q}))')

    def _parse_expr(self, src: str) -> ast.AST:
        return ast.parse(src, mode='eval').body