from __future__ import annotations

from abc import ABC, abstractmethod
from inspect import Signature
from typing import AsyncIterator, Callable, get_args, get_origin, overload

from humps import camelize

from .proto.qlik.sse import BundledRows, DataType, Dual, FunctionDefinition, FunctionType, Parameter, Row
from .qlik_sse_types import (
    SingleDataType,
    WrappableAggregationFunction,
    WrappableFunction,
    WrappableScalarFunction,
    WrappableTensorFunction,
)
from .util import dual_getters, map_data_type, py_to_data_type, to_dual


class QlikSSEFunction(ABC):
    @abstractmethod
    def get_definition(self, function_id: int) -> FunctionDefinition:
        raise NotImplementedError

    @abstractmethod
    def handle_bundle(self, bundle: BundledRows) -> BundledRows:
        raise NotImplementedError

    async def __call__(self, bundles: AsyncIterator[BundledRows]) -> AsyncIterator[BundledRows]:
        async for bundle in bundles:
            yield self.handle_bundle(bundle)


class QlikSSEFunctionWrapper(QlikSSEFunction, ABC):
    _f: WrappableFunction
    function_type: FunctionType
    parameter_types: list[Parameter]
    return_type: DataType

    def __init__(self, f: WrappableFunction) -> None:
        self._f = f
        signature = Signature.from_callable(f)
        assert len(signature.parameters) > 0, "at least one parameter is required"
        for k in ["function_type", "parameter_types", "return_type"]:
            if getattr(self, k, None) is None:
                setattr(self, k, getattr(self, f"_infer_{k}")(signature))

    @classmethod
    def _infer_function_type(cls, signature: Signature) -> FunctionType:
        if get_origin(signature.return_annotation) is list:
            return FunctionType.TENSOR
        else:
            if get_origin(list(signature.parameters.values())[0].annotation) is list:
                return FunctionType.AGGREGATION
            else:
                return FunctionType.SCALAR

    @classmethod
    @abstractmethod
    def _infer_parameter_types(cls, signature: Signature) -> list[Parameter]:
        raise NotImplementedError

    @classmethod
    def _infer_return_type(cls, signature: Signature) -> DataType:
        return py_to_data_type[signature.return_annotation]

    @overload
    def __new__(cls, f: WrappableScalarFunction) -> QlikSSEScalarFunctionWrapper: ...
    @overload
    def __new__(cls, f: WrappableAggregationFunction) -> QlikSSEAggregationFunctionWrapper: ...
    @overload
    def __new__(cls, f: WrappableTensorFunction) -> QlikSSETensorFunctionWrapper: ...
    def __new__(cls, f: WrappableFunction) -> QlikSSEFunctionWrapper:
        function_type = cls._infer_function_type(Signature.from_callable(f))
        if function_type == FunctionType.SCALAR:
            return super().__new__(QlikSSEScalarFunctionWrapper)
        if function_type == FunctionType.AGGREGATION:
            return super().__new__(QlikSSEAggregationFunctionWrapper)
        if function_type == FunctionType.TENSOR:
            return super().__new__(QlikSSETensorFunctionWrapper)
        raise NotImplementedError

    def get_definition(self, function_id: int) -> FunctionDefinition:
        return FunctionDefinition(
            function_id=function_id,
            name=camelize(self._f.__name__).title(),
            function_type=self.function_type,
            return_type=self.return_type,
            params=self.parameter_types,
        )


class QlikSSEScalarFunctionWrapper(QlikSSEFunctionWrapper):
    function_type: FunctionType = FunctionType.SCALAR
    _f: WrappableScalarFunction
    _param_getters: list[Callable[[Dual], SingleDataType]]

    def __init__(self, f: WrappableScalarFunction):
        super().__init__(f)
        self._param_getters = [dual_getters[param.data_type] for param in self.parameter_types]

    @classmethod
    def _infer_parameter_types(cls, signature: Signature) -> list[Parameter]:
        return [
            Parameter(name=parameter.name, data_type=map_data_type(parameter.annotation))
            for parameter in signature.parameters.values()
        ]

    def _parse_params(self, duals: list[Dual]) -> list[SingleDataType]:
        return [getter(dual) for getter, dual in zip(self._param_getters, duals)]

    def handle_bundle(self, bundle: BundledRows) -> BundledRows:
        return BundledRows(
            rows=[Row(duals=[to_dual(self._f(*self._parse_params(row.duals)))]) for row in bundle.rows],
        )


class QlikSSEMultiFunctionWrapper(QlikSSEFunctionWrapper, ABC):
    _f: WrappableAggregationFunction | WrappableTensorFunction
    _param_getters: list[Callable[[Dual], SingleDataType]]

    def __init__(self, f: WrappableAggregationFunction | WrappableTensorFunction):
        super().__init__(f)
        self._param_getters = [dual_getters[param.data_type] for param in self.parameter_types]

    @classmethod
    def _infer_parameter_types(cls, signature: Signature) -> list[Parameter]:
        return [
            Parameter(
                name=parameter.name,
                data_type=map_data_type(get_args(parameter.annotation)[0]),
            )
            for parameter in signature.parameters.values()
        ]

    def _parse_params(self, bundle: BundledRows) -> list[list[SingleDataType]]:
        return [[getter(row.duals[idx]) for row in bundle.rows] for idx, getter in enumerate(self._param_getters)]


class QlikSSEAggregationFunctionWrapper(QlikSSEMultiFunctionWrapper):
    function_type: FunctionType = FunctionType.AGGREGATION
    _f: WrappableAggregationFunction

    def __init__(self, f: WrappableAggregationFunction):
        super().__init__(f)

    def handle_bundle(self, bundle: BundledRows) -> BundledRows:
        return BundledRows(rows=[Row(duals=[to_dual(self._f(*self._parse_params(bundle)))])])


class QlikSSETensorFunctionWrapper(QlikSSEMultiFunctionWrapper):
    function_type: FunctionType = FunctionType.TENSOR
    _f: WrappableTensorFunction

    def __init__(self, f: WrappableTensorFunction):
        super().__init__(f)

    @classmethod
    def _infer_return_type(cls, signature: Signature) -> DataType:
        return_type = get_args(signature.return_annotation)[0]
        assert get_origin(return_type) is list
        return py_to_data_type[get_args(return_type)[0]]

    def handle_bundle(self, bundle: BundledRows) -> BundledRows:
        p = self._parse_params(bundle)
        result = self._f(*p)
        return BundledRows(
            rows=[Row(duals=[to_dual(value) for value in row]) for row in result],
        )
