from __future__ import annotations
import ast
import random
from collections.abc import Sequence

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 _rotate_left(self, data: bytes, amount: int) -> bytes:
        if not data:
            return data
        amount %= len(data)
        return data[amount:] + data[:amount]

    def _split_chunks(self, data: bytes) -> list[bytes]:
        if len(data) <= 2:
            return [data]
        max_chunks = min(8, len(data))
        min_chunks = 2 if len(data) > 4 else 1
        count = self.rng.randint(min_chunks, max_chunks)
        if count <= 1:
            return [data]
        cuts = sorted(self.rng.sample(range(1, len(data)), count - 1))
        positions = [0, *cuts, len(data)]
        return [data[positions[i]:positions[i + 1]] for i in range(len(positions) - 1)]

    def _shuffle_chunks(self, chunks: Sequence[bytes]) -> tuple[list[bytes], list[int]]:
        if len(chunks) <= 1:
            return list(chunks), [0]
        shuffled_indexes = list(range(len(chunks)))
        self.rng.shuffle(shuffled_indexes)
        shuffled = [chunks[i] for i in shuffled_indexes]
        order = [shuffled_indexes.index(i) for i in range(len(chunks))]
        return shuffled, order

    def _bytes_tuple(self, values: Sequence[bytes]) -> ast.Tuple:
        return ast.Tuple(elts=[ast.Constant(value=value) for value in values], ctx=ast.Load())

    def _int_tuple(self, values: Sequence[int]) -> ast.Tuple:
        return ast.Tuple(elts=[ast.Constant(value=value) for value in values], ctx=ast.Load())

    def _encoded_args(self, data: bytes, key: bytes) -> list[ast.AST]:
        enc = self._xor(data, key)
        rotation = self.rng.randrange(len(enc)) if enc else 0
        rotated = self._rotate_left(enc, rotation)
        data_chunks, data_order = self._shuffle_chunks(self._split_chunks(rotated))
        key_chunks, key_order = self._shuffle_chunks(self._split_chunks(key))
        return [
            self._bytes_tuple(data_chunks),
            self._bytes_tuple(key_chunks),
            self._int_tuple(data_order),
            self._int_tuple(key_order),
            ast.Constant(value=rotation),
        ]

    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()
        return ast.Call(func=ast.Name(id=self.dec_s, ctx=ast.Load()), args=self._encoded_args(data, key), keywords=[])

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

def build_decrypt_helper(decrypt_str_name: str, decrypt_bytes_name: str) -> list[ast.stmt]:
    src = f'''
def _pw_join_chunks(_chunks, _order):
    return b''.join(_chunks[_i] for _i in _order)

def _pw_unrotate(_data, _amount):
    if not _data:
        return _data
    _amount %= len(_data)
    return _data[-_amount:] + _data[:-_amount] if _amount else _data

def {decrypt_bytes_name}(_d, _k, _do, _ko, _rot):
    _data = _pw_unrotate(_pw_join_chunks(_d, _do), _rot)
    _key = _pw_join_chunks(_k, _ko)
    return bytes(_b ^ _key[_i % len(_key)] for _i, _b in enumerate(_data))

def {decrypt_str_name}(_d, _k, _do, _ko, _rot):
    return {decrypt_bytes_name}(_d, _k, _do, _ko, _rot).decode()
'''
    return ast.parse(src).body