"""``MelleaSession``: the primary entry point for running generative programs.

``MelleaSession`` wraps a ``Backend`` and a ``Context`` and exposes high-level methods
(``act``, ``instruct``, ``sample``) that drive the generate-validate-repair loop. It
also manages a global context variable (accessible via ``get_session()``) so that
nested components can reach the current session without explicit threading. Use
``start_session(...)`` as a context manager to create and automatically clean up a
session.
"""

from __future__ import annotations

import contextvars
import inspect
from copy import copy
from typing import Any, Literal, overload

from PIL import Image as PILImage

from ..backends.model_ids import IBM_GRANITE_4_MICRO_3B, ModelIdentifier
from ..core import (
    Backend,
    BaseModelSubclass,
    CBlock,
    Component,
    Context,
    FancyLogger,
    GenerateLog,
    ImageBlock,
    ModelOutputThunk,
    Requirement,
    S,
    SamplingResult,
    SamplingStrategy,
    ValidationResult,
)
from ..helpers import _run_async_in_thread
from ..plugins.manager import has_plugins, invoke_hook
from ..plugins.types import HookType
from ..stdlib import functional as mfuncs
from ..telemetry import set_span_attribute, trace_application
from .components import Message
from .context import SimpleContext
from .sampling import RejectionSamplingStrategy

# Global context variable for the context session
_context_session: contextvars.ContextVar[MelleaSession | None] = contextvars.ContextVar(
    "context_session", default=None
)


def get_session() -> MelleaSession:
    """Get the current session from context.

    Returns:
        The currently active ``MelleaSession``.

    Raises:
        RuntimeError: If no session is currently active.
    """
    session = _context_session.get()
    if session is None:
        raise RuntimeError(
            "No active session found. Use 'with start_session(...):' to create one."
        )
    return session


def backend_name_to_class(name: str) -> Any:
    """Resolves backend names to Backend classes.

    Args:
        name: Short backend name, e.g. ``"ollama"``, ``"hf"``, ``"openai"``,
            ``"watsonx"``, or ``"litellm"``.

    Returns:
        The corresponding ``Backend`` class, or ``None`` if the name is unrecognised.

    Raises:
        ImportError: If the requested backend has optional dependencies that are
            not installed (e.g. ``mellea[hf]``, ``mellea[watsonx]``, or
            ``mellea[litellm]``).
    """
    if name == "ollama":
        from ..backends.ollama import OllamaModelBackend

        return OllamaModelBackend
    elif name == "hf" or name == "huggingface":
        try:
            from mellea.backends.huggingface import LocalHFBackend

            return LocalHFBackend
        except ImportError as e:
            raise ImportError(
                "The 'hf' backend requires extra dependencies. "
                "Please install them with: pip install 'mellea[hf]'"
            ) from e
    elif name == "openai":
        from ..backends.openai import OpenAIBackend

        return OpenAIBackend
    elif name == "watsonx":
        try:
            from ..backends.watsonx import WatsonxAIBackend

            return WatsonxAIBackend
        except ImportError as e:
            raise ImportError(
                "The 'watsonx' backend requires extra dependencies. "
                "Please install them with: pip install 'mellea[watsonx]'"
            ) from e
    elif name == "litellm":
        try:
            from ..backends.litellm import LiteLLMBackend

            return LiteLLMBackend
        except ImportError as e:
            raise ImportError(
                "The 'litellm' backend requires extra dependencies. "
                "Please install them with: pip install 'mellea[litellm]'"
            ) from e
    else:
        return None


