"""This verifies zeek-client invocations."""

import io
import os
import re
import shutil
import subprocess
import tempfile
import unittest
from contextlib import contextmanager

import zeekclient as zc

TESTS = os.path.dirname(os.path.realpath(__file__))
ROOT = os.path.normpath(os.path.join(TESTS, ".."))


# A context guard to switch the current working directory.
# With 3.11 this can go and become contextlib.chdir():
@contextmanager
def setdir(path):
    origin = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(origin)


class TestCliInvocation(unittest.TestCase):
    # This invokes the zeek-client toplevel script.

    def test_help(self):
        cproc = subprocess.run(
            ["zeek-client", "--help"],
            check=True,
            capture_output=True,
        )
        self.assertEqual(cproc.returncode, 0)

    def test_show_settings(self):
        cproc = subprocess.run(
            ["zeek-client", "show-settings"],
            check=True,
            capture_output=True,
        )
        self.assertEqual(cproc.returncode, 0)


class TestBundledCliInvocation(unittest.TestCase):
    # Verify that zeek-client finds its package in Zeek-bundled install, where
    # the package will not be in the usual Python search path. As we add more
    # system-level testing, this may become a btest too, but for we stick to
    # Python. Most system-level testing happens in the zeek-testing-cluster
    # external testsuite.
    @unittest.skipUnless(
        shutil.which("cmake") and shutil.which("make"),
        "needs both cmake and make in the system path",
    )
    def test_bundled_install(self):
        with tempfile.TemporaryDirectory() as tmpdir, setdir(tmpdir):
            # Configure the package via cmake with a Python module directory, as
            # Zeek would do. Do this from the temp directory we're now in ...
            cproc = subprocess.run(
                [
                    "cmake",
                    "-D",
                    f'PY_MOD_INSTALL_DIR={os.path.join(tmpdir, "python")}',
                    f"--install-prefix={tmpdir}",
                    ROOT,
                ],
                check=True,
                capture_output=True,
            )

            # ... and install there too, into local bin/ and python/ dirs.
            cproc = subprocess.run(["make", "install"], check=True, capture_output=True)

            # We should now be able to run "./bin/zeek-client --help".
            cproc = subprocess.run(
                [os.path.join(tmpdir, "bin", "zeek-client"), "--help"],
                capture_output=True,
                check=False,
            )
            if cproc.returncode != 0:
                print("==== STDOUT ====")
                print(cproc.stdout.decode("utf-8"))
                print("==== STDERR ====")
                print(cproc.stderr.decode("utf-8"))
                self.fail("zeek-client invocation failed")


class TestCliBasics(unittest.TestCase):
    def test_create_controller(self):
        # We mock create_controller() below, so use this class to test it:
        res = zc.cli.create_controller()
        self.assertIsNotNone(res)
        self.assertIsNotNone(res.controller_broker_id)


