from collections.abc import Callable, Collection, Coroutine, Iterable
from dataclasses import KW_ONLY, dataclass
from functools import lru_cache
from typing import Any, Protocol

from fastapi import Request, Response
from fastapi.templating import Jinja2Templates

from .component_selectors import ComponentHeader as _ComponentHeader
from .core_decorators import hx, page
from .typing import ComponentSelector, MaybeAsyncFunc, P, RenderFunction, RequestComponentSelector


class JinjaContextFactory(Protocol):
    """
    Protocol definition for methods that convert a FastAPI route's result and route context
    (i.e. the route's arguments) into a Jinja context (`dict`).
    """

    def __call__(self, *, route_result: Any, route_context: dict[str, Any]) -> dict[str, Any]:
        """
        Arguments:
            route_result: The result of the route.
            route_context: Every keyword argument the route received.

        Returns:
            The Jinja context dictionary.

        Raises:
            ValueError: If converting the arguments to a Jinja context fails.
        """
        ...


class JinjaPath(str):
    """
    String subclass that can be used to mark a template path as "absolute".

    In this context "absolute" means the template path should be exempt from any prefixing behavior
    during template name resolution.

    Note: calling any of the "mutation" methods (e.g. `.lower()`) of an instance will
    result in a plain `str` object.
    """

    ...


