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

# pylint: disable=protected-access

import inspect
import logging
from collections import OrderedDict
from functools import wraps
from inspect import Parameter, signature
from pathlib import Path
from typing import Callable, Dict, List, Optional, TypeVar, Union, overload

from typing_extensions import ParamSpec

from azure.ai.ml._utils.utils import is_private_preview_enabled
from azure.ai.ml.entities import Data, Model, PipelineJob, PipelineJobSettings
from azure.ai.ml.entities._builders.pipeline import Pipeline
from azure.ai.ml.entities._inputs_outputs import Input, is_group
from azure.ai.ml.entities._job.pipeline._io import NodeOutput, PipelineInput, _GroupAttrDict
from azure.ai.ml.entities._job.pipeline._pipeline_expression import PipelineExpression
from azure.ai.ml.exceptions import (
    MultipleValueError,
    ParamValueNotExistsError,
    TooManyPositionalArgsError,
    UnexpectedKeywordError,
    UnsupportedParameterKindError,
    UserErrorException,
)

from azure.ai.ml.entities._builders import BaseNode
from azure.ai.ml.dsl._pipeline_component_builder import PipelineComponentBuilder, _is_inside_dsl_pipeline_func
from azure.ai.ml.dsl._pipeline_decorator import _validate_args
from azure.ai.ml.dsl._settings import _dsl_settings_stack
from azure.ai.ml.dsl._utils import _resolve_source_file

SUPPORTED_INPUT_TYPES = (
    PipelineInput,
    NodeOutput,
    Input,
    Model,
    Data,  # For the case use a Data object as an input, we will convert it to Input object
    Pipeline,  # For the case use a pipeline node as the input, we use its only one output as the real input.
    str,
    bool,
    int,
    float,
    PipelineExpression,
    _GroupAttrDict,
)
module_logger = logging.getLogger(__name__)

T = TypeVar("T")
P = ParamSpec("P")


# Overload the returns a decorator when func is None
@overload
def pipeline(  # type: ignore[misc]
    # TODO: Bug 2875164
    func: None = None,
    *,
    name: Optional[str] = None,
    version: Optional[str] = None,
    display_name: Optional[str] = None,
    description: Optional[str] = None,
    experiment_name: Optional[str] = None,
    tags: Optional[Dict[str, str]] = None,
    **kwargs,
) -> Callable[[Callable[P, T]], Callable[P, PipelineJob]]:
    ...


# Overload the returns a decorated function when func isn't None
@overload
def pipeline(
    func: Optional[Callable[P, T]] = None,
    *,
    name: Optional[str] = None,
    version: Optional[str] = None,
    display_name: Optional[str] = None,
    description: Optional[str] = None,
    experiment_name: Optional[str] = None,
    tags: Optional[Dict[str, str]] = None,
    **kwargs,
) -> Callable[P, PipelineJob]:
    ...