class TestCli(unittest.TestCase):
    # This tests the zeekclient.cli module. Most commands in that module create
    # a controller object, so we mock out its generation so we can enqueue the
    # various response events in its websocket.
    def setUp(self):
        # For capturing log writes done by zeekclient code
        self.logbuf = io.StringIO()
        zc.logs.configure(verbosity=2, stream=self.logbuf)

        self.controller = zc.controller.Controller()

        def mock_create_controller():
            self.controller.connect()
            return self.controller

        self.orig_create_controller = zc.cli.create_controller
        zc.cli.create_controller = mock_create_controller

        def mock_make_uuid(_prefix=""):
            return "mocked-reqid-00000"

        self.orig_make_uuid = zc.utils.make_uuid
        zc.controller.make_uuid = mock_make_uuid
        zc.types.make_uuid = mock_make_uuid
        zc.utils.make_uuid = mock_make_uuid

        # Capture regular writes made by the commands,
        # and let us adjust stdin:
        zc.cli.STDOUT = io.StringIO()
        zc.cli.STDIN = io.StringIO()

    def tearDown(self):
        zc.cli.create_controller = self.orig_create_controller

        zc.controller.make_uuid = self.orig_make_uuid
        zc.types.make_uuid = self.orig_make_uuid
        zc.utils.make_uuid = self.orig_make_uuid

    def assertLogLines(self, *patterns):  # noqa: N802
        buflines = self.logbuf.getvalue().split("\n")
        todo = list(patterns)
        for line in buflines:
            if todo and re.search(todo[0], line) is not None:
                todo.pop(0)
        msg = None
        if todo:
            msg = f"log pattern '{todo[0]}' not found; have:\n{self.logbuf.getvalue().strip()}"
        self.assertEqual(len(todo), 0, msg)

    def enqueue_response_event(self, event):
        msg = zc.brokertypes.DataMessage("dummy/topic", event.to_brokertype())
        self.controller.wsock.mock_recv_queue.append(msg.serialize())

    def mock_no_controller(self, inargs):
        # This fakes the scenario where connecting to the controller fails.
        parser = zc.cli.create_parser()
        args = parser.parse_args(inargs)

        def mock_create_controller():
            self.controller.wsock.mock_connect_exc = OSError()
            self.controller.connect()

        zc.cli.create_controller = mock_create_controller

        self.assertEqual(args.run_cmd(args), 1)
        self.assertLogLines("error: socket error in connect()")

    def mock_no_response(self, inargs):
        # This fakes the scenario where an established connection to the
        # controller starts experiencing trouble.
        parser = zc.cli.create_parser()
        args = parser.parse_args(inargs)

        def mock_create_controller():
            self.controller.connect()
            self.controller.wsock.mock_recv_exc = OSError()
            return self.controller

        zc.cli.create_controller = mock_create_controller

        self.assertEqual(args.run_cmd(args), 1)
        self.assertLogLines("error: no response received")

    def test_cmd_deploy_no_controller(self):
        self.mock_no_controller(["deploy"])

    def test_cmd_deploy_no_response(self):
        self.mock_no_response(["deploy"])

    def test_cmd_deploy(self):
        node_outputs = zc.types.NodeOutputs("problems on stdout", "problems on stderr")

        self.enqueue_response_event(
            zc.events.DeployResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.brokertypes.Vector(
                    [
                        zc.types.Result(
                            "reqid-0001",
                            data=zc.brokertypes.String("reqid-config-id"),
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0002",
                            instance="instance1",
                            node="manager",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0003",
                            instance="instance1",
                            node="logger",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0004",
                            success=False,
                            instance="instance1",
                            node="worker1",
                            data=node_outputs.to_brokertype(),
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0005",
                            success=False,
                            instance="instance1",
                            error="uh-oh",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0006",
                            instance="instance1",
                        ).to_brokertype(),
                        zc.types.Result("reqid-0007").to_brokertype(),
                    ],
                ),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["deploy"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_deploy)
        self.assertEqual(args.run_cmd(args), 1)
        self.assertEqual(
            zc.cli.STDOUT.getvalue(),
            """{
  "errors": [
    "uh-oh"
  ],
  "results": {
    "id": "reqid-config-id",
    "nodes": {
      "logger": {
        "instance": "instance1",
        "success": true
      },
      "manager": {
        "instance": "instance1",
        "success": true
      },
      "worker1": {
        "instance": "instance1",
        "stderr": "problems on stderr",
        "stdout": "problems on stdout",
        "success": false
      }
    }
  }
}
""",
        )

    def test_cmd_get_config_no_controller(self):
        self.mock_no_controller(["get-config"])

    def test_cmd_get_config_no_response(self):
        self.mock_no_response(["get-config"])

    def test_cmd_get_config_as_json(self):
        config = zc.types.Configuration()
        config.instances.append(zc.types.Instance("instance1"))
        config.nodes.append(
            zc.types.Node("worker1", "instance1", zc.types.ClusterRole.WORKER),
        )

        self.enqueue_response_event(
            zc.events.GetConfigurationResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.types.Result(
                    "reqid-0001",
                    data=config.to_brokertype(),
                ).to_brokertype(),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["get-config", "--as-json"])

        # cmd_get_config closes the output handle as part of its regular
        # processing. After a close, the contents of a StringIO object vanish,
        # so we prevent this:
        zc.cli.STDOUT.close = lambda: None

        self.assertEqual(args.run_cmd, zc.cli.cmd_get_config)
        self.assertEqual(args.run_cmd(args), 0)

        zc.cli.STDOUT.flush()

        self.assertEqual(
            zc.cli.STDOUT.getvalue(),
            """{
  "id": "mocked-reqid-00000",
  "instances": [
    {
      "name": "instance1"
    }
  ],
  "nodes": [
    {
      "cpu_affinity": null,
      "env": {},
      "instance": "instance1",
      "interface": null,
      "metrics_port": null,
      "name": "worker1",
      "options": null,
      "port": null,
      "role": "WORKER",
      "scripts": null
    }
  ]
}
""",
        )

    def test_cmd_get_config_as_ini(self):
        config = zc.types.Configuration()
        config.instances.append(zc.types.Instance("instance1"))
        config.nodes.append(
            zc.types.Node("worker1", "instance1", zc.types.ClusterRole.WORKER),
        )

        self.enqueue_response_event(
            zc.events.GetConfigurationResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.types.Result(
                    "reqid-0001",
                    data=config.to_brokertype(),
                ).to_brokertype(),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["get-config"])

        # cmd_get_config closes the output handle as part of its regular
        # processing. After a close, the contents of a StringIO object vanish,
        # so we prevent this:
        zc.cli.STDOUT.close = lambda: None

        self.assertEqual(args.run_cmd, zc.cli.cmd_get_config)
        self.assertEqual(args.run_cmd(args), 0)

        zc.cli.STDOUT.flush()

        self.assertEqual(
            zc.cli.STDOUT.getvalue(),
            """[instances]
instance1

[worker1]
instance = instance1
role = WORKER

""",
        )

    def test_cmd_get_config_failure(self):
        self.enqueue_response_event(
            zc.events.GetConfigurationResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.types.Result(
                    "reqid-0001",
                    success=False,
                    error="uh-oh",
                ).to_brokertype(),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["get-config"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_get_config)
        self.assertEqual(args.run_cmd(args), 1)

    def test_cmd_get_config_no_data(self):
        self.enqueue_response_event(
            zc.events.GetConfigurationResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.types.Result("reqid-0001").to_brokertype(),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["get-config"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_get_config)
        self.assertEqual(args.run_cmd(args), 1)

    def test_cmd_get_id_value_no_controller(self):
        self.mock_no_controller(["get-id-value", "Foo:id"])

    def test_cmd_get_id_value_no_response(self):
        self.mock_no_response(["get-id-value", "Foo:id"])

    def test_cmd_get_id_value(self):
        self.enqueue_response_event(
            zc.events.GetIdValueResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.brokertypes.Vector(
                    [
                        zc.types.Result(
                            "reqid-0001",
                            data=zc.brokertypes.String('"a-value"'),
                            node="worker1",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0002",
                            data=zc.brokertypes.String('"b-value"'),
                            node="worker2",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0003",
                            success=False,
                            error="that did not work",
                            node="worker2",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0004",
                            data=zc.brokertypes.Count(10),
                            node="worker2",
                        ).to_brokertype(),
                    ],
                ),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["get-id-value", "Foo:id"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_get_id_value)
        self.assertEqual(args.run_cmd(args), 1)
        self.assertEqual(
            zc.cli.STDOUT.getvalue(),
            """{
  "errors": [
    {
      "error": "that did not work",
      "source": "worker2"
    },
    {
      "error": "invalid result data type {\\"@data-type\\": \\"count\\", \\"data\\": 10}",
      "source": "worker2"
    }
  ],
  "results": {
    "worker1": "a-value",
    "worker2": "b-value"
  }
}
""",
        )

    def test_cmd_get_instances_no_controller(self):
        self.mock_no_controller(["get-instances"])

    def test_cmd_get_instances_no_response(self):
        self.mock_no_response(["get-instances"])

    def test_cmd_get_instances(self):
        self.enqueue_response_event(
            zc.events.GetInstancesResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.types.Result(
                    "reqid-0001",
                    data=zc.brokertypes.Vector(
                        [
                            zc.types.Instance(
                                "instance1",
                                "10.0.0.1",
                                123,
                            ).to_brokertype(),
                            zc.types.Instance(
                                "instance2",
                                "10.0.0.2",
                                234,
                            ).to_brokertype(),
                        ],
                    ),
                ).to_brokertype(),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["get-instances"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_get_instances)
        self.assertEqual(args.run_cmd(args), 0)
        self.assertEqual(
            zc.cli.STDOUT.getvalue(),
            """{
  "instance1": {
    "host": "10.0.0.1",
    "port": 123
  },
  "instance2": {
    "host": "10.0.0.2",
    "port": 234
  }
}
""",
        )

    def test_cmd_get_nodes_no_controller(self):
        self.mock_no_controller(["get-nodes"])

    def test_cmd_get_nodes_no_response(self):
        self.mock_no_response(["get-nodes"])

    def test_cmd_get_nodes(self):
        self.enqueue_response_event(
            zc.events.GetNodesResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.brokertypes.Vector(
                    [
                        zc.types.Result(
                            "reqid-0001",
                            instance="instance1",
                            data=zc.brokertypes.Vector(
                                [
                                    zc.types.NodeStatus(
                                        "manager",
                                        zc.types.State.RUNNING,
                                        zc.types.ManagementRole.NONE,
                                        zc.types.ClusterRole.MANAGER,
                                        12345,
                                        2200,
                                    ).to_brokertype(),
                                    zc.types.NodeStatus(
                                        "logger",
                                        zc.types.State.RUNNING,
                                        zc.types.ManagementRole.NONE,
                                        zc.types.ClusterRole.LOGGER,
                                        12346,
                                        2201,
                                    ).to_brokertype(),
                                ],
                            ),
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0002",
                            instance="instance2",
                            data=zc.brokertypes.Vector(
                                [
                                    zc.types.NodeStatus(
                                        "worker1",
                                        zc.types.State.RUNNING,
                                        zc.types.ManagementRole.NONE,
                                        zc.types.ClusterRole.WORKER,
                                        23456,
                                    ).to_brokertype(),
                                    zc.types.NodeStatus(
                                        "worker2",
                                        zc.types.State.RUNNING,
                                        zc.types.ManagementRole.NONE,
                                        zc.types.ClusterRole.WORKER,
                                        23457,
                                    ).to_brokertype(),
                                ],
                            ),
                        ).to_brokertype(),
                        # These cover various error conditions
                        zc.types.Result(
                            "reqid-0003",
                            success=False,
                            instance="instance3",
                            error="uh-oh",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0003",
                            instance="instance4",
                        ).to_brokertype(),
                    ],
                ),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["get-nodes"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_get_nodes)
        self.assertEqual(args.run_cmd(args), 1)
        self.assertEqual(
            zc.cli.STDOUT.getvalue(),
            """{
  "errors": [
    {
      "error": "uh-oh",
      "source": "instance3"
    },
    {
      "error": "result does not contain node status data",
      "source": "instance4"
    }
  ],
  "results": {
    "instance1": {
      "logger": {
        "cluster_role": "LOGGER",
        "mgmt_role": null,
        "pid": 12346,
        "port": 2201,
        "state": "RUNNING"
      },
      "manager": {
        "cluster_role": "MANAGER",
        "mgmt_role": null,
        "pid": 12345,
        "port": 2200,
        "state": "RUNNING"
      }
    },
    "instance2": {
      "worker1": {
        "cluster_role": "WORKER",
        "mgmt_role": null,
        "pid": 23456,
        "state": "RUNNING"
      },
      "worker2": {
        "cluster_role": "WORKER",
        "mgmt_role": null,
        "pid": 23457,
        "state": "RUNNING"
      }
    }
  }
}
""",
        )

    def test_cmd_restart_no_controller(self):
        self.mock_no_controller(["restart"])

    def test_cmd_restart_no_response(self):
        self.mock_no_response(["restart"])

    def test_cmd_restart(self):
        self.enqueue_response_event(
            zc.events.RestartResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.brokertypes.Vector(
                    [
                        zc.types.Result(
                            "reqid-0001",
                            instance="instance1",
                            node="manager",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0002",
                            instance="instance1",
                            node="logger",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0003",
                            instance="instance2",
                            node="worker1",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0004",
                            instance="instance2",
                            node="worker2",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0004",
                            success=False,
                            node="worker3",
                            error="unknown node",
                        ).to_brokertype(),
                    ],
                ),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["restart"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_restart)
        self.assertEqual(args.run_cmd(args), 1)
        self.assertEqual(
            zc.cli.STDOUT.getvalue(),
            """{
  "errors": [
    {
      "error": "unknown node",
      "source": "worker3"
    }
  ],
  "results": {
    "logger": true,
    "manager": true,
    "worker1": true,
    "worker2": true
  }
}
""",
        )

    def test_cmd_stage_config_no_controller(self):
        self.mock_no_controller(["stage-config", "-"])

    def test_cmd_stage_config_no_response(self):
        self.mock_no_response(["stage-config", "-"])

    def test_cmd_stage_config(self):
        self.enqueue_response_event(
            zc.events.StageConfigurationResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.brokertypes.Vector(
                    [
                        zc.types.Result(
                            "reqid-0001",
                            data=zc.brokertypes.String("reqid-config-id"),
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0002",
                            success=False,
                            error="uh-oh",
                        ).to_brokertype(),
                    ],
                ),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["stage-config", "-"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_stage_config)
        self.assertEqual(args.run_cmd(args), 1)
        self.assertEqual(
            zc.cli.STDOUT.getvalue(),
            """{
  "errors": [
    "uh-oh"
  ],
  "results": {
    "id": "reqid-config-id"
  }
}
""",
        )

    def test_cmd_deploy_config(self):
        self.enqueue_response_event(
            zc.events.StageConfigurationResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.brokertypes.Vector(
                    [
                        zc.types.Result(
                            "reqid-0001",
                            data=zc.brokertypes.String("reqid-config-id"),
                        ).to_brokertype(),
                    ],
                ),
            ),
        )

        self.enqueue_response_event(
            zc.events.DeployResponse(
                zc.brokertypes.String(zc.utils.make_uuid()),
                zc.brokertypes.Vector(
                    [
                        zc.types.Result(
                            "reqid-0001",
                            data=zc.brokertypes.String("reqid-config-id"),
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0002",
                            instance="instance1",
                            node="manager",
                        ).to_brokertype(),
                        zc.types.Result(
                            "reqid-0003",
                            instance="instance1",
                            node="logger",
                        ).to_brokertype(),
                    ],
                ),
            ),
        )

        parser = zc.cli.create_parser()
        args = parser.parse_args(["deploy-config", "-"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_deploy_config)
        self.assertEqual(args.run_cmd(args), 0)
        self.assertEqual(
            zc.cli.STDOUT.getvalue(),
            """{
  "errors": [],
  "results": {
    "id": "reqid-config-id",
    "nodes": {
      "logger": {
        "instance": "instance1",
        "success": true
      },
      "manager": {
        "instance": "instance1",
        "success": true
      }
    }
  }
}
""",
        )

    def test_show_settings(self):
        parser = zc.cli.create_parser()
        args = parser.parse_args(["show-settings"])

        self.assertEqual(args.run_cmd, zc.cli.cmd_show_settings)
        self.assertEqual(args.run_cmd(args), 0)

        # The output should be an ini file that we can load back into our
        # configuration:
        zc.CONFIG.read_string(zc.cli.STDOUT.getvalue())