class JinjaContext:
    """
    Core `JinjaContextFactory` implementations.
    """

    @classmethod
    def unpack_object(cls, obj: Any) -> dict[str, Any]:
        """
        Utility function that unpacks an object into a `dict`.

        Supports `dict` and `Collection` instances, plus anything with `__dict__` or `__slots__`
        attributes, for example Pydantic models, dataclasses, or "standard" class instances.

        Conversion rules:

        - `dict`: returned as is.
        - `Collection`: returned as `{"items": obj}`, available in templates as `items`.
        - Objects with `__dict__` or `__slots__`: known keys are taken from `__dict__` or `__slots__`
          and the context will be created as `{key: getattr(route_result, key) for key in keys}`,
          omitting property names starting with an underscore. For Pydantic models, computed
          fields will also be included.
        - `None` is converted into an empty context.

        Raises:
            ValueError: If the given object can not be handled by any of the conversion rules.
        """
        if isinstance(obj, dict):
            return obj

        # Covers lists, tuples, sets, etc..
        if isinstance(obj, Collection):
            return {"items": obj}

        object_keys: Iterable[str] | None = None

        # __dict__ should take priority if an object has both this and __slots__.
        if hasattr(obj, "__dict__"):
            # Covers Pydantic models and standard classes.
            object_keys = obj.__dict__.keys()
            cls = type(obj)
            if hasattr(cls, "model_computed_fields"):  # Pydantic computed fields support.
                object_keys = [
                    *(() if object_keys is None else object_keys),
                    *cls.model_computed_fields,
                ]
        elif hasattr(obj, "__slots__"):
            # Covers classes with with __slots__.
            object_keys = obj.__slots__

        if object_keys is not None:
            return {key: getattr(obj, key) for key in object_keys if not key.startswith("_")}

        if obj is None:
            # Convert no response to empty context.
            return {}

        raise ValueError("Result conversion failed, unknown result type.")

    @classmethod
    def unpack_result(cls, *, route_result: Any, route_context: dict[str, Any]) -> dict[str, Any]:
        """
        Jinja context factory that tries to reasonably convert non-`dict` route results
        to valid Jinja contexts (the `route_context` argument is ignored).

        Supports everything that `JinjaContext.unpack_object()` does and follows the same
        conversion rules.

        Raises:
            ValueError: If `route_result` can not be handled by any of the conversion rules.
        """
        return cls.unpack_object(route_result)

    @classmethod
    def unpack_result_with_route_context(
        cls,
        *,
        route_result: Any,
        route_context: dict[str, Any],
    ) -> dict[str, Any]:
        """
        Jinja context factory that tries to reasonably convert non-`dict` route results
        to valid Jinja contexts, also including every key-value pair from `route_context`.

        Supports everything that `JinjaContext.unpack_object()` does and follows the same
        conversion rules.

        Raises:
            ValueError: If `JinjaContext.unpack_result()` raises an error or if there's
                a key conflict between `route_result` and `route_context`.
        """
        result = cls.unpack_result(route_result=route_result, route_context=route_context)
        if len(set(result.keys()) & set(route_context.keys())) > 0:
            raise ValueError("Overlapping keys in route result and route context.")

        # route_context is the keyword args of the route collected into a dict. Update and
        # return this dict rather than result, as the result might be the same object that
        # was returned by the route and someone may have a reference to it.
        route_context.update(result)
        return route_context

    @classmethod
    def use_converters(
        cls,
        convert_route_result: Callable[[Any], dict[str, Any]] | None,
        convert_route_context: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
    ) -> JinjaContextFactory:
        """
        Creates a `JinjaContextFactory` that uses the provided functions to convert
        `route_result` and `route_context` to a Jinja context.

        The returned `JinjaContextFactory` raises a `ValueError` if the overlapping keys are found.

        Arguments:
            convert_route_result: Function that takes `route_result` and converts it into a `dict`.
                See `JinjaContextFactory` for `route_result` details.
            convert_route_context: Function that takes `route_context` and converts it into a `dict`.
                See `JinjaContextFactory` for `route_context` details.

        Returns:
            The created `JinjaContextFactory`.
        """

        def make_jinja_context(*, route_result: Any, route_context: dict[str, Any]) -> dict[str, Any]:
            rr = {} if convert_route_result is None else convert_route_result(route_result)
            rc = {} if convert_route_context is None else convert_route_context(route_context)
            if len(set(rr.keys()) & set(rc.keys())) > 0:
                raise ValueError("Overlapping keys in route result and route context.")

            rr.update(rc)
            return rr

        return make_jinja_context

    @classmethod
    @lru_cache
    def wrap_as(cls, result_key: str, context_key: str | None = None) -> JinjaContextFactory:
        """
        Creates a `JinjaContextFactory` that wraps the route's result and optionally the route
        context under user-specified keys.

        `result_key` and `context_key` must be different.

        Arguments:
            result_key: The key by which the `route_result` should be accessible in templates.
                See `JinjaContextFactory` for `route_result` details.
            context_key: The key by which the `route_context` should be accessible in templates.
                If `None` (the default), then the `route_context` will not be accessible.
                See `JinjaContextFactory` for `route_context` details.

        Returns:
            The created `JinjaContextFactory`.

        Raises:
            ValueError: If `result_key` and `context_key` are equal.
        """

        if result_key == context_key:
            raise ValueError("The two keys must be different, merging is not supported.")

        def wrap(*, route_result: Any, route_context: dict[str, Any]) -> dict[str, Any]:
            result = {result_key: route_result}
            if context_key is not None:
                result[context_key] = route_context

            return result

        return wrap


class ComponentHeader(_ComponentHeader[str]):
    """
    `RequestComponentSelector` for Jinja templates that selects the rendered template
    based on a request header.
    """

    ...


