# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------

# pylint: disable=protected-access

import json
import logging
import os
import re
import subprocess
import sys
import time
from typing import Any, Dict, Iterable, List, Optional, TextIO, Union

from azure.ai.ml._artifacts._artifact_utilities import get_datastore_info, list_logs_in_datastore
from azure.ai.ml._restclient.runhistory.models import Run, RunDetails, TypedAssetReference
from azure.ai.ml._restclient.v2022_02_01_preview.models import DataType
from azure.ai.ml._restclient.v2022_02_01_preview.models import JobType as RestJobType
from azure.ai.ml._restclient.v2022_02_01_preview.models import ModelType
from azure.ai.ml._restclient.v2022_10_01.models import JobBase
from azure.ai.ml._utils._http_utils import HttpPipeline
from azure.ai.ml._utils.utils import create_requests_pipeline_with_retry, download_text_from_url
from azure.ai.ml.constants._common import GitProperties
from azure.ai.ml.constants._job.job import JobLogPattern, JobType
from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, JobException
from azure.ai.ml.operations._dataset_dataplane_operations import DatasetDataplaneOperations
from azure.ai.ml.operations._datastore_operations import DatastoreOperations
from azure.ai.ml.operations._model_dataplane_operations import ModelDataplaneOperations
from azure.ai.ml.operations._run_history_constants import JobStatus, RunHistoryConstants
from azure.ai.ml.operations._run_operations import RunOperations

STATUS_KEY = "status"

module_logger = logging.getLogger(__name__)


def _get_sorted_filtered_logs(
    logs_iterable: Iterable[str],
    job_type: str,
    processed_logs: Optional[Dict[str, int]] = None,
    only_streamable: bool = True,
) -> List[str]:
    """Filters log file names, sorts, and returns list starting with where we left off last iteration.

    :param logs_iterable: An iterable of log paths.
    :type logs_iterable: Iterable[str]
    :param job_type: the job type to filter log files
    :type job_type: str
    :param processed_logs: dictionary tracking the state of how many lines of each file have been written out
    :type processed_logs: dict[str, int]
    :param only_streamable: Whether to only get streamable logs
    :type only_streamable: bool
    :return: List of logs to continue from
    :rtype: list[str]
    """
    processed_logs = processed_logs if processed_logs else {}
    # First, attempt to read logs in new Common Runtime form
    output_logs_pattern = (
        JobLogPattern.COMMON_RUNTIME_STREAM_LOG_PATTERN
        if only_streamable
        else JobLogPattern.COMMON_RUNTIME_ALL_USER_LOG_PATTERN
    )
    logs = list(logs_iterable)
    filtered_logs = [x for x in logs if re.match(output_logs_pattern, x)]

    # fall back to legacy log format
    if filtered_logs is None or len(filtered_logs) == 0:
        job_type = job_type.lower()
        if job_type in JobType.COMMAND:
            output_logs_pattern = JobLogPattern.COMMAND_JOB_LOG_PATTERN
        elif job_type in JobType.PIPELINE:
            output_logs_pattern = JobLogPattern.PIPELINE_JOB_LOG_PATTERN
        elif job_type in JobType.SWEEP:
            output_logs_pattern = JobLogPattern.SWEEP_JOB_LOG_PATTERN

    filtered_logs = [x for x in logs if re.match(output_logs_pattern, x)]
    filtered_logs.sort()
    previously_printed_index = 0
    for i, v in enumerate(filtered_logs):
        if processed_logs.get(v):
            previously_printed_index = i
        else:
            break
    # Slice inclusive from the last printed log (can be updated before printing new files)
    return filtered_logs[previously_printed_index:]