def start_session(
    backend_name: Literal["ollama", "hf", "openai", "watsonx", "litellm"] = "ollama",
    model_id: str | ModelIdentifier = IBM_GRANITE_4_MICRO_3B,
    ctx: Context | None = None,
    *,
    model_options: dict | None = None,
    plugins: list[Any] | None = None,
    **backend_kwargs,
) -> MelleaSession:
    """Start a new Mellea session. Can be used as a context manager or called directly.

    This function creates and configures a new Mellea session with the specified backend
    and model. When used as a context manager (with `with` statement), it automatically
    sets the session as the current active session for use with convenience functions
    like `instruct()`, `chat()`, `query()`, and `transform()`. When called directly,
    it returns a session object that can be used directly.

    Args:
        backend_name: The backend to use. Options are:
            - "ollama": Use Ollama backend for local models
            - "hf" or "huggingface": Use HuggingFace transformers backend
            - "openai": Use OpenAI API backend
            - "watsonx": Use IBM WatsonX backend
            - "litellm": Use the LiteLLM backend
        model_id: Model identifier or name. Can be a `ModelIdentifier` from
            mellea.backends.model_ids or a string model name.
        ctx: Context manager for conversation history. Defaults to SimpleContext().
            Use ChatContext() for chat-style conversations.
        model_options: Additional model configuration options that will be passed
            to the backend (e.g., temperature, max_tokens, etc.).
        plugins: Optional list of plugins scoped to this session. Accepts
            ``@hook``-decorated functions, ``@plugin``-decorated class instances,
            ``MelleaPlugin`` instances, or ``PluginSet`` instances.
        **backend_kwargs: Additional keyword arguments passed to the backend constructor.

    Returns:
        MelleaSession: A session object that can be used as a context manager
        or called directly with session methods.

    Raises:
        Exception: If ``backend_name`` is not one of the recognised backend
            identifiers.
        ImportError: If the requested backend requires optional dependencies
            that are not installed.

    Examples:
        ```python
        # Basic usage with default settings
        with start_session() as session:
            response = session.instruct("Explain quantum computing")

        # Using OpenAI with custom model options
        with start_session("openai", "gpt-4", model_options={"temperature": 0.7}):
            response = session.chat("Write a poem")

        # Using HuggingFace with ChatContext for conversations
        from mellea.stdlib.base import ChatContext
        with start_session("hf", "microsoft/DialoGPT-medium", ctx=ChatContext()):
            session.chat("Hello!")
            session.chat("How are you?")  # Remembers previous message

        # Direct usage.
        session = start_session()
        response = session.instruct("Explain quantum computing")
        session.cleanup()
        ```
    """
    logger = FancyLogger.get_logger()

    # Get model_id string for logging and tracing
    if isinstance(model_id, ModelIdentifier):
        backend_to_attr = {
            "ollama": "ollama_name",
            "hf": "hf_model_name",
            "huggingface": "hf_model_name",
            "openai": "openai_name",
            "watsonx": "watsonx_name",
            "litellm": "hf_model_name",
        }
        attr = backend_to_attr.get(backend_name, "hf_model_name")
        model_id_str = (
            getattr(model_id, attr, None) or model_id.hf_model_name or str(model_id)
        )
    else:
        model_id_str = str(model_id)

    with trace_application(
        "start_session",
        backend=backend_name,
        model_id=model_id_str,
        context_type=ctx.__class__.__name__ if ctx else "SimpleContext",
    ):
        # --- session_pre_init hook ---
        if has_plugins(HookType.SESSION_PRE_INIT):
            from ..plugins.hooks.session import SessionPreInitPayload

            pre_payload = SessionPreInitPayload(
                backend_name=backend_name,
                model_id=model_id_str,
                model_options=model_options,
                context_type=ctx.__class__.__name__ if ctx else "SimpleContext",
            )
            _, pre_payload = _run_async_in_thread(
                invoke_hook(HookType.SESSION_PRE_INIT, pre_payload)
            )
            # Apply writable field modifications
            model_id_str = pre_payload.model_id
            model_options = pre_payload.model_options

        backend_class = backend_name_to_class(backend_name)
        if backend_class is None:
            raise Exception(
                f"Backend name {backend_name} unknown. Please see the docstring for `mellea.stdlib.session.start_session` for a list of options."
            )
        assert backend_class is not None
        backend = backend_class(model_id, model_options=model_options, **backend_kwargs)

        if ctx is None:
            ctx = SimpleContext()

        logger.info(
            f"Starting Mellea session: backend={backend_name}, model={model_id_str}, "
            f"context={ctx.__class__.__name__}"
            + (f", model_options={model_options}" if model_options else "")
        )

        session = MelleaSession(backend, ctx)

        # Register session-scoped plugins
        if plugins:
            from ..plugins.registry import register as register_plugins

            register_plugins(plugins, session_id=session.id)

        # --- session_post_init hook ---
        if has_plugins(HookType.SESSION_POST_INIT):
            from ..plugins.hooks.session import SessionPostInitPayload

            post_payload = SessionPostInitPayload(
                session_id=session.id, model_id=model_id_str, context=session.ctx
            )
            _run_async_in_thread(
                invoke_hook(HookType.SESSION_POST_INIT, post_payload, backend=backend)
            )

        return session


