import functools
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Annotated, Any, Generic, TypeAlias, TypeVar, cast

from pydantic import BaseModel

from dagster import _check as check
from dagster._core.definitions.asset_checks.asset_check_spec import AssetCheckSpec
from dagster._core.definitions.assets.definition.asset_spec import AssetSpec
from dagster.components.resolved.context import ResolutionContext
from dagster.components.resolved.core_models import (
    AssetAttributesModel,
    resolve_asset_spec_update_kwargs_to_mapping,
)
from dagster.components.resolved.model import Resolver

if TYPE_CHECKING:
    from dagster.components.component.component import Component


TRANSLATOR_MERGE_ATTRIBUTES = {"metadata", "tags"}


@dataclass
class TranslatorResolvingInfo:
    asset_attributes: str | BaseModel
    resolution_context: ResolutionContext
    model_key: str = "asset_attributes"

    def _resolve_asset_attributes(self, context: Mapping[str, Any]) -> AssetSpec | BaseModel:
        """Resolves the user-specified asset attributes into an AssetAttributesModel, or an AssetSpec
        if the UDF returns one.
        """
        if not isinstance(self.asset_attributes, str):
            return self.asset_attributes

        resolved_asset_attributes = (
            self.resolution_context.at_path(self.model_key)
            .with_scope(**context)
            .resolve_value(self.asset_attributes)
        )

        if isinstance(resolved_asset_attributes, AssetSpec):
            return resolved_asset_attributes
        elif isinstance(resolved_asset_attributes, AssetAttributesModel):
            return resolved_asset_attributes
        elif isinstance(resolved_asset_attributes, dict):
            return AssetAttributesModel(**(resolved_asset_attributes))
        else:
            check.failed(
                f"Unexpected return value for asset_attributes UDF: {type(resolved_asset_attributes)}"
            )

    def get_asset_spec(self, base_spec: AssetSpec, context: Mapping[str, Any]) -> AssetSpec:
        """Returns an AssetSpec that combines the base spec with attributes resolved using the provided context.

        Usage:

        ```python
        class WrappedDagsterXTranslator(DagsterXTranslator):
            def __init__(self, *, base_translator, resolving_info: TranslatorResolvingInfo):
                self.base_translator = base_translator
                self.resolving_info = resolving_info

            def get_asset_spec(self, base_spec: AssetSpec, x_params: Any) -> AssetSpec:
                return self.resolving_info.get_asset_spec(
                    base_spec, {"x_params": x_params}
                )

        ```
        """
        resolved_asset_attributes = self._resolve_asset_attributes(context)
        if isinstance(resolved_asset_attributes, AssetSpec):
            return resolved_asset_attributes

        resolved_attributes = dict(
            resolve_asset_spec_update_kwargs_to_mapping(
                model=resolved_asset_attributes,
                context=self.resolution_context.at_path(self.model_key).with_scope(**context),
            )
        )
        if "code_version" in resolved_attributes:
            resolved_attributes = {
                **resolved_attributes,
                "code_version": str(resolved_attributes["code_version"]),
            }

        if "key_prefix" in resolved_attributes:
            prefix = resolved_attributes.pop("key_prefix")
            if "key" in resolved_attributes:
                key = resolved_attributes.pop("key")
            else:
                key = base_spec.key
            key = key.with_prefix(prefix)
            resolved_attributes["key"] = key

        return base_spec.replace_attributes(
            **{k: v for k, v in resolved_attributes.items() if k not in TRANSLATOR_MERGE_ATTRIBUTES}
        ).merge_attributes(
            **{k: v for k, v in resolved_attributes.items() if k in TRANSLATOR_MERGE_ATTRIBUTES}
        )


T = TypeVar("T")

TranslationFn: TypeAlias = Callable[[AssetSpec, T], AssetSpec]


def _build_translation_fn(
    template_vars_for_translation_fn: Callable[[T], Mapping[str, Any]],
) -> Callable[[ResolutionContext, T], TranslationFn[T]]:
    def resolve_translation(
        context: ResolutionContext,
        model,
    ) -> TranslationFn[T]:
        info = TranslatorResolvingInfo(
            asset_attributes=model,
            resolution_context=context,
            model_key="translation",
        )
        return lambda base_asset_spec, data: info.get_asset_spec(
            base_asset_spec,
            {
                **(template_vars_for_translation_fn(data)),
                "spec": base_asset_spec,
            },
        )

    return resolve_translation


class TranslationFnResolver(Resolver, Generic[T]):
    """Resolver which builds a TranslationFn from input AssetAttributesModel, injecting
    the provided template vars into the resolution process.

    This is useful to expose nested fields within the input data as template vars.

    Example:
    .. code-block:: python

        class MyApiData(BaseModel):
            name: str

        class MyApiComponent(Component, Resolvable):
            translation: Annotated[
                TranslationFn[MyApiData],
                TranslationFnResolver[MyApiData](
                    template_vars_for_translation_fn=lambda data: {"name": data.name}
                ),
            ]

    .. code-block:: yaml

        type: my_module.MyApiComponent

        attributes:
            translation:
                key: {{ name }}

    """

    def __init__(
        self,
        template_vars_for_translation_fn: Callable[[T], Mapping[str, Any]],
        model_field_type: type = AssetAttributesModel,
    ):
        super().__init__(
            _build_translation_fn(
                template_vars_for_translation_fn=template_vars_for_translation_fn
            ),
            model_field_type=str | model_field_type,
            inject_before_resolve=False,
        )


# Common case of a TranslationFn that takes a single underlying entity and passes it through
# as a template var called "data".
ResolvedTranslationFn: TypeAlias = (
    Annotated[TranslationFn[T], TranslationFnResolver[T](lambda data: {"data": data})] | None
)

T_Component = TypeVar("T_Component", bound="Component", covariant=True)
T_Translator = TypeVar("T_Translator")


class ComponentTranslator(Generic[T_Component]):
    """To support python versions < 3.10, we need to use a Protocol to tell the type system that
    these generated classes have a component property.
    """

    def __init__(self, component: T_Component, *args, **kwargs):
        self._component = component
        super().__init__(*args, **kwargs)

    @property
    def component(self) -> T_Component:
        return self._component


def create_component_translator_cls(
    base_component_cls: type[T_Component], base_translator_cls: type[T_Translator]
) -> type[T_Translator]:
    class _GeneratedComponentTranslator(base_translator_cls, ComponentTranslator[T_Component]):  # type: ignore
        def _shim_method(self, method_name: str) -> Callable:
            component_base_method = getattr(base_component_cls, method_name)
            component_instance_method = getattr(self._component.__class__, method_name)
            if component_base_method is not component_instance_method:
                # if the user has overridden the method, we invoke the instance method
                return functools.partial(component_instance_method, self._component)
            else:
                # we never invoke the component_base_method directly. instead, if the user
                # has not overridden the method, we invoke the original translator class method
                return getattr(super(), method_name)

        def get_asset_spec(self, *args, **kwargs) -> AssetSpec:
            return self._shim_method("get_asset_spec")(*args, **kwargs)

        def get_asset_check_spec(self, *args, **kwargs) -> AssetCheckSpec | None:
            return self._shim_method("get_asset_check_spec")(*args, **kwargs)

    return cast("type[T_Translator]", _GeneratedComponentTranslator)