def _incremental_print(log: str, processed_logs: Dict[str, int], current_log_name: str, fileout: TextIO) -> None:
    """Incremental print.

    :param log:
    :type log: str
    :param processed_logs: The record of how many lines have been written for each log file
    :type processed_logs: dict[str, int]
    :param current_log_name: the file name being read out, used in header writing and accessing processed_logs
    :type current_log_name: str
    :param fileout:
    :type fileout: TestIOWrapper
    """
    lines = log.splitlines()
    doc_length = len(lines)
    if doc_length == 0:
        # If a file is empty, skip writing out.
        # This addresses issue where batch endpoint jobs can create log files before they are needed.
        return
    previous_printed_lines = processed_logs.get(current_log_name, 0)
    # when a new log is first being written to console, print spacing and label
    if previous_printed_lines == 0:
        fileout.write("\n")
        fileout.write("Streaming " + current_log_name + "\n")
        fileout.write("=" * (len(current_log_name) + 10) + "\n")
        fileout.write("\n")
    # otherwise, continue to log the file where we left off
    for line in lines[previous_printed_lines:]:
        fileout.write(line + "\n")
    # update state to record number of lines written for this log file
    processed_logs[current_log_name] = doc_length


def _get_last_log_primary_instance(logs: List) -> Any:
    """Return last log for primary instance.

    :param logs:
    :type logs: builtin.list
    :return: Returns the last log primary instance.
    :rtype:
    """
    primary_ranks = ["rank_0", "worker_0"]
    rank_match_re = re.compile(r"(.*)_(.*?_.*?)\.txt")
    last_log_name = logs[-1]

    last_log_match = rank_match_re.match(last_log_name)
    if not last_log_match:
        return last_log_name

    last_log_prefix = last_log_match.group(1)
    matching_logs = sorted(filter(lambda x: x.startswith(last_log_prefix), logs))

    # we have some specific ranks that denote the primary, use those if found
    for log_name in matching_logs:
        match = rank_match_re.match(log_name)
        if not match:
            continue
        if match.group(2) in primary_ranks:
            return log_name

    # no definitively primary instance, just return the highest sorted
    return matching_logs[0]


def _wait_before_polling(current_seconds: float) -> int:
    if current_seconds < 0:
        msg = "current_seconds must be positive"
        raise JobException(
            message=msg,
            target=ErrorTarget.JOB,
            no_personal_data_message=msg,
            error_category=ErrorCategory.USER_ERROR,
        )
    import math

    # Sigmoid that tapers off near the_get_logs max at ~ 3 min
    duration = int(
        int(RunHistoryConstants._WAIT_COMPLETION_POLLING_INTERVAL_MAX) / (1.0 + 100 * math.exp(-current_seconds / 20.0))
    )
    return max(int(RunHistoryConstants._WAIT_COMPLETION_POLLING_INTERVAL_MIN), duration)


def list_logs(run_operations: RunOperations, job_resource: JobBase) -> Dict:
    details: RunDetails = run_operations.get_run_details(job_resource.name)
    logs_dict = details.log_files
    keys = _get_sorted_filtered_logs(logs_dict, job_resource.properties.job_type)
    return {key: logs_dict[key] for key in keys}


