# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
"""datq query test class."""
from contextlib import redirect_stdout
import io
from unittest.mock import patch
import pytest
import pytest_check as check

import pandas as pd

from adal.adal_error import AdalError
from Kqlmagic.kql_response import KqlError
from Kqlmagic.kql_engine import KqlEngineError
from Kqlmagic.my_aad_helper import AuthenticationError

from msticpy.common.exceptions import (
    MsticpyKqlConnectionError,
    MsticpyNotConnectedError,
    MsticpyNoDataSourceError,
    MsticpyDataQueryError,
)
from msticpy.data.data_providers import KqlDriver

from ...unit_test_lib import get_test_data_path


_TEST_DATA = get_test_data_path()

GET_IPYTHON_PATCH = KqlDriver.__module__ + ".get_ipython"


# pylint: disable=too-many-branches, too-many-return-statements
# pylint: disable=no-self-use


class KqlResultTest:
    """Test Kql result class."""

    def __init__(self, code=0, partial=False, status="success"):
        """Create instance."""
        self.completion_query_info = {"StatusCode": code, "StatusDescription": status}
        self.is_partial_table = partial

    def to_dataframe(self):
        """Convert dataframe."""
        return pd.DataFrame()


class _MockIPython:
    """IPython get_ipython mock."""

    def find_magic(self, magic):
        """Return None if magic isn't == kql."""
        if magic == "kql":
            return "Kqlmagic"
        return None

    def run_line_magic(self, magic, line):
        """Mock run line magic."""
        return self._run_magic(magic, line)

    def run_cell_magic(self, magic, line, cell):
        """Mock run cell magic."""
        content = cell or line
        return self._run_magic(magic, content)

    @staticmethod  # noqa: MC0001
    def _run_magic(magic, content):
        if magic == "reload_ext":
            return None
        if magic == "config":
            if "=" in content:
                return "dummy_setting"
            return True

        check.equal(magic, "kql")
        if "KqlErrorUnk" in content:
            resp = '{"error": {"code": "UnknownError"}}'
            raise KqlError(http_response=resp, message=resp)
        if "KqlErrorWS" in content:
            resp = '{"error": {"code": "WorkspaceNotFoundError"}}'
            raise KqlError(http_response=resp, message=resp)
        if "KqlEngineError" in content:
            raise KqlEngineError("Test Error")
        if "AdalErrorUnk" in content:
            resp = {"error_description": "unknown error"}
            raise AdalError("Test Error", error_response=resp)
        if "AdalErrorNR" in content:
            raise AdalError("Test Error")
        if "AdalErrorPoll" in content:
            raise AdalError("Unexpected polling state code_expired")
        if "AuthenticationError" in content:
            raise AuthenticationError("Test Error")

        if content == "--schema":
            return {
                "table1": {"field1": int, "field2": str},
                "table2": {"field1": int, "field2": str},
            }

        if "query_partial" in content:
            return KqlResultTest(code=0, partial=True, status="partial")
        if "query_failed" in content:
            return KqlResultTest(code=1, partial=False, status="failed")

        return KqlResultTest(code=0, partial=False, status="success")


