# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
"""query_schema test class."""
import unittest
import json
import os

from ..msticpy.nbtools.entityschema import (
    Entity,
    Account,
    Host,
    Process,
    File,
    IpAddress,
)
from ..msticpy.common.wsconfig import WorkspaceConfig
from ..msticpy.nbtools.query_defns import DataFamily, DataEnvironment


_test_data_folders = [
    d for d, _, _ in os.walk(os.getcwd()) if d.endswith("/tests/testdata")
]
if len(_test_data_folders) == 1:
    _TEST_DATA = _test_data_folders[0]
else:
    _TEST_DATA = "./tests/testdata"


class Testnbtools(unittest.TestCase):
    """Unit test class."""

    def test_entity_creation(self):
        try:
            file = input_file = os.path.join(_TEST_DATA, "entities.json")
            with open(file, "r") as file_handle:
                txt = file_handle.read()
                entity_dict = json.loads(txt)

            parsed_entities = []
            for _, entity in entity_dict.items():
                e = Entity.instantiate_entity(entity)
                self.assertIsInstance(e, Entity)

                if e["Type"] == "account":
                    self.assertIsInstance(e, Account)
                    self.assertTrue("Name" in e)
                    self.assertGreater(len(e.Name), 0)
                elif e["Type"] == "host":
                    self.assertIsInstance(e, Host)
                    self.assertTrue("HostName" in e)
                    self.assertGreater(len(e.HostName), 0)
                elif e["Type"] == "process":
                    self.assertIsInstance(e, Process)
                    self.assertTrue("ProcessId" in e)
                    self.assertGreater(len(e.ProcessId), 0)
                elif e["Type"] == "file":
                    self.assertIsInstance(e, File)
                    self.assertTrue("Name" in e)
                    self.assertGreater(len(e.Name), 0)
                elif e["Type"] == "ipaddress":
                    self.assertIsInstance(e, IpAddress)
                    self.assertTrue("Address" in e)
                    self.assertGreater(len(e.Address), 0)

                parsed_entities.append(e)

            self.assertGreaterEqual(len(parsed_entities), 7)

        except Exception as ex:
            self.fail(msg="Exception {}".format(str(ex)))

    def test_wsconfig(self):
        file = input_file = os.path.join(_TEST_DATA, "config.json")

        ws_conf = WorkspaceConfig(file)
        self.assertEqual(ws_conf["tenant_id"], "99928fd7-42a5-48bc-a619-af56397b9f28")
        self.assertEqual(ws_conf["subscription_id"], "My Sub Id")
        self.assertEqual(ws_conf["resource_group"], "OMSWorkspaceRG")
        self.assertEqual(
            ws_conf["workspace_id"], "9997809c-8142-43e1-96b3-4ad87cfe95a3"
        )
        self.assertEqual(ws_conf["workspace_name"], "OMSWorkspace")

        ws_conf["workspace_name"] = "My other workspace"
        self.assertEqual(ws_conf["workspace_name"], "My other workspace")

    def test_query_defns(self):

        # WindowsSecurity = 1
        # LinuxSecurity = 2
        # SecurityAlert = 3
        self.assertEqual(
            DataFamily.WindowsSecurity, DataFamily.parse("WindowsSecurity")
        )
        self.assertEqual(DataFamily.LinuxSecurity, DataFamily.parse("LinuxSecurity"))
        self.assertEqual(DataFamily.SecurityAlert, DataFamily.parse("SecurityAlert"))
        self.assertEqual(DataFamily.WindowsSecurity, DataFamily.parse(1))
        self.assertEqual(DataFamily.LinuxSecurity, DataFamily.parse(2))
        self.assertEqual(DataFamily.SecurityAlert, DataFamily.parse(3))

        # LogAnalytics = 1
        # Kusto = 2
        self.assertEqual(
            DataEnvironment.LogAnalytics, DataEnvironment.parse("LogAnalytics")
        )
        self.assertEqual(DataEnvironment.Kusto, DataEnvironment.parse("Kusto"))
        self.assertEqual(DataEnvironment.LogAnalytics, DataEnvironment.parse(1))
        self.assertEqual(DataEnvironment.Kusto, DataEnvironment.parse(2))


if __name__ == "__main__":
    unittest.main()
    print("bye")
