#!/usr/bin/env python3

import os
import sys
import random
import psutil
import time


sys.path.append('../../')
from MyToolKit import MySubprocess

# Hardcode the paths here.
STAGE1_OUTPUT_PATH: str = '/mnt/m/Stage1Output'
STAGE2_OUTPUT_PATH: str = '/mnt/m/Stage2Output'

ORDERED_TABLE_RESTORE_PATH: str = os.path.join(STAGE2_OUTPUT_PATH, 'test.ot')
with open(ORDERED_TABLE_RESTORE_PATH, 'wb') as f:
  f.write(b'')

ORDERED_TABLE_INDEX_RESTORE_PATH: str = os.path.join(STAGE2_OUTPUT_PATH, 'test.oti')
with open(ORDERED_TABLE_INDEX_RESTORE_PATH, 'wb') as f:
  f.write(b'')

proc_cmd: 'list[str]' = ['../build/src/orderedtable']
p: MySubprocess = MySubprocess(proc_cmd)

# Read the stage 1's outputs.
data_fn_list: 'list[str]' = os.listdir(STAGE1_OUTPUT_PATH)
data_fn_list = [fn for fn in data_fn_list if fn.startswith('blk') and fn.endswith('.out')]
random.shuffle(data_fn_list)

NUM_DAT_FILE_PER_CHUNK = 128
NUM_TEST_FILE = 32

# Set the restore paths.
p.send('path\n')
p.send(f'{ORDERED_TABLE_RESTORE_PATH}\n')
resp = p.readline()
assert(resp == 'OK')

p.send('indexpath\n')
p.send(f'{ORDERED_TABLE_INDEX_RESTORE_PATH}\n')
resp = p.readline()
assert(resp == 'OK')

# Adding all pairs in work list to the red-black tree.
work_list = data_fn_list[:NUM_DAT_FILE_PER_CHUNK]
for i, fn in enumerate(work_list):  
  print('#' * 60)
  print(f'Progress: {i+1} / {len(work_list)}')
  file_path: str = os.path.join(STAGE1_OUTPUT_PATH, fn)
  
  print(f'Adding information for {fn}')
  p.send('add\n')
  p.send(f'{file_path}\n')
  resp = p.readline()
  assert(resp == 'OK')
  
  p.send('size\n')
  p.wait_for_keyword('OK')
  
  mem_used = psutil.Process(p.get_pid()).memory_info().rss
  print(f'Memory used: {mem_used/1024/1024} MB')

  p.send('test\n')
  p.send(f'{file_path}\n')
  resp = p.readline()
  (a, b) = resp.split('=')[-1].split('/')
  assert(a.strip() == b.strip())
  resp = p.readline()
  assert(resp == 'OK')
  
  print('#' * 60)


# Test if or not transactions that doesn't get added to the ordered table
# are not in it.
test_list = data_fn_list[NUM_DAT_FILE_PER_CHUNK:NUM_DAT_FILE_PER_CHUNK+NUM_TEST_FILE]
for i, fn in enumerate(test_list):
  print('#' * 60)
  print(f'Testing Progress: {i+1} / {len(test_list)}')
  file_path: str = os.path.join(STAGE1_OUTPUT_PATH, fn)
  
  p.send('test\n')
  p.send(f'{file_path}\n')
  resp = p.readline()
  (a, b) = resp.split('=')[-1].split('/')
  assert(int(a.strip()) == 0)
  resp = p.readline()
  assert(resp == 'OK')

  print('#' * 60)

# Store the index and table to disk.
p.send('indexgenandstore\n')
p.wait_for_keyword('OK')
p.send('store\n')
p.wait_for_keyword('OK')

# Kill the process and restart it.
print('Restarting the subprocess.')
p.kill()
time.sleep(2)

p = MySubprocess(proc_cmd)
p.send('path\n')
p.send(f'{ORDERED_TABLE_RESTORE_PATH}\n')
resp = p.readline()
assert(resp == 'OK')

p.send('indexpath\n')
p.send(f'{ORDERED_TABLE_INDEX_RESTORE_PATH}\n')
p.wait_for_keyword('OK')

p.send('load\n')
p.wait_for_keyword('OK')

# Test if the previously added transactions are in the restored ordered table.
for i, fn in enumerate(work_list):
  print('#' * 60)
  print(f'Testing Progress: {i+1} / {len(work_list)}')
  file_path: str = os.path.join(STAGE1_OUTPUT_PATH, fn)
  
  p.send('test\n')
  p.send(f'{file_path}\n')
  resp = p.readline()
  (a, b) = resp.split('=')[-1].split('/')
  assert(a.strip() == b.strip())
  resp = p.readline()
  assert(resp == 'OK')

  print('#' * 60)


# Examine the output index.
print('Test the output ordered table index.')

f = open(ORDERED_TABLE_INDEX_RESTORE_PATH, 'rb')
f_content: bytes = f.read()
f.close()

assert(len(f_content) % 40 == 0)
num_pairs = int(len(f_content) / 40)
prev_hash_val = -1
prev_offset = -16 * 40

for pair_index in range(num_pairs):
  pair_bytes: bytes = f_content[pair_index*40:(pair_index+1)*40]
  
  hash_val = 0
  hash_val_bytes: bytes = pair_bytes[0:32]
  for byte in hash_val_bytes:
    hash_val = hash_val << 8
    hash_val |= byte
  
  offset = 0
  offset_bytes: bytes = pair_bytes[32:40]
  for byte in offset_bytes:
    offset = offset << 8
    offset |= byte
    
  print(f'Pair {pair_index+1}: {hash_val_bytes.hex()} -> {offset}')
  assert(prev_hash_val < hash_val)
  assert(prev_offset + 16 * 40 == offset)
  prev_hash_val = hash_val
  prev_offset = offset

print('Test OK!')
