#!/usr/bin/env python3

"""
Local testing tool for Leaking Santa's Secrets.

Disclaimer: This is *not* the same code used to test your solution when it is
submitted. The purpose of this tool is to help with debugging the interactive
problem and it has no ambitions to extensively test all possibilities that
are allowed by the problem statement. While the tool tries to yield the same
results as the real judging system, this is not guaranteed and the result
may differ if the tested program does not use the correct formatting or
exhibits other incorrect behavior. It also *does not* apply the time and
memory limits that are applied to submitted solutions, nor does it limit the
number of queries.

The behavior is controlled by an input file written as such: The first line of
input is an integer N, the length of the permutation. The next line is the
hidden permutation. For example:

    3
    3 1 2

To run the testing tool, run:

    pypy3 testing_tool.py <input file> <program> <arguments>

where `arguments` are optional arguments to the program to run. The following
show examples for different languages:

    pypy3 testing_tool.py test.in ./myprogram
    pypy3 testing_tool.py test.in java -cp . MyProgram
    pypy3 testing_tool.py test.in pypy3 myprogram.py

You can also pass `--verbose` before the input file name to see a log of the
complete interaction.
"""

import argparse
import subprocess
import sys
from typing import TextIO


def count_cycles(p: tuple[int, ...]) -> int:
    vis = [False for _ in p]
    ans = 0
    for x in p:
        ans += not vis[x]
        while not vis[x]:
            vis[x] = True
            x = p[x]

    return ans


class WrongAnswer(RuntimeError):
    """Raised whenever an incorrect answer is received."""

    pass


def vprint(*args, verbose: bool, file: TextIO, **kwargs) -> None:
    """Print to `file`, and also to stdout if `verbose is true."""
    if verbose:
        print("< ", end="")
        print(*args, **kwargs)
        sys.stdout.flush()
    print(*args, file=file, flush=True, **kwargs)


def vreadline(data: TextIO, *, verbose: bool) -> str:
    """Read a line from `data`, and also log it to stdout if `verbose` is true."""
    line = data.readline()
    if verbose and line:
        print(">", line.rstrip("\n"))
    return line


def interact(
    process: subprocess.Popen, secret_perm: tuple[int, ...], *, verbose: bool
) -> int:
    n = len(secret_perm)
    vprint(n, file=process.stdin, verbose=verbose)
    queries = 0
    try:
        while True:
            queries += 1
            query = vreadline(process.stdout, verbose=verbose)
            perm_raw = tuple(map(int, query.split()))
            if sorted(perm_raw) != sorted(range(1, n + 1)):
                raise WrongAnswer("Invalid query: not a permutation")
            # Convert to 0-indexed for internal use
            perm = tuple(x - 1 for x in perm_raw)
            if perm == secret_perm:
                vprint(0, file=process.stdin, verbose=verbose)
                break

            ans = count_cycles(tuple(secret_perm[x] for x in perm))
            vprint(ans, file=process.stdin, verbose=verbose)
    except BrokenPipeError:
        raise WrongAnswer(
            "Error when sending response to team program. Possibly exited."
        )

    print(f"Used {queries} queries")


def main() -> int:
    parser = argparse.ArgumentParser(
        usage="%(prog)s [--verbose] test.in program [args...]"
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true", help="Show interactions"
    )
    parser.add_argument("test")
    parser.add_argument("program", nargs=argparse.REMAINDER)

    args = parser.parse_args()
    if not args.program:
        parser.error("Must specify program to run")

    with open(args.test, "r") as data:
        n = int(data.readline().strip())
        a = tuple(map(int, data.readline().strip().split()))

    if sorted(a) != sorted(range(1, n + 1)):
        raise ValueError(f"Invalid data file, not a permutation of length {n=}")

    # Convert to 0-indexed for internal use
    a = tuple(x - 1 for x in a)

    process = subprocess.Popen(
        args.program,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        encoding="utf-8",
        errors="surrogateescape",
    )
    try:
        interact(process, a, verbose=args.verbose)
        process.wait()
    except WrongAnswer as exc:
        print(f"Wrong answer:\n{exc}")
        return 1
    finally:
        if process.poll() is None:
            try:
                process.terminate()
            except ProcessLookupError:  # Should be impossible, but just to be safe
                pass
        process.wait()

    if process.returncode != 0:
        print(f"Run-time error (process exited with status {process.returncode})")
        return process.returncode

    return 0


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