@dataclass(frozen=True, slots=True)
class Jinja:
    """Jinja2 renderer utility with FastAPI route decorators."""

    templates: Jinja2Templates
    """The Jinja2 templates of the application."""

    make_context: JinjaContextFactory = JinjaContext.unpack_result
    """
    Function that will be used by default to convert a route's return value into
    a Jinja rendering context. The default value is `JinjaContext.unpack_result`.
    """

    _: KW_ONLY

    no_data: bool = False
    """
    If set, `hx()` routes will only accept HTMX requests.

    Note that if this property is `True`, then the `hx()` decorator's `no_data` argument
    will have no effect.
    """

    def hx(
        self,
        template: ComponentSelector[str],
        *,
        error_template: ComponentSelector[str] | None = None,
        make_context: JinjaContextFactory | None = None,
        no_data: bool = False,
        prefix: str | None = None,
    ) -> Callable[[MaybeAsyncFunc[P, Any]], Callable[P, Coroutine[None, None, Any | Response]]]:
        """
        Decorator for rendering a route's result if the request was an HTMX one.

        Arguments:
            template: The Jinja2 template selector to use.
            error_template: The Jinja2 template selector to use for route error rendering.
            make_context: Route-specific override for the `make_context` property.
            no_data: If set, the route will only accept HTMX requests.
            prefix: Optional template name prefix.

        Returns:
            The rendered HTML for HTMX requests, otherwise the route's unchanged return value.
        """
        if make_context is None:
            # No route-specific override.
            make_context = self.make_context

        render_func = self._make_render_function(template, make_context=make_context, prefix=prefix)
        error_render_func = (
            None
            if error_template is None
            else self._make_render_function(
                error_template, make_context=make_context, prefix=prefix, error_renderer=True
            )
        )
        # TODO: check why mypy doesn't resolve the generics correctly.
        return hx(render_func, render_error=error_render_func, no_data=self.no_data or no_data)  # type: ignore[return-value]

    def page(
        self,
        template: ComponentSelector[str],
        *,
        error_template: ComponentSelector[str] | None = None,
        make_context: JinjaContextFactory | None = None,
        prefix: str | None = None,
    ) -> Callable[[MaybeAsyncFunc[P, Any]], Callable[P, Coroutine[None, None, Any | Response]]]:
        """
        Decorator for rendering a route's result.

        This decorator triggers HTML rendering regardless of whether the request was HTMX or not.

        Arguments:
            template: The Jinja2 template selector to use.
            error_template: The Jinja2 template selector to use for route error rendering.
            make_context: Route-specific override for the `make_context` property.
            prefix: Optional template name prefix.
        """
        if make_context is None:
            # No route-specific override.
            make_context = self.make_context

        render_func = self._make_render_function(template, make_context=make_context, prefix=prefix)
        error_render_func = (
            None
            if error_template is None
            else self._make_render_function(
                error_template,
                make_context=make_context,
                prefix=prefix,
                error_renderer=True,
            )
        )
        # TODO: check why mypy doesn't resolve the generics correctly.
        return page(render_func, render_error=error_render_func)  # type: ignore[return-value]

    def _make_render_function(
        self,
        template: ComponentSelector[str],
        *,
        make_context: JinjaContextFactory,
        prefix: str | None,
        error_renderer: bool = False,
    ) -> RenderFunction[Any]:
        """
        Creates a `RenderFunction` with the given configuration.

        Arguments:
            template: The template the renderer function should use.
            make_context: The Jinja rendering context factory to use.
            prefix: Optional template name prefix.
            error_renderer: Whether this is an error renderer creation.
        """

        def render(result: Any, *, context: dict[str, Any], request: Request) -> str:
            template_name = self._resolve_template_name(
                template,
                error=result if error_renderer else None,
                prefix=prefix,
                request=request,
            )
            result = self.templates.TemplateResponse(
                name=template_name,
                context=make_context(route_result=result, route_context=context),
                request=request,
            )
            return bytes(result.body).decode(result.charset)

        return render

    def _resolve_template_name(
        self,
        template: ComponentSelector[str],
        *,
        error: Exception | None = None,
        prefix: str | None,
        request: Request,
    ) -> str:
        """
        Resolves the template selector into a full template name.

        Arguments:
            template: The template selector.
            error: The error raised by the route.
            prefix: Optional template name prefix.
            request: The current request.

        Returns:
            The resolved, full template name.

        Raises:
            ValueError: If template resolution failed.
        """
        if isinstance(template, RequestComponentSelector):
            try:
                result = template.get_component(request, error)
            except KeyError as e:
                raise ValueError("Failed to resolve template name from request.") from e
        elif isinstance(template, str):
            result = template
        else:
            raise ValueError("Unknown template selector.")

        prefix = None if isinstance(result, JinjaPath) else prefix
        result = result.lstrip("/")
        return f"{prefix}/{result}" if prefix else result
