#!/usr/bin/env python3
"""
marian test fixture

Marian can be mocked for tests when the actual binary doesn't need to be run.
It outputs the training corpus to marian.input.txt and saves the arguments to marian.args.txt.
"""

import json
import os
import sys

artifacts_dir = os.environ.get("TEST_ARTIFACTS")

if not artifacts_dir:
    raise Exception("TEST_ARTIFACTS was not set.")

if not os.path.exists(artifacts_dir):
    raise Exception("The TEST_ARTIFACTS directory did not exist")

# Write the arguments to disk
arguments = sys.argv[1:]
with open(os.path.join(artifacts_dir, "marian.args.txt"), "w") as file:
    json.dump(arguments, file)

# Output the input for N times and exit
counter = 0
with open(os.path.join(artifacts_dir, "marian.input.txt"), "w") as file:
    for line in sys.stdin:
        file.write(line)
        counter += 1
        if counter == 100:
            # create "model" files
            open(os.path.join(artifacts_dir, 'artifacts', "model.npz.best-chrf.npz"), 'a').close()
            open(os.path.join(artifacts_dir, 'artifacts', "model.npz.best-chrf.npz.decoder.yml"), 'a').close()
            sys.exit(0)

