"""Utility for testing InteractiveModelRunner.

Use it from pass-specific tests by providing a main .py which calls this library's
`run_interactive` with an appropriate callback to provide advice.

From .ll tests, just call the above-mentioned main as a prefix to the opt/llc
invocation (with the appropriate flags enabling the interactive mode)

Examples:
test/Transforms/Inline/ML/interactive-mode.ll
test/CodeGen/MLRegAlloc/interactive-mode.ll
"""

import ctypes
import log_reader
import io
import math
import os
import subprocess
from typing import Callable, List, Union


def send(f: io.BufferedWriter, value: Union[int, float], spec: log_reader.TensorSpec):
    """Send the `value` - currently just a scalar - formatted as per `spec`."""

    # just int64 for now
    assert spec.element_type == ctypes.c_int64
    to_send = ctypes.c_int64(int(value))
    assert f.write(bytes(to_send)) == ctypes.sizeof(spec.element_type) * math.prod(
        spec.shape
    )
    f.flush()


def run_interactive(
    temp_rootname: str,
    make_response: Callable[[List[log_reader.TensorValue]], Union[int, float]],
    process_and_args: List[str],
):
    """Host the compiler.
    Args:
      temp_rootname: the base file name from which to construct the 2 pipes for
      communicating with the compiler.
      make_response: a function that, given the current tensor values, provides a
      response.
      process_and_args: the full commandline for the compiler. It it assumed it
      contains a flag poiting to `temp_rootname` so that the InteractiveModeRunner
      would attempt communication on the same pair as this function opens.

    This function sets up the communication with the compiler - via 2 files named
    `temp_rootname`.in and `temp_rootname`.out - prints out the received features,
    and sends back to the compiler an advice (which it gets from `make_response`).
    It's used for testing, and also to showcase how to set up communication in an
    interactive ML ("gym") environment.
    """
    to_compiler = temp_rootname + ".in"
    from_compiler = temp_rootname + ".out"
    try:
        os.mkfifo(to_compiler, 0o666)
        os.mkfifo(from_compiler, 0o666)
        compiler_proc = subprocess.Popen(
            process_and_args, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL
        )
        with io.BufferedWriter(io.FileIO(to_compiler, "wb")) as tc:
            with io.BufferedReader(io.FileIO(from_compiler, "rb")) as fc:
                tensor_specs, _, advice_spec = log_reader.read_header(fc)
                context = None
                while compiler_proc.poll() is None:
                    next_event = fc.readline()
                    if not next_event:
                        break
                    (
                        last_context,
                        observation_id,
                        features,
                        _,
                    ) = log_reader.read_one_observation(
                        context, next_event, fc, tensor_specs, None
                    )
                    if last_context != context:
                        print(f"context: {last_context}")
                    context = last_context
                    print(f"observation: {observation_id}")
                    tensor_values = []
                    for fv in features:
                        log_reader.pretty_print_tensor_value(fv)
                        tensor_values.append(fv)
                    send(tc, make_response(tensor_values), advice_spec)
        _, err = compiler_proc.communicate()
        print(err.decode("utf-8"))
        compiler_proc.wait()

    finally:
        os.unlink(to_compiler)
        os.unlink(from_compiler)