class MelleaSession:
    """Mellea sessions are a THIN wrapper around `m` convenience functions with NO special semantics.

    Using a Mellea session is not required, but it does represent the "happy path" of Mellea programming. Some nice things about ussing a `MelleaSession`:
    1. In most cases you want to keep a Context together with the Backend from which it came.
    2. You can directly run an instruction or a send a chat, instead of first creating the `Instruction` or `Chat` object and then later calling backend.generate on the object.
    3. The context is "threaded-through" for you, which allows you to issue a sequence of commands instead of first calling backend.generate on something and then appending it to your context.

    These are all relatively simple code hygiene and state management benefits, but they add up over time.
    If you are doing complicating programming (e.g., non-trivial inference scaling) then you might be better off forgoing `MelleaSession`s and managing your Context and Backend directly.

    Note: we put the `instruct`, `validate`, and other convenience functions here instead of in `Context` or `Backend` to avoid import resolution issues.

    Args:
        backend (Backend): The backend to use for all model inference in this
            session.
        ctx (Context | None): The conversation context. Defaults to a new
            ``SimpleContext`` if ``None``.

    Attributes:
        ctx (Context): The active conversation context; never ``None`` (defaults
            to a fresh ``SimpleContext`` when ``None`` is passed). Updated after
            every call that produces model output.
        id (str): Unique session UUID assigned at construction.
    """

    ctx: Context

    def __init__(self, backend: Backend, ctx: Context | None = None):
        """Initialize MelleaSession with a backend and optional conversation context."""
        import uuid

        self.id = str(uuid.uuid4())
        self.backend = backend
        self.ctx: Context = ctx if ctx is not None else SimpleContext()
        self._session_logger = FancyLogger.get_logger()
        self._context_token = None
        self._session_span = None

    def __enter__(self):
        """Enter context manager and set this session as the current global session."""
        # Start a session span that will last for the entire context manager lifetime
        self._session_span = trace_application(
            "session_context",
            backend=self.backend.__class__.__name__,
            context_type=self.ctx.__class__.__name__,
        ).__enter__()
        self._context_token = _context_session.set(self)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit context manager and cleanup session."""
        self.cleanup()
        if self._context_token is not None:
            _context_session.reset(self._context_token)
            self._context_token = None
        if self._session_span is not None:
            self._session_span.__exit__(exc_type, exc_val, exc_tb)
            self._session_span = None

    def __copy__(self):
        """Use self.clone. Copies the current session but keeps references to the backend and context."""
        new = MelleaSession(backend=self.backend, ctx=self.ctx)
        new._session_logger = self._session_logger
        # Explicitly don't copy over the _context_token.

        return new

    def clone(self):
        """Useful for running multiple generation requests while keeping the context at a given point in time.

        Returns:
            a copy of the current session. Keeps the context, backend, and session logger.

        Examples:
            ```python
            >>> from mellea import start_session
            >>> m = start_session()
            >>> m.instruct("What is 2x2?")
            >>>
            >>> m1 = m.clone()
            >>> out = m1.instruct("Multiply that by 2")
            >>> print(out)
            ... 8
            >>>
            >>> m2 = m.clone()
            >>> out = m2.instruct("Multiply that by 3")
            >>> print(out)
            ... 12
            ```
        """
        return copy(self)

    def reset(self):
        """Reset the context state to a fresh, empty context of the same type.

        Fires the ``SESSION_RESET`` plugin hook if any plugins are registered, then
        replaces ``self.ctx`` with the result of ``ctx.reset_to_new()``, discarding
        all accumulated conversation history.
        """
        if has_plugins(HookType.SESSION_RESET):
            from ..plugins.hooks.session import SessionResetPayload

            payload = SessionResetPayload(previous_context=self.ctx)
            _run_async_in_thread(
                invoke_hook(HookType.SESSION_RESET, payload, backend=self.backend)
            )
        self.ctx = self.ctx.reset_to_new()

    def cleanup(self) -> None:
        """Clean up session resources and deregister session-scoped plugins."""
        if has_plugins(HookType.SESSION_CLEANUP):
            from ..plugins.hooks.session import SessionCleanupPayload

            payload = SessionCleanupPayload(
                context=self.ctx, interaction_count=len(self.ctx.as_list())
            )
            _run_async_in_thread(
                invoke_hook(HookType.SESSION_CLEANUP, payload, backend=self.backend)
            )

        # Deregister session-scoped plugins — must run whenever plugins are
        # enabled, regardless of whether any plugin subscribes to SESSION_CLEANUP.
        if has_plugins():
            from ..plugins.manager import deregister_session_plugins

            deregister_session_plugins(self.id)

    @overload
    def act(
        self,
        action: Component[S],
        *,
        requirements: list[Requirement] | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: Literal[False] = False,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk[S]: ...

    @overload
    def act(
        self,
        action: Component[S],
        *,
        requirements: list[Requirement] | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: Literal[True],
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> SamplingResult[S]: ...

    def act(
        self,
        action: Component[S],
        *,
        requirements: list[Requirement] | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: bool = False,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk[S] | SamplingResult:
        """Runs a generic action, and adds both the action and the result to the context.

        Args:
            action: the Component from which to generate.
            requirements: used as additional requirements when a sampling strategy is provided
            strategy: a SamplingStrategy that describes the strategy for validating and repairing/retrying for the instruct-validate-repair pattern. None means that no particular sampling strategy is used.
            return_sampling_results: attach the (successful and failed) sampling attempts to the results.
            format: if set, the BaseModel to use for constrained decoding.
            model_options: additional model options, which will upsert into the model/backend's defaults.
            tool_calls: if true, tool calling is enabled.

        Returns:
            A ModelOutputThunk if `return_sampling_results` is `False`, else returns a `SamplingResult`.
        """
        r = mfuncs.act(
            action,
            context=self.ctx,
            backend=self.backend,
            requirements=requirements,
            strategy=strategy,
            return_sampling_results=return_sampling_results,
            format=format,
            model_options=model_options,
            tool_calls=tool_calls,
        )  # type: ignore

        if isinstance(r, SamplingResult):
            self.ctx = r.result_ctx
            return r
        else:
            result, context = r
            self.ctx = context
            return result

    @overload
    def instruct(
        self,
        description: str,
        *,
        images: list[ImageBlock] | list[PILImage.Image] | None = None,
        requirements: list[Requirement | str] | None = None,
        icl_examples: list[str | CBlock] | None = None,
        grounding_context: dict[str, str | CBlock | Component] | None = None,
        user_variables: dict[str, str] | None = None,
        prefix: str | CBlock | None = None,
        output_prefix: str | CBlock | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: Literal[False] = False,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk[str]: ...

    @overload
    def instruct(
        self,
        description: str,
        *,
        images: list[ImageBlock] | list[PILImage.Image] | None = None,
        requirements: list[Requirement | str] | None = None,
        icl_examples: list[str | CBlock] | None = None,
        grounding_context: dict[str, str | CBlock | Component] | None = None,
        user_variables: dict[str, str] | None = None,
        prefix: str | CBlock | None = None,
        output_prefix: str | CBlock | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: Literal[True],
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> SamplingResult[str]: ...

    def instruct(
        self,
        description: str,
        *,
        images: list[ImageBlock] | list[PILImage.Image] | None = None,
        requirements: list[Requirement | str] | None = None,
        icl_examples: list[str | CBlock] | None = None,
        grounding_context: dict[str, str | CBlock | Component] | None = None,
        user_variables: dict[str, str] | None = None,
        prefix: str | CBlock | None = None,
        output_prefix: str | CBlock | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: bool = False,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk[str] | SamplingResult:
        """Generates from an instruction.

        Args:
            description: The description of the instruction.
            requirements: A list of requirements that the instruction can be validated against.
            icl_examples: A list of in-context-learning examples that the instruction can be validated against.
            grounding_context: A list of grounding contexts that the instruction can use. They can bind as variables using a (key: str, value: str | ContentBlock) tuple.
            user_variables: A dict of user-defined variables used to fill in Jinja placeholders in other parameters. This requires that all other provided parameters are provided as strings.
            prefix: A prefix string or ContentBlock to use when generating the instruction.
            output_prefix: A string or ContentBlock that defines a prefix for the output generation. Usually you do not need this.
            strategy: A SamplingStrategy that describes the strategy for validating and repairing/retrying for the instruct-validate-repair pattern. None means that no particular sampling strategy is used.
            return_sampling_results: attach the (successful and failed) sampling attempts to the results.
            format: If set, the BaseModel to use for constrained decoding.
            model_options: Additional model options, which will upsert into the model/backend's defaults.
            tool_calls: If true, tool calling is enabled.
            images: A list of images to be used in the instruction or None if none.

        Returns:
            A ``ModelOutputThunk`` if ``return_sampling_results`` is ``False``,
            else a ``SamplingResult``.
        """
        r = mfuncs.instruct(
            description,
            context=self.ctx,
            backend=self.backend,
            images=images,
            requirements=requirements,
            icl_examples=icl_examples,
            grounding_context=grounding_context,
            user_variables=user_variables,
            prefix=prefix,
            output_prefix=output_prefix,
            strategy=strategy,
            return_sampling_results=return_sampling_results,  # type: ignore
            format=format,
            model_options=model_options,
            tool_calls=tool_calls,
        )

        if isinstance(r, SamplingResult):
            self.ctx = r.result_ctx
            return r
        else:
            # It's a tuple[ModelOutputThunk, Context].
            result, context = r
            self.ctx = context
            return result

    def chat(
        self,
        content: str,
        role: Message.Role = "user",
        *,
        images: list[ImageBlock] | list[PILImage.Image] | None = None,
        user_variables: dict[str, str] | None = None,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> Message:
        """Sends a simple chat message and returns the response. Adds both messages to the Context.

        Args:
            content: The message text to send.
            role: The role for the outgoing message (default ``"user"``).
            images: Optional list of images to include in the message.
            user_variables: Optional Jinja variable substitutions applied to ``content``.
            format: Optional Pydantic model for constrained decoding of the response.
            model_options: Additional model options to merge with backend defaults.
            tool_calls: If true, tool calling is enabled.

        Returns:
            The assistant ``Message`` response.
        """
        result, context = mfuncs.chat(
            content=content,
            context=self.ctx,
            backend=self.backend,
            role=role,
            images=images,
            user_variables=user_variables,
            format=format,
            model_options=model_options,
            tool_calls=tool_calls,
        )

        self.ctx = context
        return result

    def validate(
        self,
        reqs: Requirement | list[Requirement],
        *,
        output: CBlock | None = None,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        generate_logs: list[GenerateLog] | None = None,
        input: CBlock | None = None,
    ) -> list[ValidationResult]:
        """Validates a set of requirements over the output (if provided) or the current context (if the output is not provided).

        Args:
            reqs: A single ``Requirement`` or a list of them to validate.
            output: Optional model output ``CBlock`` to validate against instead of the context.
            format: Optional Pydantic model for constrained decoding.
            model_options: Additional model options to merge with backend defaults.
            generate_logs: Optional list to append generation logs to.
            input: Optional input ``CBlock`` to include alongside ``output`` when validating.

        Returns:
            List of ``ValidationResult`` objects, one per requirement.
        """
        return mfuncs.validate(
            reqs=reqs,
            context=self.ctx,
            backend=self.backend,
            output=output,
            format=format,
            model_options=model_options,
            generate_logs=generate_logs,
            input=input,
        )

    def query(
        self,
        obj: Any,
        query: str,
        *,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk:
        """Query method for retrieving information from an object.

        Args:
            obj : The object to be queried. It should be an instance of MObject or can be converted to one if necessary.
            query:  The string representing the query to be executed against the object.
            format:  format for output parsing.
            model_options: Model options to pass to the backend.
            tool_calls: If true, the model may make tool calls. Defaults to False.

        Returns:
            ModelOutputThunk: The result of the query as processed by the backend.
        """
        result, context = mfuncs.query(
            obj=obj,
            query=query,
            context=self.ctx,
            backend=self.backend,
            format=format,
            model_options=model_options,
            tool_calls=tool_calls,
        )
        self.ctx = context
        return result

    def transform(
        self,
        obj: Any,
        transformation: str,
        *,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
    ) -> ModelOutputThunk | Any:
        """Transform method for creating a new object with the transformation applied.

        Args:
            obj : The object to be queried. It should be an instance of MObject or can be converted to one if necessary.
            transformation:  The string representing the query to be executed against the object.
            format: format for output parsing; usually not needed with transform.
            model_options: Model options to pass to the backend.

        Returns:
            ModelOutputThunk|Any: The result of the transformation as processed by the backend. If no tools were called,
            the return type will be always be ModelOutputThunk. If a tool was called, the return type will be the return type
            of the function called, usually the type of the object passed in.
        """
        result, context = mfuncs.transform(
            obj=obj,
            transformation=transformation,
            context=self.ctx,
            backend=self.backend,
            format=format,
            model_options=model_options,
        )
        self.ctx = context
        return result

    @overload
    async def aact(
        self,
        action: Component[S],
        *,
        requirements: list[Requirement] | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: Literal[False] = False,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk[S]: ...

    @overload
    async def aact(
        self,
        action: Component[S],
        *,
        requirements: list[Requirement] | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: Literal[True],
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> SamplingResult[S]: ...

    async def aact(
        self,
        action: Component[S],
        *,
        requirements: list[Requirement] | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: bool = False,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk[S] | SamplingResult:
        """Runs a generic action, and adds both the action and the result to the context.

        Args:
            action: the Component from which to generate.
            requirements: used as additional requirements when a sampling strategy is provided
            strategy: a SamplingStrategy that describes the strategy for validating and repairing/retrying for the instruct-validate-repair pattern. None means that no particular sampling strategy is used.
            return_sampling_results: attach the (successful and failed) sampling attempts to the results.
            format: if set, the BaseModel to use for constrained decoding.
            model_options: additional model options, which will upsert into the model/backend's defaults.
            tool_calls: if true, tool calling is enabled.

        Returns:
            A ModelOutputThunk if `return_sampling_results` is `False`, else returns a `SamplingResult`.
        """
        r = await mfuncs.aact(
            action,
            context=self.ctx,
            backend=self.backend,
            requirements=requirements,
            strategy=strategy,
            return_sampling_results=return_sampling_results,
            format=format,
            model_options=model_options,
            tool_calls=tool_calls,
        )  # type: ignore

        if isinstance(r, SamplingResult):
            self.ctx = r.result_ctx
            return r
        else:
            result, context = r
            self.ctx = context
            return result

    @overload
    async def ainstruct(
        self,
        description: str,
        *,
        images: list[ImageBlock] | list[PILImage.Image] | None = None,
        requirements: list[Requirement | str] | None = None,
        icl_examples: list[str | CBlock] | None = None,
        grounding_context: dict[str, str | CBlock | Component] | None = None,
        user_variables: dict[str, str] | None = None,
        prefix: str | CBlock | None = None,
        output_prefix: str | CBlock | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: Literal[False] = False,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk[str]: ...

    @overload
    async def ainstruct(
        self,
        description: str,
        *,
        images: list[ImageBlock] | list[PILImage.Image] | None = None,
        requirements: list[Requirement | str] | None = None,
        icl_examples: list[str | CBlock] | None = None,
        grounding_context: dict[str, str | CBlock | Component] | None = None,
        user_variables: dict[str, str] | None = None,
        prefix: str | CBlock | None = None,
        output_prefix: str | CBlock | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: Literal[True],
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> SamplingResult[str]: ...

    async def ainstruct(
        self,
        description: str,
        *,
        images: list[ImageBlock] | list[PILImage.Image] | None = None,
        requirements: list[Requirement | str] | None = None,
        icl_examples: list[str | CBlock] | None = None,
        grounding_context: dict[str, str | CBlock | Component] | None = None,
        user_variables: dict[str, str] | None = None,
        prefix: str | CBlock | None = None,
        output_prefix: str | CBlock | None = None,
        strategy: SamplingStrategy | None = RejectionSamplingStrategy(loop_budget=2),
        return_sampling_results: bool = False,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk[str] | SamplingResult[str]:
        """Generates from an instruction.

        Args:
            description: The description of the instruction.
            requirements: A list of requirements that the instruction can be validated against.
            icl_examples: A list of in-context-learning examples that the instruction can be validated against.
            grounding_context: A list of grounding contexts that the instruction can use. They can bind as variables using a (key: str, value: str | ContentBlock) tuple.
            user_variables: A dict of user-defined variables used to fill in Jinja placeholders in other parameters. This requires that all other provided parameters are provided as strings.
            prefix: A prefix string or ContentBlock to use when generating the instruction.
            output_prefix: A string or ContentBlock that defines a prefix for the output generation. Usually you do not need this.
            strategy: A SamplingStrategy that describes the strategy for validating and repairing/retrying for the instruct-validate-repair pattern. None means that no particular sampling strategy is used.
            return_sampling_results: attach the (successful and failed) sampling attempts to the results.
            format: If set, the BaseModel to use for constrained decoding.
            model_options: Additional model options, which will upsert into the model/backend's defaults.
            tool_calls: If true, tool calling is enabled.
            images: A list of images to be used in the instruction or None if none.

        Returns:
            A ``ModelOutputThunk`` if ``return_sampling_results`` is ``False``,
            else a ``SamplingResult``.
        """
        r = await mfuncs.ainstruct(
            description,
            context=self.ctx,
            backend=self.backend,
            images=images,
            requirements=requirements,
            icl_examples=icl_examples,
            grounding_context=grounding_context,
            user_variables=user_variables,
            prefix=prefix,
            output_prefix=output_prefix,
            strategy=strategy,
            return_sampling_results=return_sampling_results,  # type: ignore
            format=format,
            model_options=model_options,
            tool_calls=tool_calls,
        )

        if isinstance(r, SamplingResult):
            self.ctx = r.result_ctx
            return r
        else:
            # It's a tuple[ModelOutputThunk, Context].
            result, context = r
            self.ctx = context
            return result

    async def achat(
        self,
        content: str,
        role: Message.Role = "user",
        *,
        images: list[ImageBlock] | list[PILImage.Image] | None = None,
        user_variables: dict[str, str] | None = None,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> Message:
        """Sends a simple chat message and returns the response. Adds both messages to the Context.

        Args:
            content: The message text to send.
            role: The role for the outgoing message (default ``"user"``).
            images: Optional list of images to include in the message.
            user_variables: Optional Jinja variable substitutions applied to ``content``.
            format: Optional Pydantic model for constrained decoding of the response.
            model_options: Additional model options to merge with backend defaults.
            tool_calls: If true, tool calling is enabled.

        Returns:
            The assistant ``Message`` response.
        """
        result, context = await mfuncs.achat(
            content=content,
            context=self.ctx,
            backend=self.backend,
            role=role,
            images=images,
            user_variables=user_variables,
            format=format,
            model_options=model_options,
            tool_calls=tool_calls,
        )

        self.ctx = context
        return result

    async def avalidate(
        self,
        reqs: Requirement | list[Requirement],
        *,
        output: CBlock | None = None,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        generate_logs: list[GenerateLog] | None = None,
        input: CBlock | None = None,
    ) -> list[ValidationResult]:
        """Validates a set of requirements over the output (if provided) or the current context (if the output is not provided).

        Args:
            reqs: A single ``Requirement`` or a list of them to validate.
            output: Optional model output ``CBlock`` to validate against instead of the context.
            format: Optional Pydantic model for constrained decoding.
            model_options: Additional model options to merge with backend defaults.
            generate_logs: Optional list to append generation logs to.
            input: Optional input ``CBlock`` to include alongside ``output`` when validating.

        Returns:
            List of ``ValidationResult`` objects, one per requirement.
        """
        return await mfuncs.avalidate(
            reqs=reqs,
            context=self.ctx,
            backend=self.backend,
            output=output,
            format=format,
            model_options=model_options,
            generate_logs=generate_logs,
            input=input,
        )

    async def aquery(
        self,
        obj: Any,
        query: str,
        *,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> ModelOutputThunk:
        """Query method for retrieving information from an object.

        Args:
            obj : The object to be queried. It should be an instance of MObject or can be converted to one if necessary.
            query:  The string representing the query to be executed against the object.
            format:  format for output parsing.
            model_options: Model options to pass to the backend.
            tool_calls: If true, the model may make tool calls. Defaults to False.

        Returns:
            ModelOutputThunk: The result of the query as processed by the backend.
        """
        result, context = await mfuncs.aquery(
            obj=obj,
            query=query,
            context=self.ctx,
            backend=self.backend,
            format=format,
            model_options=model_options,
            tool_calls=tool_calls,
        )
        self.ctx = context
        return result

    async def atransform(
        self,
        obj: Any,
        transformation: str,
        *,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
    ) -> ModelOutputThunk | Any:
        """Transform method for creating a new object with the transformation applied.

        Args:
            obj: The object to be queried. It should be an instance of MObject or can be converted to one if necessary.
            transformation:  The string representing the query to be executed against the object.
            format: format for output parsing; usually not needed with transform.
            model_options: Model options to pass to the backend.

        Returns:
            ModelOutputThunk|Any: The result of the transformation as processed by the backend. If no tools were called,
            the return type will be always be ModelOutputThunk. If a tool was called, the return type will be the return type
            of the function called, usually the type of the object passed in.
        """
        result, context = await mfuncs.atransform(
            obj=obj,
            transformation=transformation,
            context=self.ctx,
            backend=self.backend,
            format=format,
            model_options=model_options,
        )
        self.ctx = context
        return result

    @classmethod
    def powerup(cls, powerup_cls: type):
        """Appends methods in a class object `powerup_cls` to MelleaSession.

        Iterates over all functions defined on ``powerup_cls`` and attaches each
        one as a method on the ``MelleaSession`` class, effectively extending
        the session with domain-specific helpers at runtime.

        Args:
            powerup_cls (type): A class whose functions should be added to
                ``MelleaSession`` as instance methods.
        """
        for name, fn in inspect.getmembers(powerup_cls, predicate=inspect.isfunction):
            setattr(cls, name, fn)

    # ###############################
    #  Convenience functions
    # ###############################
    def last_prompt(self) -> str | list[dict] | None:
        """Returns the last prompt that has been called from the session context.

        Returns:
            A string if the last prompt was a raw call to the model OR a list of messages (as role-msg-dicts). Is None if none could be found.
        """
        op = self.ctx.last_output()
        if op is None:
            return None
        log = op._generate_log
        if isinstance(log, GenerateLog):
            return log.prompt
        elif isinstance(log, list):
            last_el = log[-1]
            if isinstance(last_el, GenerateLog):
                return last_el.prompt
        return None
