import json
import logging
import os
import time
from unittest.mock import MagicMock, patch
from uuid import uuid4
import tempfile
import jwt
import pandas as pd
import pathlib
import pytest
from azure.ai.evaluation._azure._token_manager import AzureMLTokenManager

import azure.ai.evaluation._evaluate._utils as ev_utils
from azure.ai.evaluation._evaluate._eval_run import EvalRun, RunStatus
from azure.ai.evaluation._exceptions import EvaluationException


def generate_mock_token():
    expiration_time = time.time() + 3600  # 1 hour in the future
    return jwt.encode({"exp": expiration_time}, "secret", algorithm="HS256")


@pytest.mark.unittest
@patch.object(AzureMLTokenManager, "get_token", return_value=generate_mock_token())
class TestEvalRun:
    """Unit tests for the eval-run object."""

    _MOCK_CREDS = dict(
        tracking_uri=(
            "https://region.api.azureml.ms/mlflow/v2.0/subscriptions"
            "/000000-0000-0000-0000-0000000/resourceGroups/mock-rg-region"
            "/providers/Microsoft.MachineLearningServices"
            "/workspaces/mock-ws-region"
        ),
        subscription_id="000000-0000-0000-0000-0000000",
        group_name="mock-rg-region",
        workspace_name="mock-ws-region",
        management_client=MagicMock(),
    )

    def _get_mock_create_response(self, status=200):
        """Return the mock create request"""
        mock_response = MagicMock()
        mock_response.status_code = status
        if status != 200:
            mock_response.text = lambda: "Mock error"
        else:
            mock_response.json.return_value = {
                "run": {"info": {"run_id": str(uuid4()), "experiment_id": str(uuid4()), "run_name": str(uuid4())}}
            }
        return mock_response

    def _get_mock_end_response(self, status=200):
        """Get the mock end run response."""
        mock_response = MagicMock()
        mock_response.status_code = status
        mock_response.text = lambda: "Everything good" if status == 200 else "Everything bad"
        return mock_response

    @pytest.mark.parametrize(
        "status,should_raise", [("KILLED", False), ("WRONG_STATUS", True), ("FINISHED", False), ("FAILED", False)]
    )
    def test_end_raises(self, token_mock, status, should_raise, caplog):
        """Test that end run raises exception if incorrect status is set."""
        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request", return_value=self._get_mock_create_response()
        ), caplog.at_level(logging.INFO):
            with EvalRun(run_name=None, **TestEvalRun._MOCK_CREDS) as run:
                if should_raise:
                    with pytest.raises(EvaluationException) as cm:
                        run._end_run(status)
                    assert status in cm.value.args[0]
                else:
                    run._end_run(status)
                    assert len(caplog.records) == 0

    def test_run_logs_if_terminated(self, token_mock, caplog):
        """Test that run warn user if we are trying to terminate it twice."""
        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request", return_value=self._get_mock_create_response()
        ), caplog.at_level(logging.INFO):
            logger = logging.getLogger(EvalRun.__module__)
            # All loggers, having promptflow. prefix will have "promptflow" logger
            # as a parent. This logger does not propagate the logs and cannot be
            # captured by caplog. Here we will skip this logger to capture logs.
            logger.parent = logging.root
            run = EvalRun(
                run_name=None,
                tracking_uri="www.microsoft.com",
                subscription_id="mock",
                group_name="mock",
                workspace_name="mock",
                management_client=MagicMock(),
            )
            run._start_run()
            run._end_run("KILLED")
            run._end_run("KILLED")
            assert len(caplog.records) == 1
            assert "Unable to stop run due to Run status=RunStatus.TERMINATED." in caplog.records[0].message

    def test_end_logs_if_fails(self, token_mock, caplog):
        """Test that if the terminal status setting was failed, it is logged."""
        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request",
            side_effect=[self._get_mock_create_response(), self._get_mock_end_response(500)],
        ), caplog.at_level(logging.INFO):
            logger = logging.getLogger(EvalRun.__module__)
            # All loggers, having promptflow. prefix will have "promptflow" logger
            # as a parent. This logger does not propagate the logs and cannot be
            # captured by caplog. Here we will skip this logger to capture logs.
            logger.parent = logging.root
            with EvalRun(
                run_name=None,
                tracking_uri="www.microsoft.com",
                subscription_id="mock",
                group_name="mock",
                workspace_name="mock",
                management_client=MagicMock(),
            ):
                pass
            assert len(caplog.records) == 1
            assert "Unable to terminate the run." in caplog.records[0].message

    def test_start_run_fails(self, token_mock, caplog):
        """Test that there are log messges if run was not started."""
        mock_response_start = MagicMock()
        mock_response_start.status_code = 500
        mock_response_start.text = lambda: "Mock internal service error."
        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request", return_value=mock_response_start
        ), caplog.at_level(logging.INFO):
            logger = logging.getLogger(EvalRun.__module__)
            # All loggers, having promptflow. prefix will have "promptflow" logger
            # as a parent. This logger does not propagate the logs and cannot be
            # captured by caplog. Here we will skip this logger to capture logs.
            logger.parent = logging.root
            run = EvalRun(
                run_name=None,
                tracking_uri="www.microsoft.com",
                subscription_id="mock",
                group_name="mock",
                workspace_name="mock",
                management_client=MagicMock(),
            )
            run._start_run()
            assert len(caplog.records) == 1
            assert "500" in caplog.records[0].message
            assert mock_response_start.text() in caplog.records[0].message
            assert "The results will be saved locally" in caplog.records[0].message
            caplog.clear()
            # Log artifact
            run.log_artifact("test")
            assert len(caplog.records) == 1
            assert "Unable to log artifact due to Run status=RunStatus.BROKEN." in caplog.records[0].message
            caplog.clear()
            # Log metric
            run.log_metric("a", 42)
            assert len(caplog.records) == 1
            assert "Unable to log metric due to Run status=RunStatus.BROKEN." in caplog.records[0].message
            caplog.clear()
            # End run
            run._end_run("FINISHED")
            assert len(caplog.records) == 1
            assert "Unable to stop run due to Run status=RunStatus.BROKEN." in caplog.records[0].message
            caplog.clear()

    def test_run_name(self, token_mock):
        """Test that the run name is the same as ID if name is not given."""
        mock_response = self._get_mock_create_response()
        with patch("azure.ai.evaluation._http_utils.HttpPipeline.request", return_value=mock_response):
            with EvalRun(
                run_name=None,
                tracking_uri="www.microsoft.com",
                subscription_id="mock",
                group_name="mock",
                workspace_name="mock",
                management_client=MagicMock(),
            ) as run:
                pass
        assert run.info.run_id == mock_response.json.return_value["run"]["info"]["run_id"]
        assert run.info.experiment_id == mock_response.json.return_value["run"]["info"]["experiment_id"]
        assert run.info.run_name == mock_response.json.return_value["run"]["info"]["run_name"]

    def test_run_with_name(self, token_mock):
        """Test that the run name is not the same as id if it is given."""
        mock_response = self._get_mock_create_response()
        mock_response.json.return_value["run"]["info"]["run_name"] = "test"
        with patch("azure.ai.evaluation._http_utils.HttpPipeline.request", return_value=mock_response):
            with EvalRun(
                run_name="test",
                tracking_uri="www.microsoft.com",
                subscription_id="mock",
                group_name="mock",
                workspace_name="mock",
                management_client=MagicMock(),
            ) as run:
                pass
        assert run.info.run_id == mock_response.json.return_value["run"]["info"]["run_id"]
        assert run.info.experiment_id == mock_response.json.return_value["run"]["info"]["experiment_id"]
        assert run.info.run_name == "test"
        assert run.info.run_name != run.info.run_id

    def test_get_urls(self, token_mock):
        """Test getting url-s from eval run."""
        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request", return_value=self._get_mock_create_response()
        ):
            with EvalRun(run_name="test", **TestEvalRun._MOCK_CREDS) as run:
                pass
        assert run.get_run_history_uri() == (
            "https://region.api.azureml.ms/history/v1.0/subscriptions"
            "/000000-0000-0000-0000-0000000/resourceGroups/mock-rg-region"
            "/providers/Microsoft.MachineLearningServices"
            "/workspaces/mock-ws-region/experimentids/"
            f"{run.info.experiment_id}/runs/{run.info.run_id}"
        ), "Wrong RunHistory URL"
        assert run.get_artifacts_uri() == (
            "https://region.api.azureml.ms/history/v1.0/subscriptions"
            "/000000-0000-0000-0000-0000000/resourceGroups/mock-rg-region"
            "/providers/Microsoft.MachineLearningServices"
            "/workspaces/mock-ws-region/experimentids/"
            f"{run.info.experiment_id}/runs/{run.info.run_id}"
            "/artifacts/batch/metadata"
        ), "Wrong Artifacts URL"
        assert run.get_metrics_url() == (
            "https://region.api.azureml.ms/mlflow/v2.0/subscriptions"
            "/000000-0000-0000-0000-0000000/resourceGroups/mock-rg-region"
            "/providers/Microsoft.MachineLearningServices"
            "/workspaces/mock-ws-region/api/2.0/mlflow/runs/log-metric"
        ), "Wrong Metrics URL"

    @pytest.mark.parametrize(
        "log_function,expected_str", [("log_artifact", "register artifact"), ("log_metric", "save metrics")]
    )
    def test_log_artifacts_logs_error(self, token_mock, tmp_path, caplog, log_function, expected_str):
        """Test that the error is logged."""
        mock_response = MagicMock()
        mock_response.status_code = 404
        mock_response.text = lambda: "Mock not found error."
        if log_function == "log_artifact":
            with open(os.path.join(tmp_path, "test.json"), "w") as fp:
                json.dump({"f1": 0.5}, fp)

        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request",
            side_effect=[
                self._get_mock_create_response(),
                mock_response,
                self._get_mock_end_response(),
            ],
        ), caplog.at_level(logging.INFO):
            logger = logging.getLogger(EvalRun.__module__)
            # All loggers, having promptflow. prefix will have "promptflow" logger
            # as a parent. This logger does not propagate the logs and cannot be
            # captured by caplog. Here we will skip this logger to capture logs.
            logger.parent = logging.root
            with EvalRun(run_name="test", **TestEvalRun._MOCK_CREDS) as run:
                fn = getattr(run, log_function)
                if log_function == "log_artifact":
                    with open(os.path.join(tmp_path, EvalRun.EVALUATION_ARTIFACT), "w") as fp:
                        fp.write("42")
                    kwargs = {"artifact_folder": tmp_path}
                else:
                    kwargs = {"key": "f1", "value": 0.5}
                with patch("azure.ai.evaluation._evaluate._eval_run.BlobServiceClient", return_value=MagicMock()):
                    fn(**kwargs)

        assert len(caplog.records) == 1
        assert mock_response.text() in caplog.records[0].message
        assert "404" in caplog.records[0].message
        assert expected_str in caplog.records[0].message

    @pytest.mark.parametrize(
        "dir_exists,dir_empty,expected_error",
        [
            (True, True, "The path to the artifact is empty."),
            (False, True, "The path to the artifact is either not a directory or does not exist."),
            (True, False, "The run results file was not found, skipping artifacts upload."),
        ],
    )
    def test_wrong_artifact_path(
        self,
        token_mock,
        tmp_path,
        caplog,
        dir_exists,
        dir_empty,
        expected_error,
    ):
        """Test that if artifact path is empty, or dies not exist we are logging the error."""
        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request", return_value=self._get_mock_create_response()
        ), caplog.at_level(logging.INFO):
            with EvalRun(run_name="test", **TestEvalRun._MOCK_CREDS) as run:
                logger = logging.getLogger(EvalRun.__module__)
                # All loggers, having promptflow. prefix will have "promptflow" logger
                # as a parent. This logger does not propagate the logs and cannot be
                # captured by caplog. Here we will skip this logger to capture logs.
                logger.parent = logging.root
                artifact_folder = tmp_path if dir_exists else "wrong_path_567"
                if not dir_empty:
                    with open(os.path.join(tmp_path, "test.txt"), "w") as fp:
                        fp.write("42")
                run.log_artifact(artifact_folder)
            assert len(caplog.records) == 1
            assert expected_error in caplog.records[0].message

    def test_store_multi_modal_no_images(self, token_mock, caplog):
        data_path = os.path.join(pathlib.Path(__file__).parent.resolve(), "data")
        data_file = os.path.join(data_path, "generated_qa_chat_conv.jsonl")
        data_convo = pd.read_json(data_file, lines=True)
        with tempfile.TemporaryDirectory() as tmpdir:
            for value in data_convo["messages"]:
                ev_utils._store_multimodal_content(value, tmpdir)

    def test_store_multi_modal_image_urls(self, token_mock, caplog):
        data_path = os.path.join(pathlib.Path(__file__).parent.resolve(), "data")
        data_file = os.path.join(data_path, "generated_conv_image_urls.jsonl")
        data_convo = pd.read_json(data_file, lines=True)
        with tempfile.TemporaryDirectory() as tmpdir:
            for value in data_convo["messages"]:
                ev_utils._store_multimodal_content(value, tmpdir)

    def test_store_multi_modal_images(self, token_mock, caplog):
        data_path = os.path.join(pathlib.Path(__file__).parent.resolve(), "data")
        data_file = os.path.join(data_path, "generated_conv_images.jsonl")
        data_convo = pd.read_json(data_file, lines=True)
        with tempfile.TemporaryDirectory() as tmpdir:
            for value in data_convo["messages"]:
                ev_utils._store_multimodal_content(value, tmpdir)

    def test_log_metrics_and_instance_results_logs_error(self, token_mock, caplog):
        """Test that we are logging the error when there is no trace destination."""
        logger = logging.getLogger(ev_utils.__name__)
        # All loggers, having promptflow. prefix will have "promptflow" logger
        # as a parent. This logger does not propagate the logs and cannot be
        # captured by caplog. Here we will skip this logger to capture logs.
        logger.parent = logging.root

        with caplog.at_level(logging.DEBUG):
            ev_utils._log_metrics_and_instance_results(
                metrics=None,
                instance_results=None,
                trace_destination=None,
                run=None,
                name_map={},
                evaluation_name=None,
            )
        assert len(caplog.records) == 1
        assert (
            "Skip uploading evaluation results to AI Studio since no trace destination was provided."
            in caplog.records[0].message
        )

    def test_run_broken_if_no_tracking_uri(self, token_mock, caplog):
        """Test that if no tracking URI is provirded, the run is being marked as broken."""
        logger = logging.getLogger(ev_utils.__name__)
        # All loggers, having promptflow. prefix will have "promptflow" logger
        # as a parent. This logger does not propagate the logs and cannot be
        # captured by caplog. Here we will skip this logger to capture logs.
        logger.parent = logging.root
        with caplog.at_level(logging.INFO), EvalRun(
            run_name=None,
            tracking_uri=None,
            subscription_id="mock",
            group_name="mock",
            workspace_name="mock",
            management_client=MagicMock(),
        ) as run:
            assert len(caplog.records) == 1
            assert "The results will be saved locally, but will not be logged to Azure." in caplog.records[0].message
            with patch("azure.ai.evaluation._evaluate._eval_run.EvalRun.request_with_retry") as mock_request:
                run.log_artifact("mock_dir")
                run.log_metric("foo", 42)
                run.write_properties_to_run_history({"foo": "bar"})
            mock_request.assert_not_called()

    @pytest.mark.parametrize(
        "status_code,pf_run",
        [
            (401, False),
            (200, False),
            (401, True),
            (200, True),
        ],
    )
    def test_lifecycle(self, token_mock, status_code, pf_run):
        """Test the run statuses throughout its life cycle."""
        pf_run_mock = None
        if pf_run:
            pf_run_mock = MagicMock()
            pf_run_mock.name = "mock_pf_run"
            pf_run_mock._experiment_name = "mock_pf_experiment"
        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request",
            return_value=self._get_mock_create_response(status_code),
        ):
            run = EvalRun(run_name="test", **TestEvalRun._MOCK_CREDS, promptflow_run=pf_run_mock)
            assert run.status == RunStatus.NOT_STARTED, f"Get {run.status}, expected {RunStatus.NOT_STARTED}"
            run._start_run()
            if status_code == 200 or pf_run:
                assert run.status == RunStatus.STARTED, f"Get {run.status}, expected {RunStatus.STARTED}"
            else:
                assert run.status == RunStatus.BROKEN, f"Get {run.status}, expected {RunStatus.BROKEN}"
            run._end_run("FINISHED")
            if status_code == 200 or pf_run:
                assert run.status == RunStatus.TERMINATED, f"Get {run.status}, expected {RunStatus.TERMINATED}"
            else:
                assert run.status == RunStatus.BROKEN, f"Get {run.status}, expected {RunStatus.BROKEN}"

    def test_local_lifecycle(self, token_mock):
        """Test that the local run have correct statuses."""
        run = EvalRun(
            run_name=None,
            tracking_uri=None,
            subscription_id="mock",
            group_name="mock",
            workspace_name="mock",
            management_client=MagicMock(),
        )
        assert run.status == RunStatus.NOT_STARTED, f"Get {run.status}, expected {RunStatus.NOT_STARTED}"
        run._start_run()
        assert run.status == RunStatus.BROKEN, f"Get {run.status}, expected {RunStatus.BROKEN}"
        run._end_run("FINISHED")
        assert run.status == RunStatus.BROKEN, f"Get {run.status}, expected {RunStatus.BROKEN}"

    @pytest.mark.parametrize("status_code", [200, 401])
    def test_write_properties(self, token_mock, caplog, status_code):
        """Test writing properties to the evaluate run."""
        mock_write = MagicMock()
        mock_write.status_code = status_code
        mock_write.text = lambda: "Mock error"
        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request",
            side_effect=[self._get_mock_create_response(), mock_write, self._get_mock_end_response()],
        ), caplog.at_level(logging.INFO):
            with EvalRun(run_name="test", **TestEvalRun._MOCK_CREDS) as run:
                run.write_properties_to_run_history({"foo": "bar"})
        if status_code != 200:
            assert len(caplog.records) == 1
            assert "Fail writing properties" in caplog.records[0].message
            assert mock_write.text() in caplog.records[0].message
        else:
            assert len(caplog.records) == 0

    def test_write_properties_to_run_history_logs_error(self, token_mock, caplog):
        """Test that we are logging the error when there is no trace destination."""
        logger = logging.getLogger(EvalRun.__module__)
        # All loggers, having promptflow. prefix will have "promptflow" logger
        # as a parent. This logger does not propagate the logs and cannot be
        # captured by caplog. Here we will skip this logger to capture logs.
        logger.parent = logging.root
        with caplog.at_level(logging.INFO), EvalRun(
            run_name=None,
            tracking_uri=None,
            subscription_id="mock",
            group_name="mock",
            workspace_name="mock",
            management_client=MagicMock(),
        ) as run:
            run.write_properties_to_run_history({"foo": "bar"})
        assert len(caplog.records) == 3
        assert "tracking_uri was not provided," in caplog.records[0].message
        assert "Unable to write properties due to Run status=RunStatus.BROKEN." in caplog.records[1].message
        assert "Unable to stop run due to Run status=RunStatus.BROKEN." in caplog.records[2].message

    @pytest.mark.parametrize(
        "function_literal,args,expected_action",
        [
            ("write_properties_to_run_history", ({"foo": "bar"}), "write properties"),
            ("log_metric", ("foo", 42), "log metric"),
            ("log_artifact", ("mock_folder",), "log artifact"),
        ],
    )
    def test_logs_if_not_started(self, token_mock, caplog, function_literal, args, expected_action):
        """Test that all public functions are raising exception if run is not started."""
        logger = logging.getLogger(ev_utils.__name__)
        # All loggers, having promptflow. prefix will have "promptflow" logger
        # as a parent. This logger does not propagate the logs and cannot be
        # captured by caplog. Here we will skip this logger to capture logs.
        logger.parent = logging.root
        run = EvalRun(run_name=None, **TestEvalRun._MOCK_CREDS)
        with caplog.at_level(logging.INFO):
            getattr(run, function_literal)(*args)
        assert len(caplog.records) == 1
        assert expected_action in caplog.records[0].message, caplog.records[0].message
        assert (
            f"Unable to {expected_action} due to Run status=RunStatus.NOT_STARTED" in caplog.records[0].message
        ), caplog.records[0].message

    @pytest.mark.parametrize("status", [RunStatus.STARTED, RunStatus.BROKEN, RunStatus.TERMINATED])
    def test_starting_started_run(self, token_mock, status):
        """Test exception if the run was already started"""
        run = EvalRun(run_name=None, **TestEvalRun._MOCK_CREDS)
        with patch(
            "azure.ai.evaluation._http_utils.HttpPipeline.request",
            return_value=self._get_mock_create_response(500 if status == RunStatus.BROKEN else 200),
        ):
            run._start_run()
            if status == RunStatus.TERMINATED:
                run._end_run("FINISHED")
        with pytest.raises(EvaluationException) as cm:
            run._start_run()
        assert f"Unable to start run due to Run status={status}" in cm.value.args[0], cm.value.args[0]