from __future__ import annotations
import ast
import random

class StringEncryptor(ast.NodeTransformer):

    def __init__(self, rng: random.Random, decrypt_str_name: str, decrypt_bytes_name: str):
        self.rng = rng
        self.dec_s = decrypt_str_name
        self.dec_b = decrypt_bytes_name

    def _strip_docstring(self, body: list[ast.stmt]) -> list[ast.stmt]:
        if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant) and isinstance(body[0].value.value, str):
            return body[1:] or [ast.Pass()]
        return body

    def visit_Module(self, node: ast.Module) -> ast.AST:
        node.body = self._strip_docstring(node.body)
        self.generic_visit(node)
        return node

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

    def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST:
        node.body = self._strip_docstring(node.body)
        self.generic_visit(node)
        return node

    def visit_JoinedStr(self, node: ast.JoinedStr) -> ast.AST:
        new_values: list[ast.AST] = []
        for v in node.values:
            if isinstance(v, ast.FormattedValue):
                v.value = self.visit(v.value)
                if v.format_spec is not None:
                    v.format_spec = self.visit(v.format_spec)
                new_values.append(v)
            else:
                new_values.append(v)
        node.values = new_values
        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 visit_Constant(self, node: ast.Constant) -> ast.AST:
        v = node.value
        if isinstance(v, bool):
            return node
        if isinstance(v, str):
            return self._encrypt_str(v) if v else node
        if isinstance(v, bytes):
            return self._encrypt_bytes(v) if v else node
        return node

    def _make_key(self, n_min: int=4, n_max: int=24) -> bytes:
        return bytes((self.rng.randrange(256) for _ in range(self.rng.randint(n_min, n_max))))

    def _xor(self, data: bytes, key: bytes) -> bytes:
        return bytes((b ^ key[i % len(key)] for i, b in enumerate(data)))

    def _encrypt_str(self, s: str) -> ast.AST:
        try:
            data = s.encode('utf-8')
        except UnicodeEncodeError:
            return ast.Constant(value=s)
        key = self._make_key()
        enc = self._xor(data, key)
        return ast.Call(func=ast.Name(id=self.dec_s, ctx=ast.Load()), args=[ast.Constant(value=enc), ast.Constant(value=key)], keywords=[])

    def _encrypt_bytes(self, b: bytes) -> ast.AST:
        key = self._make_key()
        enc = self._xor(b, key)
        return ast.Call(func=ast.Name(id=self.dec_b, ctx=ast.Load()), args=[ast.Constant(value=enc), ast.Constant(value=key)], keywords=[])

def build_decrypt_helper(decrypt_str_name: str, decrypt_bytes_name: str) -> list[ast.stmt]:
    src = f'def {decrypt_str_name}(d, k):\n    return bytes(b ^ k[i % len(k)] for i, b in enumerate(d)).decode()\ndef {decrypt_bytes_name}(d, k):\n    return bytes(b ^ k[i % len(k)] for i, b in enumerate(d))\n'
    return ast.parse(src).body