#!/usr/bin/env python
import argparse
import asyncio
import random


async def grab_output(p: asyncio.subprocess.Process) -> float | None:
    out = (await p.communicate())[0].decode()
    if p.returncode:
        print("Error:", out)
        return None
    print(out.strip())
    return float(out.rsplit(':')[-1])


async def run_custom_exec(p_s: str, sample: int, search: int, insertion: int) -> float | None:
    p = (await asyncio.create_subprocess_exec(p_s, str(sample), str(search), str(insertion),
                                              stdout=asyncio.subprocess.PIPE))
    await p.wait()
    return await grab_output(p)


async def run_exec(p_s: str, search: int, insertion: int) -> float | None:
    p = await asyncio.create_subprocess_exec(p_s, str(search), str(insertion), stdout=asyncio.subprocess.PIPE)
    await p.wait()
    return await grab_output(p)


def grab_limit(limit: str) -> tuple[int, int]:
    if '-' in limit:
        return tuple(map(lambda x: int(x.strip()), limit.split("-", 1)))
    return 10, int(limit)


async def main(matches: argparse.Namespace) -> None:
    def pow2(n: int) -> int:
        return 2 ** n

    search_limit_low, search_limit_high = grab_limit(matches.limit)
    limit_low, limit_high = grab_limit(matches.sample_limit)
    total_run = 0
    fault_run = 0
    result = {}

    for limit in range(limit_low, limit_high):
        for search in range(search_limit_low, search_limit_high):
            insertion = pow2(search + random.randint(1, 3))
            search = pow2(search)
            sample = pow2(limit) - 1
            # print('\r', insertion, search, sample, end='')
            if sample not in result:
                result.update({sample: {}})
            if search not in result[sample]:
                result[sample].update({search: {'normal': [], 'custom': []}})
            for _ in range(3):
                retries = 3
                while (normal := await run_exec(matches.exec2, search, insertion)) is None:
                    retries -= 1
                    fault_run += 1
                    if not retries:
                        break
                total_run += 4 - retries
                retries = 3
                while (custom := await run_custom_exec(matches.exec1, sample, search, insertion)) is None:
                    retries -= 1
                    fault_run += 1
                    if not retries:
                        break
                total_run += 4 - retries
                result[sample][search]['normal'].append(normal)
                result[sample][search]['custom'].append(custom)
    print(fault_run, '/', total_run)
    print(result)


if __name__ == '__main__':
    arg_ = argparse.ArgumentParser()
    arg_.add_argument("exec1")
    arg_.add_argument("exec2")
    arg_.add_argument("sample_limit")
    arg_.add_argument("limit")
    parser_ = arg_.parse_args()
    asyncio.run(main(parser_))