import inspect
from collections.abc import Generator, Iterator, Sequence
from dataclasses import dataclass
from typing import Callable, Optional, Union, get_args, get_origin

from dvcx.lib.feature import Feature, FeatureType, FeatureTypeNames
from dvcx.lib.signal_schema import SignalSchema
from dvcx.lib.utils import DvcxParamsError


class UdfSignatureError(DvcxParamsError):
    def __init__(self, chain: str, msg):
        suffix = f"(dataset '{chain}')" if chain else ""
        super().__init__(f"processor signature error{suffix}: {msg}")


@dataclass
class UdfSignature:
    func: Callable
    params: Sequence[str]
    output_schema: SignalSchema

    @classmethod
    def parse(
        cls,
        chain: str,
        signal_map: dict[str, Callable],
        func: Optional[Callable] = None,
        params: Union[None, str, Sequence[str]] = None,
        output: Union[None, FeatureType, Sequence[str], dict[str, FeatureType]] = None,
        is_generator: bool = True,
    ) -> "UdfSignature":
        keys = ", ".join(signal_map.keys())
        if len(signal_map) > 1:
            raise UdfSignatureError(
                chain,
                f"multiple signals '{keys}' are not supported in processors."
                f" Chain multiple processors instead.",
            )
        if len(signal_map) == 1:
            if func is not None:
                raise UdfSignatureError(
                    chain,
                    f"processor can't have signal '{keys}' with function '{func}'",
                )
            signal_name, udf_func = next(iter(signal_map.items()))
        else:
            if func is None:
                raise UdfSignatureError(chain, "user function is not defined")
            udf_func = func
            signal_name = None
        if not callable(udf_func):
            raise UdfSignatureError(chain, f"function '{func}' is not callable")
        func_params_map_sign, func_outs_sign, is_iterator = (
            UdfSignature._func_signature(chain, udf_func)
        )
        if params:
            udf_params = [params] if isinstance(params, str) else params
        elif not func_params_map_sign:
            udf_params = []
        else:
            udf_params = list(func_params_map_sign.keys())
        if output:
            udf_output_map = UdfSignature._validate_output(
                chain, signal_name, func, func_outs_sign, output
            )
        else:
            if not func_outs_sign:
                raise UdfSignatureError(
                    chain,
                    f"outputs are not defined in function '{udf_func.__name__}'"
                    f" hints or 'output'",
                )

            if not signal_name:
                raise UdfSignatureError(
                    chain,
                    "signal name is not specified."
                    " Define it as signal name 's1=func() or in 'output'",
                )

            if is_generator and not is_iterator:
                raise UdfSignatureError(
                    chain,
                    f"function '{func}' cannot be used in generator/aggregator"
                    f" because it returns a type that is not Iterator/Generator."
                    f" Instead, it returns '{func_outs_sign}'",
                )

            if isinstance(func_outs_sign, tuple):
                udf_output_map = {
                    signal_name + f"_{num}": typ
                    for num, typ in enumerate(func_outs_sign)
                }
            else:
                udf_output_map = {signal_name: func_outs_sign[0]}

        return cls(
            func=udf_func,
            params=udf_params,
            output_schema=SignalSchema(udf_output_map),
        )

    @staticmethod
    def _validate_output(chain, signal_name, func, func_outs_sign, output):
        if isinstance(output, str):
            output = [output]
        if isinstance(output, Sequence):
            if not func_outs_sign:
                raise UdfSignatureError(
                    chain,
                    "output types are not specified. Specify types in 'output' as"
                    " a dict or as function return value hint.",
                )

            if len(func_outs_sign) != len(output):
                raise UdfSignatureError(
                    chain,
                    f"length of outputs names ({len(output)}) and function '{func}'"
                    f" return type length ({len(func_outs_sign)}) does not match",
                )

            udf_output_map = dict(zip(output, func_outs_sign))
        elif isinstance(output, dict):
            for key, value in output.items():
                if not isinstance(key, str):
                    raise UdfSignatureError(
                        chain,
                        f"output signal '{key}' has type '{type(key)}'"
                        f" while 'str' is expected",
                    )
                if not Feature.is_feature_type(value):
                    raise UdfSignatureError(
                        chain,
                        f"output type '{value.__name__}' of signal '{key}' is not"
                        f" supported. Please use Feature types: {FeatureTypeNames}",
                    )

            udf_output_map = output
        elif Feature.is_feature_type(output):
            udf_output_map = {signal_name: output}
        else:
            raise UdfSignatureError(
                chain,
                f"unknown output type: {output}. List of signals or dict of signals"
                f" to function are expected.",
            )
        return udf_output_map

    def __eq__(self, other) -> bool:
        return (
            self.func == other.func
            and self.params == other.params
            and self.output_schema.values == other.output_schema.values
        )

    @staticmethod
    def _func_signature(
        chain: str, func: Callable
    ) -> tuple[dict[str, type], Sequence[type], bool]:
        sign = inspect.signature(func)

        input_map = {prm.name: prm.annotation for prm in sign.parameters.values()}
        is_iterator = False

        anno = sign.return_annotation
        if anno == inspect.Signature.empty:
            output_types: list[type] = []
        else:
            orig = get_origin(anno)
            if inspect.isclass(orig) and issubclass(orig, Iterator):
                args = get_args(anno)
                if len(args) > 1 and not (
                    issubclass(orig, Generator) and len(args) == 3
                ):
                    raise UdfSignatureError(
                        chain,
                        f"function '{func}' should return iterator with a single"
                        f" value while '{args}' are specified",
                    )
                is_iterator = True
                anno = args[0]
                orig = get_origin(anno)

            if orig and orig == tuple:
                output_types = tuple(get_args(anno))  # type: ignore[assignment]
            else:
                output_types = [anno]

        return input_map, output_types, is_iterator