def pipeline(
    func: Optional[Callable[P, T]] = None,
    *,
    name: Optional[str] = None,
    version: Optional[str] = None,
    display_name: Optional[str] = None,
    description: Optional[str] = None,
    experiment_name: Optional[str] = None,
    tags: Optional[Dict[str, str]] = None,
    **kwargs,
) -> Union[Callable[[Callable[P, T]], Callable[P, PipelineJob]], Callable[P, PipelineJob]]:
    """Build a pipeline which contains all component nodes defined in this function.

    :param func: The user pipeline function to be decorated.
    :type func: types.FunctionType
    :keyword name: The name of pipeline component, defaults to function name.
    :paramtype name: str
    :keyword version: The version of pipeline component, defaults to "1".
    :paramtype version: str
    :keyword display_name: The display name of pipeline component, defaults to function name.
    :paramtype display_name: str
    :keyword description: The description of the built pipeline.
    :paramtype description: str
    :keyword experiment_name: Name of the experiment the job will be created under, \
                if None is provided, experiment will be set to current directory.
    :paramtype experiment_name: str
    :keyword tags: The tags of pipeline component.
    :paramtype tags: dict[str, str]
    :keyword kwargs: A dictionary of additional configuration parameters.
    :paramtype kwargs: dict

    .. admonition:: Example:

        .. literalinclude:: ../../../../samples/ml_samples_pipeline_job_configurations.py
            :start-after: [START configure_pipeline]
            :end-before: [END configure_pipeline]
            :language: python
            :dedent: 8
            :caption: Shows how to create a pipeline using this decorator.
    :return: Either
      * A decorator, if `func` is None
      * The decorated `func`
    :rtype: Union[
        Callable[[Callable], Callable[..., PipelineJob]],
        Callable[P, PipelineJob]
      ]
    """
    get_component = kwargs.get("get_component", False)

    def pipeline_decorator(func: Callable[P, T]) -> Callable[P, PipelineJob]:
        # pylint: disable=isinstance-second-argument-not-valid-type
        if not isinstance(func, Callable):  # type: ignore
            raise UserErrorException(f"Dsl pipeline decorator accept only function type, got {type(func)}.")

        non_pipeline_inputs = kwargs.get("non_pipeline_inputs", []) or kwargs.get("non_pipeline_parameters", [])
        # compute variable names changed from default_compute_targe -> compute -> default_compute -> none
        # to support legacy usage, we support them with priority.
        compute = kwargs.get("compute", None)
        default_compute_target = kwargs.get("default_compute_target", None)
        default_compute_target = kwargs.get("default_compute", None) or default_compute_target
        continue_on_step_failure = kwargs.get("continue_on_step_failure", None)
        on_init = kwargs.get("on_init", None)
        on_finalize = kwargs.get("on_finalize", None)

        default_datastore = kwargs.get("default_datastore", None)
        force_rerun = kwargs.get("force_rerun", None)
        job_settings = {
            "default_datastore": default_datastore,
            "continue_on_step_failure": continue_on_step_failure,
            "force_rerun": force_rerun,
            "default_compute": default_compute_target,
            "on_init": on_init,
            "on_finalize": on_finalize,
        }
        func_entry_path = _resolve_source_file()
        if not func_entry_path:
            func_path = Path(inspect.getfile(func))
            # in notebook, func_path may be a fake path and will raise error when trying to resolve this fake path
            if func_path.exists():
                func_entry_path = func_path.resolve().absolute()

        job_settings = {k: v for k, v in job_settings.items() if v is not None}
        pipeline_builder = PipelineComponentBuilder(
            func=func,
            name=name,
            version=version,
            display_name=display_name,
            description=description,
            default_datastore=default_datastore,
            tags=tags,
            source_path=str(func_entry_path),
            non_pipeline_inputs=non_pipeline_inputs,
        )

        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> PipelineJob:
            # Default args will be added here.
            # pylint: disable=abstract-class-instantiated
            # Node: push/pop stack here instead of put it inside build()
            # Because we only want to enable dsl settings on top level pipeline
            _dsl_settings_stack.push()  # use this stack to track on_init/on_finalize settings
            try:
                # Convert args to kwargs
                provided_positional_kwargs = _validate_args(func, args, kwargs, non_pipeline_inputs)

                # When pipeline supports variable params, update pipeline component to support the inputs in **kwargs.
                pipeline_parameters = {
                    k: v for k, v in provided_positional_kwargs.items() if k not in non_pipeline_inputs
                }
                pipeline_builder._update_inputs(pipeline_parameters)

                non_pipeline_params_dict = {
                    k: v for k, v in provided_positional_kwargs.items() if k in non_pipeline_inputs
                }

                # TODO: cache built pipeline component
                pipeline_component = pipeline_builder.build(
                    user_provided_kwargs=provided_positional_kwargs,
                    non_pipeline_inputs_dict=non_pipeline_params_dict,
                    non_pipeline_inputs=non_pipeline_inputs,
                )
            finally:
                # use `finally` to ensure pop operation from the stack
                dsl_settings = _dsl_settings_stack.pop()

            # update on_init/on_finalize settings if init/finalize job is set
            if dsl_settings.init_job_set:
                job_settings["on_init"] = dsl_settings.init_job_name(pipeline_component.jobs)
            if dsl_settings.finalize_job_set:
                job_settings["on_finalize"] = dsl_settings.finalize_job_name(pipeline_component.jobs)

            # TODO: pass compute & default_compute separately?
            common_init_args = {
                "experiment_name": experiment_name,
                "component": pipeline_component,
                "inputs": pipeline_parameters,
                "tags": tags,
            }
            if _is_inside_dsl_pipeline_func() or get_component:
                # on_init/on_finalize is not supported for pipeline component
                if job_settings.get("on_init") is not None or job_settings.get("on_finalize") is not None:
                    raise UserErrorException("On_init/on_finalize is not supported for pipeline component.")
                # Build pipeline node instead of pipeline job if inside dsl.
                built_pipeline = Pipeline(_from_component_func=True, **common_init_args)
                if job_settings:
                    module_logger.warning(
                        ("Job settings %s on pipeline function %r are ignored when using inside PipelineJob."),
                        job_settings,
                        func.__name__,
                    )
            else:
                built_pipeline = PipelineJob(
                    jobs=pipeline_component.jobs,
                    compute=compute,
                    settings=PipelineJobSettings(**job_settings),
                    **common_init_args,
                )

            return built_pipeline
        
        wrapper._is_dsl_func = True  # type: ignore[attr-defined]
        wrapper._job_settings = job_settings  # type: ignore[attr-defined]
        wrapper._pipeline_builder = pipeline_builder  # type: ignore[attr-defined]
        return wrapper

    # enable use decorator without "()" if all arguments are default values
    if func is not None:
        return pipeline_decorator(func)
    return pipeline_decorator
