# pylint: disable=invalid-name,no-self-use,protected-access
import os
import json
import sys
import pathlib
import tempfile
import io
from contextlib import redirect_stdout


from allennlp.commands import main
from allennlp.common.testing import AllenNlpTestCase


class TestPrintResults(AllenNlpTestCase):

    def setUp(self):
        super().setUp()

        self.out_dir1 = pathlib.Path(tempfile.mkdtemp(prefix="hi"))
        self.out_dir2 = pathlib.Path(tempfile.mkdtemp(prefix="hi"))

        self.directory1 = self.TEST_DIR / "results1"
        self.directory2 = self.TEST_DIR / "results2"
        self.directory3 = self.TEST_DIR / "results3"
        os.makedirs(self.directory1)
        os.makedirs(self.directory2)
        os.makedirs(self.directory3)
        json.dump({"train": 1, "test": 2, "dev": 3}, open(os.path.join(self.directory1 / "metrics.json"), "w+"))
        json.dump({"train": 4, "dev": 5}, open(os.path.join(self.directory2 / "metrics.json"), "w+"))
        json.dump({"train": 6, "dev": 7}, open(os.path.join(self.directory3 / "cool_metrics.json"), "w+"))


    def test_print_results(self):
        kebab_args = ["run.py", "print-results", str(self.TEST_DIR), "--keys", "train", "dev", "test"]
        sys.argv = kebab_args
        with io.StringIO() as buf, redirect_stdout(buf):
            main()
            output = buf.getvalue()

        lines = output.strip().split("\n")
        assert lines[0] == 'model_run, train, dev, test'

        expected_results = {
                (str(self.directory1) + "/metrics.json", "1", "3", "2"),
                (str(self.directory2) + "/metrics.json", "4", "5", "N/A")
        }
        results = {tuple(line.split(", ")) for line in lines[1:]}
        assert results == expected_results

    def test_print_results_with_metrics_filename(self):
        kebab_args = ["run.py", "print-results", str(self.TEST_DIR), "--keys",
                      "train", "dev", "test", "--metrics-filename", "cool_metrics.json"]
        sys.argv = kebab_args
        with io.StringIO() as buf, redirect_stdout(buf):
            main()
            output = buf.getvalue()

        lines = output.strip().split("\n")
        assert lines[0] == 'model_run, train, dev, test'

        expected_results = {
                (str(self.directory3) + "/cool_metrics.json", "6", "7", "N/A")
        }
        results = {tuple(line.split(", ")) for line in lines[1:]}
        assert results == expected_results
