import socket
import re
import json
import sys
from pathlib import Path

source = Path('yudao-server/src/main/resources/application-local.yaml').read_text(encoding='utf-8')
block = re.search(r'(?m)^    redis:\s*\n((?:      [^\n]*\n)+)', source).group(1)
redis = {key: value.split(' #')[0].strip().strip('\"\'') for key, value in re.findall(r'(?m)^      (host|port|database|password):\s*([^\n]*)', block)}
connection = socket.create_connection((redis['host'], int(redis.get('port', 6379))), timeout=10)
stream = connection.makefile('rb')

def read():
    kind = stream.read(1)
    value = stream.readline().rstrip(b'\r\n')
    if kind == b'-':
        raise RuntimeError(value.decode())
    if kind == b':':
        return int(value)
    if kind == b'$':
        size = int(value)
        if size < 0:
            return None
        result = stream.read(size)
        stream.read(2)
        return result
    if kind == b'*':
        return [read() for _ in range(int(value))]
    return value

def command(*parts):
    values = [str(p).encode() if not isinstance(p, bytes) else p for p in parts]
    connection.sendall(b'*' + str(len(values)).encode() + b'\r\n' + b''.join(b'$' + str(len(v)).encode() + b'\r\n' + v + b'\r\n' for v in values))
    return read()

if redis.get('password'):
    command('AUTH', redis['password'])
command('SELECT', redis.get('database', 0))
plan = json.loads(Path('scripts/permission-fix-plan.json').read_text(encoding='utf-8'))
suffixes = tuple(':' + key for key in set(plan['allKeys'] + [pair[0] for pair in plan['aliases']]))
keys = []
cursor = b'0'
while True:
    cursor, batch = command('SCAN', cursor, 'MATCH', '*permission_menu_ids*', 'COUNT', 100)
    keys.extend(key for key in batch if key.decode().endswith(suffixes))
    if cursor == b'0':
        break
menu_suffixes = tuple(':' + str(int(value)) for value in sys.argv[1:])
if menu_suffixes:
    cursor = b'0'
    while True:
        cursor, batch = command('SCAN', cursor, 'MATCH', '*menu_role_ids*', 'COUNT', 100)
        keys.extend(key for key in batch if key.decode().endswith(menu_suffixes))
        if cursor == b'0':
            break
print('Target permission cache keys:', [key.decode() for key in keys])
if keys:
    print('Deleted permission cache entries:', command('DEL', *keys))
else:
    print('No stale affected permission cache entries')
connection.close()