# pylint: disable=too-many-statements,too-many-locals
def stream_logs_until_completion(
    run_operations: RunOperations,
    job_resource: JobBase,
    datastore_operations: Optional[DatastoreOperations] = None,
    raise_exception_on_failed_job: bool = True,
    *,
    requests_pipeline: HttpPipeline
) -> None:
    """Stream the experiment run output to the specified file handle. By default the the file handle points to stdout.

    :param run_operations: The run history operations class.
    :type run_operations: RunOperations
    :param job_resource: The job to stream
    :type job_resource: JobBase
    :param datastore_operations: Optional, the datastore operations class, used to get logs from datastore
    :type datastore_operations: Optional[DatastoreOperations]
    :param raise_exception_on_failed_job: Should this method fail if job fails
    :type raise_exception_on_failed_job: Boolean
    :keyword requests_pipeline: The HTTP pipeline to use for requests.
    :type requests_pipeline: ~azure.ai.ml._utils._http_utils.HttpPipeline
    :return:
    :rtype: None
    """
    job_type = job_resource.properties.job_type
    job_name = job_resource.name
    studio_endpoint = job_resource.properties.services.get("Studio", None)
    studio_endpoint = studio_endpoint.endpoint if studio_endpoint else None
    # Feature store jobs should be linked to the Feature Store Workspace UI.
    # Todo: Consolidate this logic to service side
    if "azureml.FeatureStoreJobType" in job_resource.properties.properties:
        url_format = (
            "https://ml.azure.com/featureStore/{fs_name}/featureSets/{fset_name}/{fset_version}/matJobs/"
            "jobs/{run_id}?wsid=/subscriptions/{fs_sub_id}/resourceGroups/{fs_rg_name}/providers/"
            "Microsoft.MachineLearningServices/workspaces/{fs_name}"
        )
        studio_endpoint = url_format.format(
            fs_name=job_resource.properties.properties["azureml.FeatureStoreName"],
            fs_sub_id=run_operations._subscription_id,
            fs_rg_name=run_operations._resource_group_name,
            fset_name=job_resource.properties.properties["azureml.FeatureSetName"],
            fset_version=job_resource.properties.properties["azureml.FeatureSetVersion"],
            run_id=job_name,
        )
    elif "FeatureStoreJobType" in job_resource.properties.properties:
        url_format = (
            "https://ml.azure.com/featureStore/{fs_name}/featureSets/{fset_name}/{fset_version}/matJobs/"
            "jobs/{run_id}?wsid=/subscriptions/{fs_sub_id}/resourceGroups/{fs_rg_name}/providers/"
            "Microsoft.MachineLearningServices/workspaces/{fs_name}"
        )
        studio_endpoint = url_format.format(
            fs_name=job_resource.properties.properties["FeatureStoreName"],
            fs_sub_id=run_operations._subscription_id,
            fs_rg_name=run_operations._resource_group_name,
            fset_name=job_resource.properties.properties["FeatureSetName"],
            fset_version=job_resource.properties.properties["FeatureSetVersion"],
            run_id=job_name,
        )

    file_handle = sys.stdout
    ds_properties = None
    prefix = None
    if (
        hasattr(job_resource.properties, "outputs")
        and job_resource.properties.job_type != RestJobType.AUTO_ML
        and datastore_operations
    ):
        # Get default output location

        default_output = (
            job_resource.properties.outputs.get("default", None) if job_resource.properties.outputs else None
        )
        is_uri_folder = default_output and default_output.job_output_type == DataType.URI_FOLDER
        if is_uri_folder:
            output_uri = default_output.uri  # type: ignore
            # Parse the uri format
            output_uri = output_uri.split("datastores/")[1]
            datastore_name, prefix = output_uri.split("/", 1)
            ds_properties = get_datastore_info(datastore_operations, datastore_name)

    try:
        file_handle.write("RunId: {}\n".format(job_name))
        file_handle.write("Web View: {}\n".format(studio_endpoint))

        _current_details: RunDetails = run_operations.get_run_details(job_name)

        processed_logs: Dict = {}

        poll_start_time = time.time()
        pipeline_with_retries = create_requests_pipeline_with_retry(requests_pipeline=requests_pipeline)
        while (
            _current_details.status in RunHistoryConstants.IN_PROGRESS_STATUSES
            or _current_details.status == JobStatus.FINALIZING
        ):
            file_handle.flush()
            time.sleep(_wait_before_polling(time.time() - poll_start_time))
            _current_details = run_operations.get_run_details(job_name)  # TODO use FileWatcher
            if job_type.lower() in JobType.PIPELINE:
                legacy_folder_name = "/logs/azureml/"
            else:
                legacy_folder_name = "/azureml-logs/"
            _current_logs_dict = (
                list_logs_in_datastore(
                    ds_properties,
                    prefix=str(prefix),
                    legacy_log_folder_name=legacy_folder_name,
                )
                if ds_properties is not None
                else _current_details.log_files
            )
            # Get the list of new logs available after filtering out the processed ones
            available_logs = _get_sorted_filtered_logs(_current_logs_dict, job_type, processed_logs)
            content = ""
            for current_log in available_logs:
                content = download_text_from_url(
                    _current_logs_dict[current_log],
                    pipeline_with_retries,
                    timeout=RunHistoryConstants._DEFAULT_GET_CONTENT_TIMEOUT,
                )

                _incremental_print(content, processed_logs, current_log, file_handle)

            # TODO: Temporary solution to wait for all the logs to be printed in the finalizing state.
            if (
                _current_details.status not in RunHistoryConstants.IN_PROGRESS_STATUSES
                and _current_details.status == JobStatus.FINALIZING
                and "The activity completed successfully. Finalizing run..." in content
            ):
                break

        file_handle.write("\n")
        file_handle.write("Execution Summary\n")
        file_handle.write("=================\n")
        file_handle.write("RunId: {}\n".format(job_name))
        file_handle.write("Web View: {}\n".format(studio_endpoint))

        warnings = _current_details.warnings
        if warnings:
            messages = [x.message for x in warnings if x.message]
            if len(messages) > 0:
                file_handle.write("\nWarnings:\n")
                for message in messages:
                    file_handle.write(message + "\n")
                file_handle.write("\n")

        if _current_details.status == JobStatus.FAILED:
            error = (
                _current_details.error.as_dict()
                if _current_details.error
                else "Detailed error not set on the Run. Please check the logs for details."
            )
            # If we are raising the error later on, so we don't double print.
            if not raise_exception_on_failed_job:
                file_handle.write("\nError:\n")
                file_handle.write(json.dumps(error, indent=4))
                file_handle.write("\n")
            else:
                raise JobException(
                    message="Exception : \n {} ".format(json.dumps(error, indent=4)),
                    target=ErrorTarget.JOB,
                    no_personal_data_message="Exception raised on failed job.",
                    error_category=ErrorCategory.SYSTEM_ERROR,
                )

        file_handle.write("\n")
        file_handle.flush()
    except KeyboardInterrupt as e:
        error_message = (
            "The output streaming for the run interrupted.\n"
            "But the run is still executing on the compute target. \n"
            "Details for canceling the run can be found here: "
            "https://aka.ms/aml-docs-cancel-run"
        )
        raise JobException(
            message=error_message,
            target=ErrorTarget.JOB,
            no_personal_data_message=error_message,
            error_category=ErrorCategory.USER_ERROR,
        ) from e


