#!/usr/bin/env python3
import itertools
import os.path
import subprocess


def test(input, lines):
	p = subprocess.Popen([os.path.join(os.path.dirname(__file__), '.make-and-exec-binaries', 'youtube-extract-rapid'], text = False, stdin = subprocess.PIPE, stdout = subprocess.PIPE)
	stdout, stderr = p.communicate(input)
	assert not stderr
	stdout = stdout.split(b'\n')
	assert stdout[-1] == b'' and stdout[:-1] == lines, f'Got {stdout!r} instead of {lines!r} from {input!r}'


def is_id_char(c):
	return b'0' <= c <= b'9' or b'a' <= c <= b'z' or b'A' <= c <= b'Z' or c == b'_' or c == b'-'


def bytes_range(a, b):
	# Yields every char between a and b (inclusive) as a bytes object
	return map(lambda x: bytes([x]), range(ord(a), ord(b) + 1))


test(b'', [])
test(b'short\n', [])
test(b'01234567890', [b'v 01234567890'])
test(b'01234567890\n', [b'v 01234567890'])

# Videos
input = []
for a in map(lambda x: bytes.fromhex(f'{x:02x}'), range(256)):
	if is_id_char(a):
		continue
	for b in map(lambda x: bytes.fromhex(f'{x:02x}'), range(256)):
		if is_id_char(b):
			continue
		input.append(a + b'0aA_-1bB_-2' + b)
test(b''.join(input), [b'v 0aA_-1bB_-2'] * len(input))

# Channels
test(b'0123456789abcdeFGHIJ_-', [b'c 0123456789abcdeFGHIJ_-'])
test(b'UC0123456789abcdeFGHIJ_-', [b'c UC0123456789abcdeFGHIJ_-'])

# Pure playlists
playlists = [
	b'0123456789ABCDEF',
	b'PL0123456789ABCDEF',
	b'0123456789abcdefghijABCDEFGHIJ_-',
	b'PL0123456789abcdefghijABCDEFGHIJ_-',
	b'RDAMVM0123456789abcdeFGHIJ_-',
	b'RDGMEM0123456789abcdeFGHIJ_-',
	b'RDAO0123456789abcdeFGHIJ_-',
	b'RDEM0123456789abcdeFGHIJ_-',
	b'RDKM0123456789abcdeFGHIJ_-',
]
for playlist in playlists:
	test(playlist, [b'p ' + playlist])

# Music playlist madness
for prefix in (b'RDCLAK5uy_', b'RDTMAK5uy_', b'OLAK5uy_'):
	for c in bytes_range(b'k', b'n'):
		test(prefix + c + b'0123456789abcdefghijABCDEFGHIJ_-', [b'p ' + prefix + c + b'0123456789abcdefghijABCDEFGHIJ_-'])

# Playlists with video IDs
for prefix in (b'RD', b'UL', b'EL', b'CL', b'SL', b'LP', b'RDMM', b'RDQM', b'RDEM', b'RDLV', b'RDHC'):
	test(prefix + b'0aA_-1bB_-2', [b'p ' + prefix + b'0aA_-1bB_-2', b'v 0aA_-1bB_-2'])
for a, b in itertools.product(bytes_range(b'0', b'4'), bytes_range(b'0', b'9')):
	playlist = b'RD' + a + b + b'0aA_-1bB_-2'
	test(playlist, [b'p ' + playlist, b'v 0aA_-1bB_-2'])
playlist = b'RDGMEM' + b'0123456789abcdeFGHIJ_-' + b'VM0aA_-1bB_-2'
test(playlist, [b'p ' + playlist, b'v 0aA_-1bB_-2'])

# Playlists with channel IDs
for prefix in (b'UU', b'LL', b'FL', b'PU', b'UUSH'):
	test(prefix + b'0123456789abcdeFGHIJ_-', [b'p ' + prefix + b'0123456789abcdeFGHIJ_-', b'c 0123456789abcdeFGHIJ_-'])
test(b'RDCMUC0123456789abcdeFGHIJ_-', [b'p RDCMUC0123456789abcdeFGHIJ_-', b'c UC0123456789abcdeFGHIJ_-'])

# Some particular unrecognised IDs
ids = [
	b'0123456789ABCDEG',
	b'PL0123456789ABCDEG',
	b'RDCLAK5uy_j0123456789abcdefghijABCDEFGHIJ_-',
	b'RDCLAK5uy_o0123456789abcdefghijABCDEFGHIJ_-',
]
for id_ in ids:
	test(id_, [b'? ' + id_])

# Buffer rollover
BUFFER_SIZE = 1024 * 1024
for offset in range(-11, 1):
	test(b'?' * (BUFFER_SIZE + offset) + b'0aA_-1bB_-2', [b'v 0aA_-1bB_-2'])

# Max length exceedance
MAX_RESULT_SIZE = 1024
for length in range(MAX_RESULT_SIZE + 1, MAX_RESULT_SIZE + 15):
	test(b'0' * length, [])