@patch(GET_IPYTHON_PATCH)
def test_kql_load(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()
    check.is_true(kql_driver.loaded)

    kql_driver = KqlDriver(connection_str="la://connection")
    check.is_true(kql_driver.loaded)
    check.is_true(kql_driver.connected)


@patch(GET_IPYTHON_PATCH)
def test_kql_connect(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()
    check.is_true(kql_driver.loaded)

    kql_driver.connect(connection_str="la://connection")
    check.is_true(kql_driver.connected)


@patch(GET_IPYTHON_PATCH)
def test_kql_connect_no_cs(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()
    check.is_true(kql_driver.loaded)
    with pytest.raises(MsticpyKqlConnectionError) as mp_ex:
        kql_driver.connect()
    check.is_in("no connection string", mp_ex.value.args)


@patch(GET_IPYTHON_PATCH)
def test_kql_connect_kql_exceptions(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()

    with pytest.raises(MsticpyKqlConnectionError) as mp_ex:
        kql_driver.connect(connection_str="la://connection+KqlErrorUnk")
    check.is_in("Kql response error", mp_ex.value.args)
    check.is_false(kql_driver.connected)

    with pytest.raises(MsticpyKqlConnectionError) as mp_ex:
        kql_driver.connect(
            connection_str="la://connection.workspace('1234').tenant(KqlErrorWS)"
        )
    check.is_in("unknown workspace", mp_ex.value.args)
    check.is_false(kql_driver.connected)

    with pytest.raises(MsticpyKqlConnectionError) as mp_ex:
        kql_driver.connect(
            connection_str="la://connection.workspace('1234').tenant(KqlEngineError)"
        )
    check.is_in("kql connection error", mp_ex.value.args)
    check.is_false(kql_driver.connected)


@patch(GET_IPYTHON_PATCH)
def test_kql_connect_adal_exceptions(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()

    with pytest.raises(MsticpyKqlConnectionError) as mp_ex:
        kql_driver.connect(connection_str="la://connection+AdalErrorUnk")
    check.is_in("could not authenticate to tenant", mp_ex.value.args)
    check.is_false(kql_driver.connected)

    with pytest.raises(MsticpyKqlConnectionError) as mp_ex:
        kql_driver.connect(connection_str="la://connection+AdalErrorNR")
    check.is_in("could not authenticate to tenant", mp_ex.value.args)
    check.is_in("Full error", str(mp_ex.value.args))
    check.is_false(kql_driver.connected)

    with pytest.raises(MsticpyKqlConnectionError) as mp_ex:
        kql_driver.connect(connection_str="la://connection+AdalErrorPoll")
    check.is_in("authentication timed out", mp_ex.value.args)
    check.is_false(kql_driver.connected)


@patch(GET_IPYTHON_PATCH)
def test_kql_connect_authn_exceptions(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()

    with pytest.raises(MsticpyKqlConnectionError) as mp_ex:
        kql_driver.connect(connection_str="la://connection+AuthenticationError")
    check.is_in("authentication failed", mp_ex.value.args)
    check.is_false(kql_driver.connected)


@patch(GET_IPYTHON_PATCH)
def test_kql_schema(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()
    kql_driver.connect(connection_str="la://connection")

    check.is_in("table1", kql_driver.schema)
    check.is_in("table2", kql_driver.schema)
    check.is_in("field1", kql_driver.schema["table1"])


@patch(GET_IPYTHON_PATCH)
def test_kql_query_not_connected(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()

    with pytest.raises(MsticpyNotConnectedError) as mp_ex:
        kql_driver.query("test")
    check.is_in("not connected to a workspace.", mp_ex.value.args)
    check.is_false(kql_driver.connected)


@patch(GET_IPYTHON_PATCH)
def test_kql_query_failed(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()
    kql_driver.connect(connection_str="la://connection")

    with pytest.raises(MsticpyDataQueryError) as mp_ex:
        kql_driver.query("test query_failed")
    arg_str = "\n".join([str(arg) for arg in mp_ex.value.args])
    check.is_in("Query:", arg_str)
    check.is_in("test query_failed", arg_str)
    check.is_in("Query failed", arg_str)
    check.is_in(
        "https://msticpy.readthedocs.io/en/latest/DataAcquisition.html", arg_str
    )


@patch(GET_IPYTHON_PATCH)
def test_kql_query_success(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()
    kql_driver.connect(connection_str="la://connection")

    result_df = kql_driver.query("test query")
    check.is_instance(result_df, pd.DataFrame)


@patch(GET_IPYTHON_PATCH)
def test_kql_query_partial(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()
    kql_driver.connect(connection_str="la://connection")

    output = io.StringIO()
    with redirect_stdout(output):
        result_df = kql_driver.query("test query_partial")
    check.is_instance(result_df, pd.DataFrame)
    check.is_in("Warning - query returned partial", output.getvalue())


@patch(GET_IPYTHON_PATCH)
def test_kql_query_no_table(get_ipython):
    """Check loaded true."""
    get_ipython.return_value = _MockIPython()
    kql_driver = KqlDriver()
    kql_driver.connect(connection_str="la://connection")

    with pytest.raises(MsticpyNoDataSourceError) as mp_ex:
        query_source = {"args.table": "table3"}
        kql_driver.query("test query", query_source=query_source)

    check.is_in("table3 not found.", mp_ex.value.args)
