#!python

import sys
from termcolor import colored

def solve(inp, part2=False):
    part = '1' if not part2 else '2'
    import subprocess
    p = subprocess.Popen(["./main.py", part],
                         stderr = sys.stderr,
                         stdin = subprocess.PIPE,
                         stdout = subprocess.PIPE)
    o, err = p.communicate(inp.encode('utf-8'))
    return o.decode('utf-8').strip()

def run_test(tname, inp, exp, part2):
    print("-"*30)
    print(f"Running {tname} for part {2 if part2 else 1}:")
    sol = solve(inp, part2)
    if str(sol) != exp:
        print(colored(f"Wrong on {tname}!", 'red'))
        print(f"Expected {exp}")
        print(f"Got {sol}")
    else:
        print(colored(f"Correct on {tname}!", 'green'))
        print(f"Got {sol}")

def run_tests(part2):
    import glob
    tests = list(glob.glob("test*.in"))
    for in_fn in tests:
        inp = open(in_fn, 'r').read().strip()
        tname = in_fn.split('.')[0]
        out_fn = f'{tname}.{2 if part2 else 1}.out'
        try:
            exp = open(out_fn, 'r').read().strip()
        except:
            print(colored(f"No solution for {tname}, part {2 if part2 else 1}", 'red'))
            continue
        run_test(tname, inp, exp, part2)

def main(*argv):
    if not '2' in argv[0]:
        run_tests(False)
    if not '1' in argv[0]:
        run_tests(True)

if __name__ == "__main__":
    main(sys.argv)
