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, pool_name: str):
        self.rng = rng
        self.dec_s = decrypt_str_name
        self.dec_b = decrypt_bytes_name
        self.pool_name = pool_name
        self.entries: list[tuple[list[bytes], list[bytes], list[int], list[int], int, int]] = []

    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_entry(self, data: bytes, key: bytes) -> tuple[list[bytes], list[bytes], list[int], list[int], int, int]:
        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))
        mode = self.rng.randrange(3)
        return data_chunks, key_chunks, data_order, key_order, rotation, mode

    def _store_entry(self, data: bytes, key: bytes) -> int:
        idx = len(self.entries)
        self.entries.append(self._encoded_entry(data, key))
        return idx

    def _pool_ref(self, idx: int, decoder_name: str) -> ast.Call:
        return ast.Call(
            func=ast.Name(id=decoder_name, ctx=ast.Load()),
            args=[
                ast.Name(id=self.pool_name, ctx=ast.Load()),
                ast.Constant(value=idx),
            ],
            keywords=[],
        )

    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 self._pool_ref(self._store_entry(data, key), self.dec_s)

    def _encrypt_bytes(self, b: bytes) -> ast.AST:
        key = self._make_key()
        return self._pool_ref(self._store_entry(b, key), self.dec_b)

    def build_pool_stmt(self) -> ast.stmt | None:
        if not self.entries:
            return None
        entry_nodes = []
        for data_chunks, key_chunks, data_order, key_order, rotation, mode in self.entries:
            entry_nodes.append(
                ast.Tuple(
                    elts=[
                        self._bytes_tuple(data_chunks),
                        self._bytes_tuple(key_chunks),
                        self._int_tuple(data_order),
                        self._int_tuple(key_order),
                        ast.Constant(value=rotation),
                        ast.Constant(value=mode),
                    ],
                    ctx=ast.Load(),
                )
            )
        return ast.Assign(
            targets=[ast.Name(id=self.pool_name, ctx=ast.Store())],
            value=ast.Tuple(elts=entry_nodes, ctx=ast.Load()),
        )

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 _pw_xor_a(_data, _key):
    return bytes(_b ^ _key[_i % len(_key)] for _i, _b in enumerate(_data))

def _pw_xor_b(_data, _key):
    _out = bytearray(len(_data))
    for _i, _b in enumerate(_data):
        _out[_i] = _b ^ _key[_i % len(_key)]
    return bytes(_out)

def _pw_xor_c(_data, _key):
    return bytes(map(lambda _p: _p[1] ^ _key[_p[0] % len(_key)], enumerate(_data)))

def _pw_decode_entry(_pool, _idx):
    _d, _k, _do, _ko, _rot, _mode = _pool[_idx]
    _data = _pw_unrotate(_pw_join_chunks(_d, _do), _rot)
    _key = _pw_join_chunks(_k, _ko)
    if _mode == 1:
        return _pw_xor_b(_data, _key)
    if _mode == 2:
        return _pw_xor_c(_data, _key)
    return _pw_xor_a(_data, _key)

def {decrypt_bytes_name}(_pool, _idx):
    return _pw_decode_entry(_pool, _idx)

def {decrypt_str_name}(_pool, _idx):
    return _pw_decode_entry(_pool, _idx).decode()
'''
    return ast.parse(src).body