from typing import Any, Callable, List, Mapping, Protocol, Tuple, TypeVar, Union

from djc_core.safe_eval.eval import _exec_func_with_error_handling
from djc_core.rust import template_parser


TContext = TypeVar("TContext")

TagUnion = Union[
    template_parser.GenericTag, template_parser.ForLoopTag, template_parser.EndTag
]


class TemplateStringResolver(Protocol[TContext]):
    def __call__(
        self,
        __context: TContext,
        __source: str,
        __token: Tuple[int, int],
        __filters: Mapping[str, Callable],
        __tags: Mapping[str, Callable],
        expr: str,
    ) -> Any: ...


class VariableResolver(Protocol[TContext]):
    def __call__(
        self,
        __context: TContext,
        __source: str,
        __token: Tuple[int, int],
        __filters: Mapping[str, Callable],
        __tags: Mapping[str, Callable],
        var: str,
    ) -> Any: ...


class TranslationResolver(Protocol[TContext]):
    def __call__(
        self,
        __context: TContext,
        __source: str,
        __token: Tuple[int, int],
        __filters: Mapping[str, Callable],
        __tags: Mapping[str, Callable],
        text: str,
    ) -> Any: ...


class ExpressionResolver(Protocol[TContext]):
    def __call__(
        self,
        __context: TContext,
        __source: str,
        __token: Tuple[int, int],
        __filters: Mapping[str, Callable],
        __tags: Mapping[str, Callable],
        code: str,
    ) -> Any: ...


class FilterResolver(Protocol[TContext]):
    def __call__(
        self,
        __context: TContext,
        __source: str,
        __token: Tuple[int, int],
        __filters: Mapping[str, Callable],
        __tags: Mapping[str, Callable],
        name: str,
        value: Any,
        arg: Any,
    ) -> Any: ...


# The top-level `...var` spread may be handled either as a list of args or a dict of kwargs
# depending on the value of `var`.
# So we check at runtime whether `var` is a mapping by checking for `.keys()` or if it's an iterable.
# We use this helper function to handle this.
# This helper is called from within the compiled function generated by `compile_tag_attrs()`.
def _handle_spread(
    value: Any,
    raw_token_str: str,
    args: List[Any],
    kwargs: List[Tuple[str, Any]],
    kwarg_seen: bool,
) -> bool:
    if hasattr(value, "keys"):
        kwargs.extend(value.items())
        return True
    else:
        if kwarg_seen:
            raise SyntaxError("positional argument follows keyword argument")
        try:
            args.extend(value)
        except TypeError:
            raise TypeError(
                f"Value of '...{raw_token_str}' must be a mapping or an iterable, "
                f"not {type(value).__name__}."
            )
        return False


def compile_tag(
    tag_or_attrs: Union[template_parser.GenericTag, List[template_parser.TagAttr]],
    source: str,
    filters: Mapping[str, Callable],
    tags: Mapping[str, Callable],
    *,
    expr: ExpressionResolver[TContext],
    template_string: TemplateStringResolver[TContext],
    variable: VariableResolver[TContext],
    translation: TranslationResolver[TContext],
    filter: FilterResolver[TContext],
) -> Callable[[TContext], Tuple[List[Any], List[Tuple[str, Any]]]]:
    """
    Compile the template tag AST into a Python function.

    The AST can be a `GenericTag` or a list of `TagAttr` objects
    (typically generated by `parse_tag`).

    This function requires implementations of the following resolvers:
    - `expr`: Python expressions (e.g., `(1 + 2)`)
    - `template_string`: Template strings (e.g., `'{% lorem %}'`)
    - `translation`: Translation strings (e.g., `_("text")`)
    - `variable`: Variables (e.g., `my_var`)
    - `filter`: Filters (e.g., `value|upper:arg`)

    Each resolver receives:
    - `context`: The context dictionary passed to the compiled function
    - `source`: The full source string of the tag
    - `token`: A tuple of (start_index, end_index) indicating the token's position in source
    - `filters`: The filters mapping passed to `compile_tag`
    - `tags`: The tags mapping passed to `compile_tag`
    - Additional resolver-specific arguments (see resolver definitions)

    Returns a function that accepts the context object and returns a tuple of
    (arguments, keyword arguments).

    Example:

    ```python
    from djc_core.template_parser import parse_tag, compile_tag

    tag_content = "{% component val1 [1, 2, 3] a=b data={'key': 'value'} %}"
    tag_ast = parse_tag(tag_content)

    # Resolvers - each receives (context, source, token, filters, tags, ...specific_args)
    variable = lambda ctx, src, token, filters, tags, var: ctx.get(var)
    template_string = lambda ctx, src, token, filters, tags, expr: f"TEMPLATE:{expr}"
    translation = lambda ctx, src, token, filters, tags, text: f"TRANS:{text}"
    filter_resolver = lambda ctx, src, token, filters, tags, name, value, arg: (
        f"{value}|{name}:{arg}" if arg is not None else f"{value}|{name}"
    )
    expr = lambda ctx, src, token, filters, tags, code: eval(code, {}, ctx)

    compiled_func = compile_tag(
        tag_ast,
        source=tag_content,
        filters={},
        tags={},
        expr=expr,
        template_string=template_string,
        variable=variable,
        translation=translation,
        filter=filter_resolver,
    )

    args, kwargs = compiled_func({"val1": "foo", "b": "bar"})

    print(args)  # ['foo', [1, 2, 3]]
    print(kwargs)  # [('a', 'bar'), ('data', {'key': 'value'})]
    ```
    """
    if isinstance(tag_or_attrs, template_parser.GenericTag):
        attributes = tag_or_attrs.attrs
    else:
        attributes = tag_or_attrs

    global_scope = {
        "_handle_spread": _handle_spread,
        "expr": expr,
        "filter": filter,
        "template_string": template_string,
        "translation": translation,
        "variable": variable,
        # Pass through other metadata that we may want to access from within the resolvers.
        # We pass these into the resolvers as positional args.
        # That way we don't have to re-define the functions on each evaluation.
        "source": source,
        "filters": filters,
        "tags": tags,
    }

    func_string = template_parser.compile_tag_attrs(attributes)
    return _exec_func_with_error_handling(
        func_string, "compiled_func", source, "expression", global_scope
    )


