# Copyright 2025 Google LLC
#
# Licensed 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.
#
# SPDX-License-Identifier: Apache-2.0

"""Utilities for resolving tools and partials.

A well-defined object resolver is a callable that takes a name and returns
either an object or an awaitable (such as a future or coroutine) that returns an
object.

## Key Operations

| Function              | Description                                                                |
|-----------------------|----------------------------------------------------------------------------|
| `resolve`             | Core async function to resolve a named object using a given resolver.      |
|                       | Handles both sync/async resolvers and sync functions returning awaitables. |
| `resolve_tool`        | Helper async function specifically for resolving tool names.               |
| `resolve_partial`     | Helper async function specifically for resolving partial names.            |
| `resolve_json_schema` | Helper async function specifically for resolving JSON schemas.             |

The `resolve` function handles both sync and async resolvers. If the resolver is
sync, it is run in a thread pool to avoid blocking the event loop. If the
resolver is async, it is awaited directly.

The `resolve_*` functions are convenience wrappers around `resolve` that handle
the specific types of resolvers for tools, partials, and schemas.
"""

import inspect
from collections.abc import Awaitable, Callable
from typing import Any

from anyio.to_thread import run_sync

from dotpromptz.errors import ResolverFailedError
from dotpromptz.typing import (
    JsonSchema,
    PartialResolver,
    SchemaResolver,
    ToolDefinition,
    ToolResolver,
)

# Resolver callable type.
ResolverCallable = Callable[[str], Awaitable[Any] | Any]


async def resolve[ResolverT: ResolverCallable, DefinitionT](
    name: str,
    kind: str,
    resolver: ResolverT | None,
) -> DefinitionT:
    """Resolves a single object using the provided resolver.

    If the resolver is synchronous, it is run in a thread pool to avoid
    blocking the event loop.

    Args:
        name: The name of the object to resolve.
        kind: The kind of object to resolve.
        resolver: The object resolver callable.

    Returns:
        The resolved object.

    Raises:
        LookupError: If the resolver returns None for the object.
        ResolverFailedError: For exceptions raised by the resolver.
        TypeError: If the resolver is not callable or returns an invalid type.
        ValueError: If the resolver is not defined.
    """
    obj: DefinitionT | None = None

    if resolver is None:
        raise ValueError(f'{kind} resolver is not defined')

    if not callable(resolver):
        raise TypeError(f"{kind} resolver for '{name}' is not callable")

    try:
        # We need to check if the callable itself is async first, or if it returns an awaitable.
        #
        # ```ascii
        #         +---------------------------+
        #         | collections.abc.Awaitable |  (ABC: Can be `await`ed)
        #         +---------------------------+
        #            /|\             /|\
        #             |               | (Implements/Is-a)
        #     (Inherits/Is-a)         |
        #             |               |
        #     +----------------+  +---------------------+  <-+ (Returned by call to...)
        #     | asyncio.Future |  | types.CoroutineType |    |
        #     | (Low-level     |  | (Awaitable Object)  |    | +--------------------------+
        #     |  Awaitable)    |  | (from `async def`)  |    +-| collections.abc.Callable |
        #     +----------------+  +---------------------+      | (e.g. `async def` func)  |
        #             /|\                 ^                    +--------------------------+
        #              |                  | (Often wrapped by)
        #     (Inherits/Is-a)             |
        #              |                  |
        #     +----------------+          |
        #     |  asyncio.Task  |----------+
        #     | (Runs Coroutine|
        #     |  is a Future)  |
        #     +----------------+
        #                 ^
        #                 | (Managed by)
        #                 |
        #     +-------------------+
        #     | asyncio.TaskGroup |
        #     | (Context Manager) |
        #     +-------------------+
        # ```
        if inspect.iscoroutinefunction(resolver):
            # If resolver is async, call it directly and await.
            obj = await resolver(name)
        else:
            # If resolver is sync, run it in a thread pool and check the return
            # type after calling, as we don't know it yet. It might still return
            # an awaitable (e.g. sync function returning `asyncio.Future`) but
            # calling it sync first is necessary to check.
            result_or_awaitable = await run_sync(resolver, name)
            if inspect.isawaitable(result_or_awaitable):
                obj = await result_or_awaitable
            else:
                obj = result_or_awaitable

    except Exception as e:
        # Catch errors from both await and sync execution in thread.
        raise ResolverFailedError(name, kind, str(e)) from e

    # Raise when the resolver indicates the object doesn't exist.
    if obj is None:
        raise LookupError(f"{kind} resolver for '{name}' returned None")

    return obj


async def resolve_tool(name: str, resolver: ToolResolver | None) -> ToolDefinition:
    """Resolve a tool using the provided resolver.

    Args:
        name: The name of the tool to resolve.
        resolver: The tool resolver callable (sync or async).

    Returns:
        The resolved tool definition.

    Raises:
        LookupError: If the resolver returns None for the tool.
        ResolverFailedError: For exceptions raised by the resolver.
        TypeError: If the resolver is not callable or returns an invalid type.
        ValueError: If the resolver is not defined.
    """
    return await resolve(name, 'tool', resolver)


async def resolve_partial(name: str, resolver: PartialResolver | None) -> str:
    """Resolve a partial using the provided resolver.

    Args:
        name: The name of the partial to resolve.
        resolver: The partial resolver callable.

    Returns:
        The resolved partial.

    Raises:
        LookupError: If the resolver returns None for the partial.
        ResolverFailedError: For exceptions raised by the resolver.
        TypeError: If the resolver is not callable or returns an invalid type.
        ValueError: If the resolver is not defined.
    """
    return await resolve(name, 'partial', resolver)


async def resolve_json_schema(name: str, resolver: SchemaResolver | None) -> JsonSchema:
    """Resolve a JSON schema using the provided resolver.

    Args:
        name: The name of the JSON schema to resolve.
        resolver: The JSON schema resolver callable.

    Returns:
        The resolved JSON schema.

    Raises:
        LookupError: If the resolver returns None for the schema.
        ResolverFailedError: For exceptions raised by the resolver.
        TypeError: If the resolver is not callable or returns an invalid type.
    """
    return await resolve(name, 'schema', resolver)
