# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Compilation specifications, for example, dynamic shape inputs."""

import inspect
import typing

if typing.TYPE_CHECKING:
    from .core import Module as nn_module_class

ArgSpecType = typing.Union["Int", "Tensor"]
MethodSpecType = typing.Union["MethodSpec", dict[str, ArgSpecType]]
ModuleSpecType = typing.Union["ModuleSpec", dict[str, MethodSpecType]]
SpecAny = typing.Union["Object", "Int", "Tensor", "Tuple"]


class Int:  # pylint: disable=too-few-public-methods
    """An integer input"""

    def __init__(self) -> None:
        pass

    def __repr__(self) -> str:
        return "int"


class Tensor:  # pylint: disable=too-few-public-methods
    """A tensor input with static ndim and dtype, but can have symbolic shapes."""

    shape: list[int | str]
    dtype: str

    def __init__(self, shape: typing.Sequence[int | str], dtype: str) -> None:
        self.shape = list(shape)
        self.dtype = dtype

    def __repr__(self) -> str:
        shape = ", ".join(str(i) for i in self.shape)
        return f"Tensor([{shape}], '{self.dtype}')"


class Object:  # pylint: disable=too-few-public-methods
    """An non-tensor opaque frontend object."""

    object_type: type

    def __init__(self, object_type: type) -> None:
        self.object_type = object_type

    def __repr__(self) -> str:
        return "object"


class Tuple:  # pylint: disable=too-few-public-methods
    """A tuple input or a list input"""

    name: str
    elements: list[SpecAny] | tuple[SpecAny, ...]

    def __init__(
        self,
        name: str,
        elements: list[SpecAny] | tuple[SpecAny, ...],
    ) -> None:
        assert isinstance(elements, tuple | list), f"Unsupported container type: {type(elements)}"
        self.name = name
        self.elements = elements

    def __repr__(self) -> str:
        return self.elements.__repr__()


class MethodSpec:
    """A spec for a compiled method"""

    method: typing.Callable
    arg_names: list[str]
    arg_specs: list[ArgSpecType]
    param_mode: str  # "plain", "packed", "none"
    effect_mode: str  # "plain", "packed", "none"

    def __init__(  # pylint: disable=too-many-arguments
        self,
        method: typing.Callable,
        arg_names: list[str],
        arg_specs: list[ArgSpecType],
        param_mode: str,
        effect_mode: str,
    ):
        if param_mode not in ["plain", "packed", "none"]:
            raise ValueError(f"Invalid param_mode: {param_mode}")
        if effect_mode not in ["plain", "packed", "none"]:
            raise ValueError(f"Invalid effect_mode: {effect_mode}")
        self.method = method
        self.arg_names = arg_names
        self.arg_specs = arg_specs
        self.param_mode = param_mode
        self.effect_mode = effect_mode

    def _repr(self, name: str) -> str:
        args = ", ".join(
            f"{name}: {spec}"
            for name, spec in zip(
                self.arg_names,
                self.arg_specs,
            )
        )
        return f"{name}({args})"

    def __repr__(self) -> str:
        return self._repr(name="MethodSpec")

    @staticmethod
    def from_raw(spec: MethodSpecType, method: typing.Callable) -> "MethodSpec":
        """Create MethodSpec from raw python dictionaries.

        Examples
        --------
        .. code-block:: python

            MethodSpec.from_raw(
                spec={
                    "inputs": spec.Tensor([batch_size, "seq_len"], "int32"),
                    "total_seq_len": "int",
                },
                method=module.prefill,
            )
        """
        if isinstance(spec, MethodSpec):
            return spec
        config: dict[str, typing.Any] = spec.pop("$", {})  # type: ignore[assignment]
        param_mode = config.get("param_mode", "plain")
        effect_mode = config.get("effect_mode", "plain")
        method_signature = inspect.signature(method)
        arg_names = list(method_signature.parameters.keys())
        arg_specs = []

        def _convert_arg_spec(arg_spec, arg_name):
            if arg_spec is Int or arg_spec is int:
                return Int()
            if isinstance(arg_spec, str) and arg_spec == "int":
                return Int()
            if isinstance(arg_spec, Int | Tensor | Object):
                return arg_spec
            if isinstance(arg_spec, tuple | list | Tuple):
                return Tuple(
                    arg_name,
                    elements=type(arg_spec)(
                        [
                            _convert_arg_spec(arg_spec_i, f"{arg_name}_{i}")
                            for i, arg_spec_i in enumerate(arg_spec)
                        ]
                    ),
                )
            raise TypeError(f"Invalid spec for argument {arg_name}: {arg_spec}")

        for arg_name in arg_names:
            if arg_name in spec:
                arg_spec = spec[arg_name]
                arg_spec = _convert_arg_spec(arg_spec, arg_name)
                arg_specs.append(arg_spec)
        return MethodSpec(
            method,
            arg_names,
            arg_specs,
            param_mode=param_mode,
            effect_mode=effect_mode,
        )

    @staticmethod
    def from_torch(args: list[typing.Any], method: typing.Callable) -> "MethodSpec":
        """Converts a list of torch tensors to MethodSpec."""
        from .torch import (  # pylint: disable=import-outside-toplevel
            _method_spec_from_torch,
        )

        return _method_spec_from_torch(args, method)


class ModuleSpec:
    """A spec for a compiled nn.Module"""

    module: "nn_module_class"
    method_names: list[str]
    method_specs: list[MethodSpec]

    def __init__(
        self,
        module: "nn_module_class",
        method_names: list[str],
        method_specs: list[MethodSpec],
    ) -> None:
        self.module = module
        self.method_names = method_names
        self.method_specs = method_specs

    @staticmethod
    def from_raw(spec: ModuleSpecType, module: "nn_module_class") -> "ModuleSpec":
        """Create ModuleSpec from raw python dictionaries.

        Examples
        --------
        .. code-block:: python

            ModuleSpec.from_raw(
                spec={
                    "prefill": {
                        "inputs": spec.Tensor([batch_size, "seq_len"], "int32"),
                        "total_seq_len": int,
                    },
                    "decode": {
                        "inputs": spec.Tensor([batch_size, 1], "int32"),
                        "total_seq_len": int,
                    },
                    "softmax_with_temperature": {
                        "logits": spec.Tensor([1, 1, config.vocab_size], "float32"),
                        "temperature": spec.Tensor([], "float32"),
                    },
                },
                module=module,
            )
        """
        if isinstance(spec, ModuleSpec):
            return spec
        method_names = list(spec.keys())
        method_specs: list[MethodSpec] = []
        for method_name in method_names:
            method_spec = spec[method_name]
            if isinstance(method_spec, MethodSpec):
                pass
            else:
                method_spec = MethodSpec.from_raw(method_spec, getattr(module, method_name))
            method_specs.append(method_spec)
        return ModuleSpec(module, method_names, method_specs)

    def __repr__(self) -> str:
        return "ModuleSpec:\n" + "\n".join(
            "  " + spec._repr(name)  # pylint: disable=protected-access
            for name, spec in zip(
                self.method_names,
                self.method_specs,
            )
        )