def get_git_properties() -> Dict[str, str]:
    """Gather Git tracking info from the local environment.

    :return: Properties dictionary.
    :rtype: dict
    """

    def _clean_git_property_bool(value: Any) -> Optional[bool]:
        if value is None:
            return None
        return str(value).strip().lower() in ["true", "1"]

    def _clean_git_property_str(value: Any) -> Optional[str]:
        if value is None:
            return None
        return str(value).strip() or None

    def _run_git_cmd(args: Iterable[str]) -> Optional[str]:
        """Runs git with the provided arguments

        :param args: A iterable of arguments for a git command. Should not include leading "git"
        :type args: Iterable[str]
        :return: The output of running git with arguments, or None if it fails.
        :rtype: Optional[str]
        """
        try:
            with open(os.devnull, "wb") as devnull:
                return subprocess.check_output(["git"] + list(args), stderr=devnull).decode()
        except KeyboardInterrupt:
            raise
        except BaseException:  # pylint: disable=W0718
            return None

    # Check for environment variable overrides.
    repository_uri = os.environ.get(GitProperties.ENV_REPOSITORY_URI, None)
    branch = os.environ.get(GitProperties.ENV_BRANCH, None)
    commit = os.environ.get(GitProperties.ENV_COMMIT, None)
    dirty: Optional[Union[str, bool]] = os.environ.get(GitProperties.ENV_DIRTY, None)
    build_id = os.environ.get(GitProperties.ENV_BUILD_ID, None)
    build_uri = os.environ.get(GitProperties.ENV_BUILD_URI, None)

    is_git_repo = _run_git_cmd(["rev-parse", "--is-inside-work-tree"])
    if _clean_git_property_bool(is_git_repo):
        repository_uri = repository_uri or _run_git_cmd(["ls-remote", "--get-url"])
        branch = branch or _run_git_cmd(["symbolic-ref", "--short", "HEAD"])
        commit = commit or _run_git_cmd(["rev-parse", "HEAD"])
        dirty = dirty or _run_git_cmd(["status", "--porcelain", "."]) and True

    # Parsing logic.
    repository_uri = _clean_git_property_str(repository_uri)
    commit = _clean_git_property_str(commit)
    branch = _clean_git_property_str(branch)
    dirty = _clean_git_property_bool(dirty)
    build_id = _clean_git_property_str(build_id)
    build_uri = _clean_git_property_str(build_uri)

    # Return with appropriate labels.
    properties = {}

    if repository_uri is not None:
        properties[GitProperties.PROP_MLFLOW_GIT_REPO_URL] = repository_uri

    if branch is not None:
        properties[GitProperties.PROP_MLFLOW_GIT_BRANCH] = branch

    if commit is not None:
        properties[GitProperties.PROP_MLFLOW_GIT_COMMIT] = commit

    if dirty is not None:
        properties[GitProperties.PROP_DIRTY] = str(dirty)

    if build_id is not None:
        properties[GitProperties.PROP_BUILD_ID] = build_id

    if build_uri is not None:
        properties[GitProperties.PROP_BUILD_URI] = build_uri

    return properties


