import inspect
from enum import Enum, EnumMeta
from typing import Callable, Union, get_args, get_origin
from promptflow.contracts.tool import ConnectionType, InputDefinition, ValueType, ToolType
from promptflow.contracts.types import PromptTemplate


def value_to_str(val):
    if val is inspect.Parameter.empty:
        # For empty case, default field will be skipped when dumping to json
        return None
    if val is None:
        # Dump default: "" in json to avoid UI validation error
        return ""
    if isinstance(val, Enum):
        return val.value
    return str(val)


def resolve_annotation(anno) -> Union[str, list]:
    """Resolve the union annotation to type list."""
    origin = get_origin(anno)
    if origin != Union:
        return anno
    # Optional[Type] is Union[Type, NoneType], filter NoneType out
    args = [arg for arg in get_args(anno) if arg != type(None)]  # noqa: E721
    return args[0] if len(args) == 1 else args


def param_to_definition(param, value_type) -> (InputDefinition, bool):
    default_value = param.default
    enum = None
    custom_type = None
    # Get value type and enum from default if no annotation
    if default_value is not inspect.Parameter.empty and value_type == inspect.Parameter.empty:
        value_type = default_value.__class__ if isinstance(default_value, Enum) else type(default_value)
    # Extract enum for enum class
    if isinstance(value_type, EnumMeta):
        enum = [str(option.value) for option in value_type]
        value_type = str
    is_connection = False
    if ConnectionType.is_connection_value(value_type):
        if ConnectionType.is_custom_strong_type(value_type):
            typ = ["CustomConnection"]
            custom_type = [value_type.__name__]
        else:
            typ = [value_type.__name__]
        is_connection = True
    elif isinstance(value_type, list):
        if not all(ConnectionType.is_connection_value(t) for t in value_type):
            typ = [ValueType.OBJECT]
        else:
            custom_connection_added = False
            typ = []
            custom_type = []
            for t in value_type:
                if ConnectionType.is_custom_strong_type(t):
                    if not custom_connection_added:
                        custom_connection_added = True
                        typ.append("CustomConnection")
                    custom_type.append(t.__name__)
                else:
                    typ.append(t.__name__)
            is_connection = True
    else:
        typ = [ValueType.from_type(value_type)]
    return InputDefinition(type=typ, default=value_to_str(default_value),
                           description=None, enum=enum, custom_type=custom_type), is_connection


def function_to_interface(f: Callable, tool_type, initialize_inputs=None) -> tuple:
    sign = inspect.signature(f)
    all_inputs = {}
    input_defs = {}
    connection_types = []
    # Initialize the counter for prompt template
    prompt_template_count = 0
    # Collect all inputs from class and func
    if initialize_inputs:
        if any(k for k in initialize_inputs if k in sign.parameters):
            raise Exception(f'Duplicate inputs found from {f.__name__!r} and "__init__()"!')
        all_inputs = {**initialize_inputs}
    all_inputs.update(
        {
            k: v
            for k, v in sign.parameters.items()
            if k != "self" and v.kind != v.VAR_KEYWORD and v.kind != v.VAR_POSITIONAL  # TODO: Handle these cases
        }
    )
    # Resolve inputs to definitions.
    for k, v in all_inputs.items():
        # Get value type from annotation
        value_type = resolve_annotation(v.annotation)
        if value_type is PromptTemplate:
            # custom llm tool has prompt template as input, skip it
            prompt_template_count += 1
            continue
        input_def, is_connection = param_to_definition(v, value_type)
        input_defs[k] = input_def
        if is_connection:
            connection_types.append(input_def.type)

    # Check PromptTemplate input:
    # a. For custom llm tool, there should be exactly one PromptTemplate input
    # b. For python tool, PromptTemplate input is not supported
    if tool_type == ToolType.PYTHON and prompt_template_count > 0:
        raise Exception(f"Input of type 'PromptTemplate' not supported in python tool '{f.__name__}'. ")

    if tool_type == ToolType.CUSTOM_LLM and prompt_template_count == 0:
        raise Exception(f"No input of type 'PromptTemplate' was found in custom llm tool '{f.__name__}'. ")

    if tool_type == ToolType.CUSTOM_LLM and prompt_template_count > 1:
        raise Exception(f"Multiple inputs of type 'PromptTemplate' were found in '{f.__name__}'. "
                        "Only one input of this type is expected.")

    outputs = {}
    # Note: We don't have output definition now
    # outputs = {"output": OutputDefinition("output", [ValueType.from_type(type(sign.return_annotation))], "", True)}
    # if is_dataclass(sign.return_annotation):
    #     for f in fields(sign.return_annotation):
    #         outputs[f.name] = OutputDefinition(f.name, [ValueType.from_type(
    #             type(getattr(sign.return_annotation, f.name)))], "", False)
    return input_defs, outputs, connection_types