def compile_value(
    value: template_parser.TagValue,
    source: str,
    filters: Mapping[str, Callable],
    tags: Mapping[str, Callable],
    *,
    expr: ExpressionResolver[TContext],
    template_string: TemplateStringResolver[TContext],
    variable: VariableResolver[TContext],
    translation: TranslationResolver[TContext],
    filter: FilterResolver[TContext],
) -> Callable[[TContext], Any]:
    """
    Compile a single `TagValue` into a Python function.

    This function compiles a single value (e.g., a variable, list, dict, or expression)
    into a function that can be evaluated with a context.

    This function requires implementations of the following resolvers:
    - `expr`: Python expressions (e.g., `(1 + 2)`)
    - `template_string`: Template strings (e.g., `'{% lorem %}'`)
    - `translation`: Translation strings (e.g., `_("text")`)
    - `variable`: Variables (e.g., `my_var`)
    - `filter`: Filters (e.g., `value|upper:arg`)

    Each resolver receives:
    - `context`: The context dictionary passed to the compiled function
    - `source`: The full source string
    - `token`: A tuple of (start_index, end_index) indicating the token's position in source
    - `filters`: The filters mapping passed to `compile_value`
    - `tags`: The tags mapping passed to `compile_value`
    - Additional resolver-specific arguments (see resolver definitions)

    Returns a function that accepts the context object and returns the resolved value.

    Example:

    ```python
    from djc_core.template_parser import parse_tag, compile_value

    tag_content = "{% component my_var|upper %}"
    tag = parse_tag(tag_content)
    value = tag.attrs[0].value  # Get the TagValue for "my_var|upper"

    # Resolvers - each receives (context, source, token, filters, tags, ...specific_args)
    variable = lambda ctx, src, token, filters, tags, var: ctx.get(var)
    template_string = lambda ctx, src, token, filters, tags, expr: f"TEMPLATE:{expr}"
    translation = lambda ctx, src, token, filters, tags, text: f"TRANS:{text}"
    filter_resolver = lambda ctx, src, token, filters, tags, name, value, arg: (
        f"{value}|{name}:{arg}" if arg is not None else f"{value}|{name}"
    )
    expr = lambda ctx, src, token, filters, tags, code: eval(code, {}, ctx)

    compiled_func = compile_value(
        value,
        source=tag_content,
        filters={"upper": lambda value, arg: value.upper()},
        tags={},
        expr=expr,
        template_string=template_string,
        variable=variable,
        translation=translation,
        filter=filter_resolver,
    )

    result = compiled_func({"my_var": "hello"})
    print(result)  # "HELLO"
    ```
    """
    global_scope = {
        "expr": expr,
        "filter": filter,
        "template_string": template_string,
        "translation": translation,
        "variable": variable,
        # Pass through other metadata that we may want to access from within the resolvers.
        # We pass these into the resolvers as positional args.
        # That way we don't have to re-define the functions on each evaluation.
        "source": source,
        "filters": filters,
        "tags": tags,
    }

    # Get the compiled expression string from Rust
    compiled_expression = template_parser.compile_value(value)

    # Wrap it in a function that takes context and returns the value
    func_string = f"def compiled_func(context):\n    return {compiled_expression}"

    return _exec_func_with_error_handling(
        func_string, "compiled_func", source, "expression", global_scope
    )