def get_job_output_uris_from_dataplane(
    job_name: Optional[str],
    run_operations: RunOperations,
    dataset_dataplane_operations: DatasetDataplaneOperations,
    model_dataplane_operations: Optional[ModelDataplaneOperations],
    output_names: Optional[Union[Iterable[str], str]] = None,
) -> Dict[str, str]:
    """Returns the output path for the given output in cloud storage of the given job.

    If no output names are given, the output paths for all outputs will be returned.
    URIs obtained from the service will be in the long-form azureml:// format.

    For example:
    azureml://subscriptions/<sub>/resource[gG]roups/<rg_name>/workspaces/<ws_name>/datastores/<ds_name>/paths/<ds_path>

    :param job_name: The job name
    :type job_name: str
    :param run_operations: The RunOperations used to fetch run data for the job
    :type run_operations: RunOperations
    :param dataset_dataplane_operations: The DatasetDataplaneOperations used to fetch dataset uris
    :type dataset_dataplane_operations: DatasetDataplaneOperations
    :param model_dataplane_operations:  The ModelDataplaneOperations used to fetch dataset uris
    :type model_dataplane_operations: ModelDataplaneOperations
    :param output_names: The output name(s) to fetch. If not specified, retrieves all.
    :type output_names: Optional[Union[Iterable[str] str]]
    :return: Dictionary mapping user-defined output name to output uri
    :rtype: Dict[str, str]
    """
    run_metadata: Run = run_operations.get_run_data(str(job_name)).run_metadata
    run_outputs: Dict[str, TypedAssetReference] = run_metadata.outputs or {}

    # Create a reverse mapping from internal asset id to user-defined output name
    asset_id_to_output_name = {v.asset_id: k for k, v in run_outputs.items()}
    if not output_names:
        # Assume all outputs are needed if no output name is provided
        output_names = run_outputs.keys()
    else:
        if isinstance(output_names, str):
            output_names = [output_names]
        output_names = [o for o in output_names if o in run_outputs]

    # Collect all output ids that correspond to data assets
    dataset_ids = [
        run_outputs[output_name].asset_id
        for output_name in output_names
        if run_outputs[output_name].type in [o.value for o in DataType]
    ]

    # Collect all output ids that correspond to models
    model_ids = [
        run_outputs[output_name].asset_id
        for output_name in output_names
        if run_outputs[output_name].type in [o.value for o in ModelType]
    ]

    output_name_to_dataset_uri = {}
    if dataset_ids:
        # Get the data paths from the service
        dataset_uris = dataset_dataplane_operations.get_batch_dataset_uris(dataset_ids)
        # Map the user-defined output name to the output uri
        # The service returns a mapping from internal asset id to output metadata, so we need the reverse map
        # defined above to get the user-defined output name from the internal asset id.
        output_name_to_dataset_uri = {asset_id_to_output_name[k]: v.uri for k, v in dataset_uris.values.items()}

    # This is a repeat of the logic above for models.
    output_name_to_model_uri = {}
    if model_ids:
        model_uris = (
            model_dataplane_operations.get_batch_model_uris(model_ids)
            if model_dataplane_operations is not None
            else None
        )
        output_name_to_model_uri = {
            asset_id_to_output_name[k]: v.path for k, v in model_uris.values.items()  # type: ignore
        }
    return {**output_name_to_dataset_uri, **output_name_to_model_uri}