import getpass
import os
import shutil
import sys
import unittest
import uuid

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
__package__ = "tests"
from .context import SSHContext, SSHSession


class TestSSHContext(unittest.TestCase):
    def setUp(self):
        os.makedirs("loc", exist_ok=True)
        os.makedirs("loc/task0", exist_ok=True)
        os.makedirs("loc/task1", exist_ok=True)
        for ii in ["loc/task0", "loc/task1"]:
            with open(os.path.join(ii, "test0"), "w") as fp:
                fp.write(str(uuid.uuid4()))
            with open(os.path.join(ii, "test1"), "w") as fp:
                fp.write(str(uuid.uuid4()))
            os.makedirs(os.path.join(ii, "dir0"), exist_ok=True)
            with open(os.path.join(ii, "dir0", "test2"), "w") as fp:
                fp.write(str(uuid.uuid4()))
        os.makedirs("rmt", exist_ok=True)
        try:
            self.ssh_session = SSHSession(
                {
                    "hostname": "localhost",
                    "port": 22,
                    "username": getpass.getuser(),
                    "work_path": os.path.join(os.getcwd(), "rmt"),
                }
            )
        except Exception:
            # for tianhe-2
            try:
                self.ssh_session = SSHSession(
                    {
                        "hostname": "localhost",
                        "port": 5566,
                        "username": getpass.getuser(),
                        "work_path": os.path.join(os.getcwd(), "rmt"),
                    }
                )
            except Exception:
                self.skipTest("Network error")
        self.job = SSHContext("loc", self.ssh_session)
        self.job1 = SSHContext("loc", self.ssh_session, job_uuid=self.job.job_uuid)

    def tearDown(self):
        self.job.close()
        self.job1.close()
        shutil.rmtree("loc")
        shutil.rmtree("rmt")

    def test_upload(self):
        tasks = ["task0", "task1"]
        files = ["test0", "test1"]
        self.job.upload(tasks, files)
        for ii in tasks:
            for jj in files:
                locf = os.path.join("loc", ii, jj)
                rmtf = os.path.join("rmt", self.job.job_uuid, ii, jj)
                with open(locf) as fp:
                    locs = fp.read()
                with open(rmtf) as fp:
                    rmts = fp.read()
                self.assertEqual(locs, rmts)
        self.job.upload(tasks, ["dir0"])
        for ii in tasks:
            for jj in ["dir0"]:
                for kk in ["test2"]:
                    locf = os.path.join("loc", ii, jj, kk)
                    rmtf = os.path.join("rmt", self.job.job_uuid, ii, jj, kk)
                    with open(locf) as fp:
                        locs = fp.read()
                    with open(rmtf) as fp:
                        rmts = fp.read()
                    self.assertEqual(locs, rmts)

    def test_donwload(self):
        tasks = ["task0", "task1"]
        self.job.upload(tasks, ["test0", "dir0"])
        # generate extra donwload files
        record_uuid = []
        for ii in tasks:
            for jj in ["test4", "test5"]:
                with open(os.path.join("rmt", self.job.job_uuid, ii, jj), "w") as fp:
                    tmp = str(uuid.uuid4())
                    fp.write(tmp)
                    record_uuid.append(tmp)
        # generate extra donwload dirs and files
        for ii in tasks:
            for jj in ["dir1"]:
                os.makedirs(
                    os.path.join("rmt", self.job.job_uuid, ii, jj), exist_ok=False
                )
                for kk in ["test6"]:
                    with open(
                        os.path.join("rmt", self.job.job_uuid, ii, jj, kk), "w"
                    ) as fp:
                        tmp = str(uuid.uuid4())
                        fp.write(tmp)
                        record_uuid.append(tmp)
        # donwload
        files = ["test4", "test5", "dir1"]
        self.job.download(tasks, files)
        # check dlded
        cc = 0
        for ii in tasks:
            for jj in ["test4", "test5"]:
                with open(os.path.join("loc", ii, jj)) as fp:
                    tmp = fp.read()
                    self.assertEqual(tmp, record_uuid[cc])
                    cc += 1
        for ii in tasks:
            for jj in ["dir1"]:
                for kk in ["test6"]:
                    with open(os.path.join("loc", ii, jj, kk)) as fp:
                        tmp = fp.read()
                        self.assertEqual(tmp, record_uuid[cc])
                        cc += 1

    def test_donwload_check_mark(self):
        tasks = ["task0", "task1"]
        self.job.upload(tasks, ["test0", "dir0"])
        # generate extra donwload files
        record_uuid = []
        for ii in tasks:
            for jj in ["test6", "test7"]:
                if (ii == "task1" and jj == "test7") or (
                    ii == "task0" and jj == "test6"
                ):
                    continue
                with open(os.path.join("rmt", self.job.job_uuid, ii, jj), "w") as fp:
                    tmp = str(uuid.uuid4())
                    fp.write(tmp)
                    record_uuid.append(tmp)
        # donwload
        files = ["test6", "test7", "dir1"]
        self.job.download(tasks, files, check_exists=True, mark_failure=True)
        # check dlded
        cc = 0
        for ii in tasks:
            for jj in ["test6", "test7"]:
                if (ii == "task1" and jj == "test7") or (
                    ii == "task0" and jj == "test6"
                ):
                    self.assertFalse(
                        os.path.exists(os.path.join("loc", ii, jj)),
                        msg="found " + os.path.join("loc", ii, jj),
                    )
                    self.assertTrue(
                        os.path.exists(
                            os.path.join("loc", ii, "tag_failure_download_%s" % jj)
                        ),
                        msg="failed to find "
                        + os.path.join("loc", ii, "tag_failure_download_%s" % jj),
                    )
                    continue
                with open(os.path.join("loc", ii, jj)) as fp:
                    tmp = fp.read()
                    self.assertEqual(tmp, record_uuid[cc])
                    cc += 1
        for ii in tasks:
            for jj in ["dir1"]:
                self.assertFalse(os.path.exists(os.path.join("loc", ii, jj)))
                self.assertTrue(
                    os.path.exists(
                        os.path.join("loc", ii, "tag_failure_download_%s" % jj)
                    )
                )

    def test_donwload_check_nomark(self):
        tasks = ["task0", "task1"]
        self.job.upload(tasks, ["test0", "dir0"])
        # generate extra donwload files
        record_uuid = []
        for ii in tasks:
            for jj in ["test6", "test7"]:
                if ii == "task1" and jj == "test7":
                    continue
                if ii == "task0" and jj == "test6":
                    continue
                with open(os.path.join("rmt", self.job.job_uuid, ii, jj), "w") as fp:
                    tmp = str(uuid.uuid4())
                    fp.write(tmp)
                    record_uuid.append(tmp)
        # donwload
        files = ["test6", "test7", "dir1"]
        self.job.download(tasks, files, check_exists=True, mark_failure=False)
        # check dlded
        cc = 0
        for ii in tasks:
            for jj in ["test6", "test7"]:
                if ii == "task1" and jj == "test7":
                    self.assertFalse(
                        os.path.exists(os.path.join("loc", ii, jj)),
                        msg="found " + os.path.join("loc", ii, jj),
                    )
                    self.assertFalse(
                        os.path.exists(
                            os.path.join("loc", ii, "tag_failure_download_%s" % jj)
                        ),
                        msg="found "
                        + os.path.join("loc", ii, "tag_failure_download_%s" % jj),
                    )
                    continue
                if ii == "task0" and jj == "test6":
                    self.assertFalse(
                        os.path.exists(os.path.join("loc", ii, jj)),
                        msg="found " + os.path.join("loc", ii, jj),
                    )
                    self.assertFalse(
                        os.path.exists(
                            os.path.join("loc", ii, "tag_failure_download_%s" % jj)
                        ),
                        msg="found "
                        + os.path.join("loc", ii, "tag_failure_download_%s" % jj),
                    )
                    continue
                with open(os.path.join("loc", ii, jj)) as fp:
                    tmp = fp.read()
                    self.assertEqual(tmp, record_uuid[cc])
                    cc += 1
        for ii in tasks:
            for jj in ["dir1"]:
                self.assertFalse(os.path.exists(os.path.join("loc", ii, jj)))
                self.assertFalse(
                    os.path.exists(
                        os.path.join("loc", ii, "tag_failure_download_%s" % jj)
                    )
                )

    def test_block_call(self):
        tasks = ["task0", "task1"]
        files = ["test0", "test1"]
        self.job.upload(tasks, files)
        # ls
        code, stdin, stdout, stderr = self.job.block_call("ls")
        self.assertEqual(stdout.read(), b"task0\ntask1\n")
        self.assertEqual(code, 0)
        code, stdin, stdout, stderr = self.job.block_call("ls")
        self.assertEqual(stdout.readlines(), ["task0\n", "task1\n"])
        code, stdin, stdout, stderr = self.job.block_call("ls a")
        self.assertEqual(code, 2)
        # self.assertEqual(stderr.read().decode('utf-8'), 'ls: cannot access a: No such file or directory\n')
        err_msg = stderr.read().decode("utf-8")
        self.assertTrue("ls: cannot access" in err_msg)
        self.assertTrue("No such file or directory\n" in err_msg)

    def test_block_checkcall(self):
        tasks = ["task0", "task1"]
        files = ["test0", "test1"]
        self.job.upload(tasks, files)
        # ls
        stdin, stdout, stderr = self.job.block_checkcall("ls")
        self.assertEqual(stdout.read(), b"task0\ntask1\n")
        stdin, stdout, stderr = self.job.block_checkcall("ls")
        self.assertEqual(stdout.readlines(), ["task0\n", "task1\n"])
        with self.assertRaises(RuntimeError):
            stdin, stdout, stderr = self.job.block_checkcall("ls a")

    def test_file(self):
        self.assertFalse(self.job.check_file_exists("aaa"))
        tmp = str(uuid.uuid4())
        self.job.write_file("aaa", tmp)
        self.assertTrue(self.job.check_file_exists("aaa"))
        tmp1 = self.job.read_file("aaa")
        self.assertEqual(tmp, tmp1)
