#!/usr/bin/env python3

import sys

def main():
    filename1 = sys.argv[1]
    filename2 = sys.argv[2]
    cairo_mem = {}
    cairo_rs_mem = {}
    name = filename2.split("/")[-1]
    if name in ["blake2s_felt", "blake2s_finalize", "blake2s_hello_world_hash", "blake2s_integration_tests", "dict_squash", "squash_dict", "dict_write", "dict_read", "dict_update"]:
        with open(filename1, 'rb') as f:
            cairo_raw = f.read()
            assert len(cairo_raw) % 40 == 0, f'{filename1}: malformed memory file from Cairo VM'
            chunks = len(cairo_raw) // 40
            for i in range(0, chunks):
                chunk = cairo_raw[i*40:(i+1)*40]
                k, v = int.from_bytes(chunk[:8], 'little'), int.from_bytes(chunk[8:], 'little')
                assert k not in cairo_mem, f'{filename1}: address {k} has two values'
                cairo_mem[k] = v
            assert len(cairo_mem) * 40 == len(cairo_raw), f'{filename1}: {len(cairo_mem) * 40} != {len(cairo_raw)}'

        with open(filename2, 'rb') as f:
            cairo_rs_raw = f.read()
            assert len(cairo_rs_raw) % 40 == 0, f'{filename2}: malformed memory file from cairo-rs'
            chunks = len(cairo_rs_raw) // 40
            for i in range(0, chunks):
                chunk = cairo_rs_raw[i*40:(i+1)*40]
                k, v = int.from_bytes(chunk[:8], 'little'), int.from_bytes(chunk[8:], 'little')
                assert k not in cairo_rs_mem, f'{filename2}: address {k} has two values'
                cairo_rs_mem[k] = v
            assert len(cairo_rs_mem) * 40 == len(cairo_rs_raw), f'{filename2}: {len(cairo_rs_mem) * 40} != {len(cairo_rs_raw)}'

        assert len(cairo_mem) == len(cairo_rs_mem), f'{filename2}: len(cairo_mem)={len(cairo_mem)} len(cairo_mem)={len(cairo_rs_mem)}'
        if cairo_mem != cairo_rs_mem:
            print(f'Mismatch between {filename1} (Cairo) and {filename2} (cairo_rs)')
            print('keys in Cairo but not cairo-rs:')
            for k in cairo_mem:
                if k in cairo_rs_mem:
                    continue
                print(f'{k}:{v}')
            print('keys in cairo_rs but not Cairo:')
            for k in cairo_rs_mem:
                if k in cairo_mem:
                    continue
                print(f'{k}:{v}')
            print('mismatched values (Cairo <-> cairo_rs)):')
            for k in cairo_rs_mem:
                if k not in cairo_mem:
                    continue
                if cairo_rs_mem[k] == cairo_mem[k]:
                    continue
                print(f'{k}:({cairo_mem[k]} <-> {cairo_rs_mem[k]})')
            exit(1)


if __name__ == '__main__':
    main()

