"""A backend that uses the Huggingface Transformers library.

The purpose of the Hugginface backend is to provide a setting for implementing experimental features. If you want a performance local backend, and do not need experimental features such as Span-based context or ALoras, consider using Ollama backends instead.
"""

from __future__ import annotations

import abc
import asyncio
import dataclasses
import datetime
import functools
import inspect
import json
import threading
from collections.abc import Callable, Coroutine
from copy import deepcopy
from typing import TYPE_CHECKING, Any, cast

import granite_common
import outlines
import outlines_core
import peft
import torch
from transformers import (
    AsyncTextIteratorStreamer,
    AutoModelForCausalLM,
    AutoTokenizer,
    DynamicCache,
    PreTrainedModel,
    PreTrainedTokenizer,
    set_seed,
)
from transformers.generation.utils import GenerateDecoderOnlyOutput

from mellea.backends import BaseModelSubclass
from mellea.backends._utils import to_chat, to_tool_calls
from mellea.backends.adapters.adapter import (
    AdapterMixin,
    AdapterType,
    GraniteCommonAdapter,
    LocalHFAdapter,
    get_adapter_for_intrinsic,
)
from mellea.backends.adapters.catalog import fetch_intrinsic_metadata
from mellea.backends.cache import Cache, SimpleLRUCache
from mellea.backends.formatter import Formatter, FormatterBackend, TemplateFormatter
from mellea.backends.model_ids import ModelIdentifier
from mellea.backends.openai import OpenAIBackend
from mellea.backends.process_reward_models import PRM
from mellea.backends.tools import (
    add_tools_from_context_actions,
    add_tools_from_model_options,
    convert_tools_to_json,
)
from mellea.backends.types import ModelOption
from mellea.helpers.async_helpers import send_to_queue
from mellea.helpers.fancy_logger import FancyLogger
from mellea.stdlib.base import (
    CBlock,
    Component,
    Context,
    GenerateLog,
    GenerateType,
    ModelOutputThunk,
    ModelToolCall,
)
from mellea.stdlib.chat import Message
from mellea.stdlib.intrinsics.intrinsic import Intrinsic
from mellea.stdlib.requirement import ALoraRequirement, LLMaJRequirement, Requirement

assert outlines, "outlines needs to be present to make outlines_core work"

"""A configuration type for the unhappy path: Tokenizer * Model * torch device string

Huggingface backends can initialize themselves from a model string if the transformers `Auto*` classes can be used. Therefore, a TransformersTorchConfig usually isn't required. However, sometimes a model needs special care to instantiate properly, or a custom device type needs to bse used. Instead of trying to do a lot of partial magic, we basically have two modaliites: either the constructor can figure out everything from the model_id, or the user has to provide an entire config.
"""
TransformersTorchConfig = tuple[PreTrainedTokenizer, PreTrainedModel, torch.device]

format: None = None  # typing this variable in order to shadow the global format function and ensure mypy checks for errors


@dataclasses.dataclass
class HFAloraCacheInfo:
    """A dataclass for holding some KV cache and associated information."""

    kv_cache: DynamicCache
    merged_token_ids: Any
    merged_attention: Any
    q_end: int = -1


class LocalHFBackend(FormatterBackend, AdapterMixin):
    """The LocalHFBackend uses Huggingface's transformers library for inference, and uses a Formatter to convert `Component`s into prompts. This backend also supports Activated LoRAs (ALoras)](https://arxiv.org/pdf/2504.12397).

    This backend is designed for running an HF model for small-scale inference locally on your machine.

    This backend is NOT designed for inference scaling on CUDA-enabled hardware.
    """

    def __init__(
        self,
        model_id: str | ModelIdentifier,
        formatter: Formatter | None = None,
        *,
        use_caches: bool = True,
        cache: Cache | None = None,
        custom_config: TransformersTorchConfig | None = None,
        default_to_constraint_checking_alora: bool = True,
        model_options: dict | None = None,
    ):
        """Attempt to load model weights using the model_id by default, or using `custom_config` if provided.

        WARNING: initializing a `LocalHFBackend` will download and load the model on your *local* machine.

        Args:
            model_id (str | ModelIdentifier): Used to load the model *and tokenizer* via transformers Auto* classes, and then moves the model to the best available device (cuda > mps > cpu). If loading the model and/or tokenizer from a string will not work, or if you want to use a different device string, then you can use custom_config.
            formatter (Formatter): A mechanism for turning `stdlib` stuff into strings. Experimental Span-based models should use `mellea.backends.span.*` backends.
            use_caches (bool): If set to False, then caching will not be used even if a Cache is provided.
            cache (Optional[Cache]): The caching strategy to use. If None, `LRUCache(3)` will be used.
            custom_config (Optional[TransformersTorchConfig]): Overrides loading from the `model_id`. If set, then the specified tokenizer/model/device will be used instead of auto-loading from the model_id.
            default_to_constraint_checking_alora: If set to False then aloras will be deactivated. This is primarily for performance benchmarking and debugging.
            model_options (Optional[dict]): Default model options.
        """
        formatter = (
            formatter if formatter is not None else TemplateFormatter(model_id=model_id)
        )

        super().__init__(model_id, formatter, model_options=model_options)

        # A mapping of common options for this backend mapped to their Mellea ModelOptions equivalent.
        # These are usually values that must be extracted before hand or that are common among backend providers
        self.to_mellea_model_opts_map = {
            "system": ModelOption.SYSTEM_PROMPT,
            "max_new_tokens": ModelOption.MAX_NEW_TOKENS,
            "seed": ModelOption.SEED,
            "tools": ModelOption.TOOLS,
            "stream": ModelOption.STREAM,
        }

        # A mapping of Mellea specific ModelOptions to the specific names for this backend.
        # These options should almost always be a subset of those specified in the `to_mellea_model_opts_map`.
        # Usually, values that are intentionally extracted while prepping for the backend generate call
        # will be omitted here so that they will be removed when model_options are processed
        # for the call to the model.
        self.from_mellea_model_opts_map = {ModelOption.MAX_NEW_TOKENS: "max_new_tokens"}

        self.default_to_constraint_checking_alora = default_to_constraint_checking_alora

        # Either use the custom config or load the model from its model_id
        match model_id:
            case str():
                self._hf_model_id = model_id
            case ModelIdentifier():
                assert model_id.hf_model_name is not None, (
                    "model_id is None. This can also happen if the ModelIdentifier has no hf_model_id name set."
                )
                self._hf_model_id = model_id.hf_model_name
        match custom_config:
            case None:
                # Choose a device.
                self._device = torch.device(
                    "cuda"
                    if torch.cuda.is_available()
                    else "mps"
                    if torch.backends.mps.is_available()
                    else "cpu"
                )
                # Get the model and tokenizer.
                self._model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
                    self._hf_model_id
                ).to(self._device)  # type: ignore
                self._tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(
                    self._hf_model_id
                )
            case _:
                self._tokenizer, self._model, self._device = custom_config

        self._use_caches = use_caches
        self._cache = cache if cache is not None else SimpleLRUCache(3)

        # Adapters can be made known to the backend (added) and loaded.
        self._added_adapters: dict[str, LocalHFAdapter] = {}
        self._loaded_adapters: dict[str, LocalHFAdapter] = {}

        self._generation_lock = threading.Lock()
        """Used to force generation requests to be non-concurrent. Necessary for preventing issues with adapters."""

    async def generate_from_context(
        self,
        action: Component | CBlock,
        ctx: Context,
        *,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ):
        """Generate using the huggingface model."""
        # Upsert model options.
        model_opts = self._simplify_and_merge(model_options)

        # Requirements can be automatically rerouted to a requirement adapter.
        if isinstance(action, Requirement):
            # See docs/dev/requirement_aLoRA_rerouting.md
            reroute_to_alora = self.default_to_constraint_checking_alora
            adapter_name = "requirement_check"

            if isinstance(action, ALoraRequirement):
                reroute_to_alora = True
                adapter_name = action.intrinsic_name
                alora_action = action
            else:
                assert action.description is not None, (
                    "must have a description when generating from a requirement"
                )
                alora_action = ALoraRequirement(action.description, adapter_name)

            # Check if a requirement_check (or AloraRequirement specified) adapter
            # exists.
            alora_req_adapter = get_adapter_for_intrinsic(
                adapter_name, [AdapterType.ALORA], self._added_adapters
            )
            if alora_req_adapter is None:
                # Log a warning if using an AloraRequirement but no adapter fit.
                if reroute_to_alora and isinstance(action, ALoraRequirement):
                    FancyLogger.get_logger().warning(
                        f"attempted to use an AloraRequirement but backend {self} doesn't have the specified adapter added {adapter_name}; defaulting to regular generation"
                    )
                reroute_to_alora = False

            if issubclass(type(action), LLMaJRequirement):
                reroute_to_alora = False

            if reroute_to_alora:
                # Keep the alora requirement handling separate for now.
                mot = await self._generate_from_intrinsic(
                    alora_action, ctx, model_options=model_opts
                )
                return mot, ctx.add(alora_action).add(mot)

        elif isinstance(action, Intrinsic):
            mot = await self._generate_from_intrinsic(
                action, ctx, model_options=model_opts
            )
            return mot, ctx.add(action).add(mot)

        mot = await self._generate_from_context_standard(
            action, ctx, _format=format, model_options=model_opts, tool_calls=tool_calls
        )
        return mot, ctx.add(action).add(mot)

    def _generate_with_adapter_lock(
        self, adapter_name: str, generate_func: Callable, *args, **kwargs
    ):
        """Helper function for ensuring exclusive generation when adapters are present. Necessary to prevent generating with incorrect weights."""
        with self._generation_lock:
            if adapter_name != "":
                self.load_adapter(adapter_name)
                self._model.set_adapter(adapter_name)
            else:
                try:
                    # `._model.disable_adapters()` doesn't seem to actually disable them or
                    # remove them from the model's list of `.active_adapters()`.
                    self._model.set_adapter([])
                except ValueError as e:
                    # If no weights have been loaded, the model will raise a ValueError:
                    # `ValueError("No adapter loaded. Please load an adapter first.")`
                    if "No adapter loaded" in str(e):
                        pass
                    else:
                        raise e

            _assert_correct_adapters(adapter_name, self._model)
            out = generate_func(*args, **kwargs)
            _assert_correct_adapters(adapter_name, self._model)
            return out

    async def _generate_from_intrinsic(
        self, action: Intrinsic, ctx: Context, *, model_options: dict[str, Any]
    ) -> ModelOutputThunk:
        if not ctx.is_chat_context:
            raise Exception("Does not yet support non-chat contexts.")

        if len(model_options.items()) > 0:
            FancyLogger.get_logger().info(
                "passing in model options when generating with an adapter; some model options may be overwritten / ignored"
            )

        linearized_ctx = ctx.view_for_generation()
        assert linearized_ctx is not None, (
            "If ctx.is_chat_context, then the context should be linearizable."
        )
        ctx_as_message_list: list[Message] = self.formatter.to_chat_messages(
            linearized_ctx
        )

        conversation: list[dict] = []
        system_prompt = model_options.get(ModelOption.SYSTEM_PROMPT, "")
        if system_prompt != "":
            conversation.append({"role": "system", "content": system_prompt})

        conversation.extend(
            [OpenAIBackend.message_to_openai_message(m) for m in ctx_as_message_list]
        )

        docs = OpenAIBackend.messages_to_docs(ctx_as_message_list)

        seed = model_options.get(ModelOption.SEED, None)
        if seed is not None:
            set_seed(seed)

        if model_options.get(ModelOption.STREAM, None) is not None:
            # Intrinsics don't support streaming because of their post-processing step.
            FancyLogger.get_logger().warning(
                "intrinsics cannot use streaming; removing model option"
            )
            del model_options[ModelOption.STREAM]

        adapter = get_adapter_for_intrinsic(
            action.intrinsic_name, action.adapter_types, self._added_adapters
        )
        if adapter is None:
            raise ValueError(
                f"backend ({self}) has no adapter for processing intrinsic: {action.intrinsic_name}"
            )

        # TODO: Code below this point is mostly specific to RagIntrinsics (and granite_common).
        #       It should be refactored into a specific adapter.transform() function.
        assert isinstance(adapter, GraniteCommonAdapter), (
            "currently Mellea only supports GraniteCommonAdapters and Intrinsics"
        )

        intrinsic_config = adapter.config
        assert intrinsic_config is not None

        rewriter = granite_common.IntrinsicsRewriter(
            config_dict=intrinsic_config, model_name=adapter.name
        )
        result_processor = granite_common.IntrinsicsResultProcessor(
            config_dict=intrinsic_config
        )

        # Convert our conversation into a proper chat completions dict.
        # [{role: user, content: Hello}, {...}] -> {messages: [{role:user,...}, ...], model:..., ...}
        request_json: dict = {
            "messages": conversation,
            "extra_body": {"documents": docs},
        }

        # Convert other parameters from Mellea proprietary format to standard format.
        for model_option in model_options:
            if model_option == ModelOption.TEMPERATURE:
                request_json["temperature"] = model_options[model_option]

        rewritten = rewriter.transform(request_json, **action.intrinsic_kwargs)

        # TODO: Handle caching here. granite_common doesn't tell us what changed,
        #       so we will have to invalidate the cache on our side. This requires
        #       us having specific caching for each Component/Message.

        generate_input, other_input = (
            granite_common.util.chat_completion_request_to_transformers_inputs(
                rewritten, self._tokenizer, self._model
            )
        )

        chat_response = asyncio.to_thread(
            self._generate_with_adapter_lock,
            adapter.qualified_name,
            granite_common.util.generate_with_transformers,
            # Passed as args/kwargs to generate.
            self._tokenizer,
            self._model,
            generate_input,
            other_input,
        )

        output = ModelOutputThunk(None)
        output._context = ctx.view_for_generation()
        output._action = action
        output._model_options = model_options

        # Add another step to the processing function.
        async def granite_common_processing(
            mot: ModelOutputThunk,
            chunk: granite_common.ChatCompletionResponse,
            rewritten: granite_common.ChatCompletion,
            result_processor: granite_common.IntrinsicsResultProcessor,
            input_ids,
        ):
            res = result_processor.transform(chunk, rewritten)  # type: ignore

            # TODO: If we want to support caches, we need to get the GenerateDecoderOnlyOutput. This means we
            #       probably need to break out the pieces from `generate_with_transformers`.
            # processing expects a str or a GenerateDecoderOnlyOutput. Extract the str.
            return await self.processing(
                mot, res.choices[0].message.content, input_ids=input_ids
            )

        output._process = functools.partial(
            granite_common_processing,
            rewritten=rewritten,
            result_processor=result_processor,
            input_ids=generate_input["input_tokens"],
        )

        # TODO: Post-processing should release the lock for this generation.
        output._post_process = functools.partial(
            self.post_processing,
            conversation=conversation,
            input_ids=generate_input["input_tokens"],
            _format=None,
            tool_calls=False,
            tools={},
            seed=seed,
        )

        try:
            # To support lazy computation, will need to remove this create_task and store just the unexecuted coroutine.
            # We can also support synchronous calls by adding a flag and changing this ._generate function.

            # This function should always be called from a running event loop so we don't have to worry about
            # scheduling the task to a specific event loop here.
            output._generate = asyncio.create_task(
                send_to_queue(chat_response, output._async_queue)  # type: ignore
            )
            output._generate_type = GenerateType.ASYNC
        except RuntimeError as e:
            # Most likely cause is running this function without an event loop present.
            raise e

        return output

    async def _generate_from_context_standard(
        self,
        action: Component | CBlock,
        ctx: Context,
        *,
        _format: type[BaseModelSubclass] | None = None,
        model_options: dict[str, Any],
        tool_calls: bool = False,
    ) -> ModelOutputThunk:
        # Construct input.
        # If the Context is a ChatHistory then we will pretty-print each content as a message and then use apply_chat_template.
        # Otherwise, we will linearize the context and treat it as a raw input.
        if ctx.is_chat_context:
            system_prompt = model_options.get(ModelOption.SYSTEM_PROMPT, None)
            ctx_as_chat = to_chat(action, ctx, self.formatter, system_prompt)

            # Append tool call information if applicable.
            tools: dict[str, Callable] = dict()
            if tool_calls:
                if _format:
                    FancyLogger.get_logger().warning(
                        f"Tool calling typically uses constrained generation, but you have specified a `format` in your generate call. NB: tool calling is superseded by format; we will NOT call tools for your request: {action}"
                    )
                else:
                    add_tools_from_model_options(tools, model_options)
                    add_tools_from_context_actions(
                        tools, ctx.actions_for_available_tools()
                    )

                    # Add the tools from the action for this generation last so that
                    # they overwrite conflicting names.
                    add_tools_from_context_actions(tools, [action])
                FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")

            seed = model_options.get(ModelOption.SEED, None)
            if seed is not None:
                set_seed(seed)

            input_ids = self._tokenizer.apply_chat_template(  # type: ignore
                ctx_as_chat,
                tools=convert_tools_to_json(tools),  # type: ignore
                add_generation_prompt=True,  # If we change this, must modify huggingface granite guardian.
                return_tensors="pt",
                **self._make_backend_specific_and_remove(model_options),
            ).to(self._device)  # type: ignore

            format_kwargs = {}
            if _format:
                # outlines.generate.json always parses the resulting json into a python dict.
                # We however want to keep it as a json string for later storing it in ModelOutputThunk
                schema: dict[str, Any] = _format.model_json_schema()
                schema_json: str = json.dumps(schema)
                regex_str: str = outlines_core.fsm.json_schema.build_regex_from_schema(  # type: ignore
                    schema_json
                )

                from outlines.models.transformers import TransformerTokenizer
                from outlines.processors.structured import RegexLogitsProcessor
                from transformers import LogitsProcessorList

                format_kwargs["logits_processor"] = LogitsProcessorList(
                    [
                        RegexLogitsProcessor(
                            regex_str, tokenizer=TransformerTokenizer(self._tokenizer)
                        )
                    ]
                )

            streaming_kwargs = {}
            streamer = None
            stream = model_options.get(ModelOption.STREAM, False)
            if stream:
                try:
                    # HuggingFace uses a streaming interface that you pass to the generate call.
                    # Must be called from a running event loop. This should always be the case given the same
                    # requirement of the ._generate function below.
                    streamer = AsyncTextIteratorStreamer(
                        self._tokenizer,  # type: ignore
                        skip_prompt=True,
                        skip_special_tokens=True,
                    )
                    streaming_kwargs["streamer"] = streamer
                except RuntimeError as e:
                    # Most likely cause is creating this object without an event loop present.
                    raise e

            # Create a separate thread to handle the processing. Make it awaitable
            # for non-streaming cases and to get the final output.
            # Details: https://huggingface.co/docs/transformers/en/internal/generation_utils#transformers.AsyncTextIteratorStreamer

            # Filter out chat template-only options before passing to generate()
            generate_options = self._filter_chat_template_only_options(model_options)

            chat_response = asyncio.to_thread(
                self._generate_with_adapter_lock,
                "",  # Empty for no adapters.
                self._model.generate,  # type: ignore
                # Passed as args/kwargs to generate.
                input_ids,
                return_dict_in_generate=True,
                output_scores=True,
                **self._make_backend_specific_and_remove(generate_options),
                **streaming_kwargs,  # type: ignore
                **format_kwargs,  # type: ignore
            )

            output = ModelOutputThunk(None)
            output._context = ctx.view_for_generation()
            output._action = action
            output._model_options = model_options

            # Processing functions only pass the ModelOutputThunk (and current chunk of response). Bind the other vars necessary for
            # each processing step.
            output._process = functools.partial(self.processing, input_ids=input_ids)
            output._post_process = functools.partial(
                self.post_processing,
                conversation=ctx_as_chat,
                input_ids=input_ids,
                _format=_format,
                tool_calls=tool_calls,
                tools=tools,
                seed=seed,
            )

            try:
                # To support lazy computation, will need to remove this create_task and store just the unexecuted coroutine.
                # We can also support synchronous calls by adding a flag and changing this ._generate function.

                response: AsyncTextIteratorStreamer | Coroutine = chat_response
                if stream and streamer is not None:
                    # For streaming, we want to pass the AsyncIterator to the function. Unlike other backends,
                    # this isn't returned by the chat_response coroutine. So we handle it here.
                    response = streamer

                    # Since the async iterator isn't returned by the chat_response coroutine, we have to create a separate
                    # task for it here so that it runs in the background. Attach it to the ModelOutputThunk.
                    output._generate_extra = asyncio.create_task(chat_response)

                # This function should always be called from a running event loop so we don't have to worry about
                # scheduling the task to a specific event loop here.
                output._generate = asyncio.create_task(
                    send_to_queue(response, output._async_queue)  # type: ignore
                )
                output._generate_type = GenerateType.ASYNC
            except RuntimeError as e:
                # Most likely cause is running this function without an event loop present.
                raise e

            return output

        else:
            raise Exception("Does not yet support non-chat contexts.")

    async def processing(
        self, mot: ModelOutputThunk, chunk: str | GenerateDecoderOnlyOutput, input_ids
    ):
        """Process the returned chunks or the complete response."""
        if mot._underlying_value is None:
            mot._underlying_value = ""

        # Because we use the AsyncTextIteratorStreamer, streaming responses are of type str;
        # and already decoded.
        if isinstance(chunk, str):
            mot._underlying_value += chunk
        else:
            # Otherwise, it's a non-streaming request. Decode it here.
            mot._meta["hf_output"] = chunk
            mot._underlying_value += self._tokenizer.decode(
                chunk.sequences[0, input_ids.shape[1] :], skip_special_tokens=True
            )

    async def post_processing(
        self,
        mot: ModelOutputThunk,
        conversation: list[dict],
        _format: type[BaseModelSubclass] | None,
        tool_calls: bool,
        tools: dict[str, Callable],
        seed,
        input_ids,
    ):
        """Called when generation is done."""
        if mot._meta.get("hf_output", None) is None:
            if mot._generate_extra is not None:
                full_output = await mot._generate_extra
                assert isinstance(full_output, GenerateDecoderOnlyOutput)
                mot._meta["hf_output"] = full_output

        # The ModelOutputThunk must be computed by this point.
        assert mot.value is not None

        # Add an entry to the cache for ALora reuse.
        if self._use_caches and mot._meta.get("hf_output", None) is not None:
            output_complete = mot._meta["hf_output"].sequences[0]
            cache: DynamicCache = mot._meta["hf_output"].past_key_values  # type: ignore

            cache_info = HFAloraCacheInfo(
                kv_cache=cache,
                merged_token_ids=output_complete,
                merged_attention=torch.ones_like(output_complete).to(self._device),
                q_end=len(input_ids[0]),  # type: ignore
            )

            self.cache_put(mot.value, cache_info)

        # Only scan for tools if we are not doing structured output and tool calls were provided to the model.
        if _format is None and tool_calls:
            mot.tool_calls = to_tool_calls(tools, mot.value)

        assert mot._action is not None, (
            "ModelOutputThunks should have their action assigned during generation"
        )
        assert mot._model_options is not None, (
            "ModelOutputThunks should have their model_opts assigned during generation"
        )

        self.formatter.parse(mot._action, mot)

        # Generate the log for this ModelOutputThunk.
        generate_log = GenerateLog()
        generate_log.prompt = conversation
        generate_log.backend = f"hf::{self.model_id!s}"
        generate_log.model_options = mot._model_options
        generate_log.date = datetime.datetime.now()
        generate_log.model_output = mot.value
        generate_log.extra = {
            "format": _format,
            "tools_available": tools,
            "tools_called": mot.tool_calls,
            "seed": seed,
        }
        generate_log.action = mot._action
        generate_log.result = mot

        mot._generate_log = generate_log

    async def generate_from_raw(
        self,
        actions: list[Component | CBlock],
        ctx: Context,
        *,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
    ) -> list[ModelOutputThunk]:
        """Generate using the completions api. Gives the input provided to the model without templating."""
        if tool_calls:
            FancyLogger.get_logger().warning(
                "The raw endpoint does not support tool calling at the moment."
            )

        if self._model.device.type == "mps":
            # TODO: Remove this when we are able to update the torch package.
            #       Test this by ensuring all outputs from this call are populated when running on mps.
            #       https://github.com/pytorch/pytorch/pull/157727
            FancyLogger.get_logger().warning(
                "utilizing device mps with a `generate_from_raw` request; you may see issues when submitting batches of prompts to a huggingface backend; ensure all ModelOutputThunks have non-empty values."
            )

        model_opts = self._simplify_and_merge(model_options)
        seed = model_opts.get(ModelOption.SEED, None)
        if seed is not None:
            set_seed(seed)

        prompts = [self.formatter.print(action) for action in actions]

        # batch-encoding call is deprecated in favor of this
        inputs = self._tokenizer(prompts, return_tensors="pt", padding=True).to(
            self._device
        )

        format_kwargs = {}
        if format:
            # outlines.generate.json always parses the resulting json into a python dict.
            # We however want to keep it as a json string for later storing it in ModelOutputThunk
            schema: dict[str, Any] = format.model_json_schema()
            schema_json: str = json.dumps(schema)
            regex_str: str = outlines_core.fsm.json_schema.build_regex_from_schema(  # type: ignore
                schema_json
            )

            from outlines.models.transformers import TransformerTokenizer
            from outlines.processors.structured import RegexLogitsProcessor
            from transformers import LogitsProcessorList

            format_kwargs["logits_processor"] = LogitsProcessorList(
                [
                    RegexLogitsProcessor(
                        regex_str, tokenizer=TransformerTokenizer(self._tokenizer)
                    )
                ]
            )

        outputs = await asyncio.to_thread(
            self._generate_with_adapter_lock,
            "",  # Empty for no adapter.
            self._model.generate,  # type: ignore
            # Passed as args/kwargs to generate.
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            return_dict_in_generate=True,
            output_scores=True,
            **self._make_backend_specific_and_remove(model_opts),
            **format_kwargs,
        )

        sequences_to_decode = [
            sequence[inputs["input_ids"][i].size(0) :]  # type: ignore
            for i, sequence in enumerate(outputs.sequences)
        ]

        decoded_results = self._tokenizer.batch_decode(
            sequences_to_decode, skip_special_tokens=True
        )

        results = []
        for i, decoded_result in enumerate(decoded_results):
            n_prompt_tokens = inputs["input_ids"][i].size(0)  # type: ignore
            n_completion_tokens = len(sequences_to_decode[i])
            result = ModelOutputThunk(
                value=decoded_result,
                meta={
                    "usage": {
                        "prompt_tokens": n_prompt_tokens,  # type: ignore
                        "completion_tokens": n_completion_tokens,
                        "total_tokens": n_prompt_tokens + n_completion_tokens,
                    }
                },
            )

            self.formatter.parse(actions[i], result)

            generate_log = GenerateLog()
            generate_log.prompt = self.formatter.print(actions[i])
            generate_log.backend = f"hf::{self.model_id!s}"
            generate_log.model_options = model_opts
            generate_log.date = datetime.datetime.now()
            generate_log.model_output = decoded_result
            generate_log.extra = {"format": format, "seed": seed}
            generate_log.action = actions[i]

            result._generate_log = generate_log
            results.append(result)

        return results

    # region cache management
    def cache_get(self, id: str) -> HFAloraCacheInfo | None:
        """Retrieve from cache."""
        v = self._cache.get(id)
        assert v is None or type(v) is HFAloraCacheInfo
        return v

    def cache_put(self, id: str, v: HFAloraCacheInfo):
        """Put into cache."""
        self._cache.put(id, v)

    # endregion

    def _simplify_and_merge(
        self, model_options: dict[str, Any] | None
    ) -> dict[str, Any]:
        """Simplifies model_options to use the Mellea specific ModelOption.Option and merges the backend's model_options with those passed into this call.

        Rules:
        - Within a model_options dict, existing keys take precedence. This means remapping to mellea specific keys will maintain the value of the mellea specific key if one already exists.
        - When merging, the keys/values from the dictionary passed into this function take precedence.

        Because this function simplifies and then merges, non-Mellea keys from the passed in model_options will replace
        Mellea specific keys from the backend's model_options.

        Common model options: https://huggingface.co/docs/transformers/en/llm_tutorial#common-options

        Args:
            model_options: the model_options for this call

        Returns:
            a new dict
        """
        backend_model_opts = ModelOption.replace_keys(
            self.model_options, self.to_mellea_model_opts_map
        )

        if model_options is None:
            return backend_model_opts

        generate_call_model_opts = ModelOption.replace_keys(
            model_options, self.to_mellea_model_opts_map
        )
        return ModelOption.merge_model_options(
            backend_model_opts, generate_call_model_opts
        )

    def _make_backend_specific_and_remove(
        self, model_options: dict[str, Any]
    ) -> dict[str, Any]:
        """Maps specified Mellea specific keys to their backend specific version and removes any remaining Mellea keys.

        Args:
            model_options: the model_options for this call

        Returns:
            a new dict
        """
        backend_specific = ModelOption.replace_keys(
            model_options, self.from_mellea_model_opts_map
        )
        return ModelOption.remove_special_keys(backend_specific)

    def _filter_chat_template_only_options(
        self, model_options: dict[str, Any]
    ) -> dict[str, Any]:
        """Remove options that are only for apply_chat_template, not for generate().

        Args:
            model_options: the model_options for this call

        Returns:
            a new dict without chat template-specific options
        """
        # Options that should only go to apply_chat_template, not generate()
        chat_template_only = {
            "guardian_config",
            "think",
            "add_generation_prompt",
            "documents",
        }
        return {k: v for k, v in model_options.items() if k not in chat_template_only}

    # region Adapter loading, unloading, and utility functions.
    @property
    def base_model_name(self):
        """Returns the base_model_id of the model used by the backend. For example, `granite-3.3-8b-instruct` for `ibm-granite/granite-3.3-8b-instruct`."""
        return self._hf_model_id.split("/")[1]

    def add_adapter(self, adapter: LocalHFAdapter):
        """Adds the given adapter to the backend. Must not have been added to a different backend."""
        if adapter.backend is not None:
            if adapter.backend is self:
                FancyLogger.get_logger().warning(
                    f"attempted to add adapter {adapter.name} with type {adapter.adapter_type} to the same backend {adapter.backend}"
                )
                return
            else:
                raise Exception(
                    f"adapter {adapter.name} with type {adapter.adapter_type} has already been added to backend {adapter.backend}"
                )

        if self._added_adapters.get(adapter.qualified_name) is not None:
            FancyLogger.get_logger().warning(
                f"Client code attempted to add {adapter.name} with type {adapter.adapter_type} but {adapter.name} was already added to {self.__class__}. The backend is refusing to do this, because adapter loading is not idempotent."
            )
            return None

        adapter.path = adapter.get_local_hf_path(self.base_model_name)
        adapter.backend = self
        self._added_adapters[adapter.qualified_name] = adapter

    def load_adapter(self, adapter_qualified_name: str):
        """Loads the given adapter for the backend. Must have previously been added. Do not call when generation requests are happening."""
        adapter = self._added_adapters.get(adapter_qualified_name, None)
        if adapter is None:
            raise ValueError(
                f"could not load adapter {adapter_qualified_name} for backend {self}: adapter was not previously added"
            )

        try:
            adapter_kwargs = {}

            # Peft tries to stringify the device. If it's mps, it gets stringified as "mps:0" which causes
            # an error when loading with safetensors.torch.load_file. Force the device as a string "mps" to fix.
            if self._device == torch.device("mps"):
                adapter_kwargs["device"] = "mps"
            self._model.load_adapter(
                adapter.path, adapter.qualified_name, adapter_kwargs=adapter_kwargs
            )
        except ValueError as e:
            # If it's just that it's already loaded, ignore it.
            if f"Adapter with name {adapter_qualified_name} already exists." not in str(
                e
            ):
                raise e

        # Loading an adapter activates it. We disable adapters immediately after.
        # Prefer this over `.disable_adapters()`; the disable function doesn't always
        # seem to work.
        self._model.disable_adapters()
        self._loaded_adapters[adapter.qualified_name] = adapter

    def unload_adapter(self, adapter_qualified_name: str):
        """Unloads the given adapter from the backend."""
        # Check if the backend knows about this adapter.
        adapter = self._loaded_adapters.get(adapter_qualified_name, None)
        if adapter is None:
            FancyLogger.get_logger().info(
                f"could not unload adapter {adapter_qualified_name} for backend {self}: adapter is not loaded"
            )
            return

        self._model.delete_adapter(adapter.qualified_name)

        # Remove the adapter from the list of loaded adapters.
        del self._loaded_adapters[adapter.qualified_name]

    def list_adapters(self) -> list[str]:
        """Lists the adapters added via add_adapter().

        :returns: list of adapter names that are currently registered with this backend
        """
        return list(self._loaded_adapters.keys())


def _assert_correct_adapters(expected_state: str, model: PreTrainedModel):
    """When generating with a huggingface model, this can be used to ensure the correct adapters are active.

    Args:
        expected_state: the current state of the lock
        model: the model underlying the LocalHFBackend; this is the model the adapters are activated on
    """
    try:
        active = model.active_adapters()

        if expected_state == "":
            assert len(active) == 0, (
                f'no adapters should be active if expected state is "", got "{active[0]}"'
            )
        else:
            assert len(active) == 1, (
                f'one adapter should be active if expected state is "{expected_state}"'
            )
            assert active[0] == expected_state, (
                f'the active adapter "{active[0]}" doesn\'t match the expected state: "{expected_state}"'
            )
    except ValueError as e:
        # If no weights have been loaded, the model will raise a ValueError:
        # `ValueError("No adapter loaded. Please load an adapter first.")`
        if "No adapter loaded" in str(e):
            assert expected_state == "", (
                f'got no adapters loaded but expected state is "{expected_state}"'
            )
        else:
            raise e


class HFProcessRewardModel(PRM, abc.ABC):
    """A Process Reward Model that works with a huggingface backend."""

    def __init__(
        self, model_name_or_path: str, score_token: str, device: str | None = None
    ):
        """Initialize an PRM that works with a huggingface backend. Currently supports and tested with IBM Process Reward Models.

        Args:
            model_name_or_path (str): A local path to PRM or a huggingface PRM
            score_token (str): token who's logits correspond to the PRM score. Can be a step demarker (for non-generative PRMs) or a correctness indicator (for generative PRMs)
            device (str): device: The computational device to use ("cuda" for GPU, "mps" for Apple Silicon, or "cpu"), defaults to None. If not specified, the best available device will be automatically selected.
        """
        super().__init__(model_name_or_path)

        # auto-device if not more specific
        self._device = device
        if device is None:
            device_name: str = (
                "cuda"
                if torch.cuda.is_available()
                else "mps"
                if torch.backends.mps.is_available()
                else "cpu"
            )
            assert device_name is not None
            self._device = torch.device(device_name)  # type: ignore

        self.model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
            self.model_name_or_path, torch_dtype=torch.bfloat16
        )
        self.model.eval()
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)

        self._score_token = score_token
        self._score_token_id = self.tokenizer.encode(
            self._score_token, add_special_tokens=False
        )[0]

    def stepify(self, content: str, step_separator: str) -> list[str]:
        """Splits the assistant response into steps to score.

        Args:
            content: assistant response to score
            step_separator: string on which to separate the content into steps
        """
        # convert assistant message into a list of steps
        list_of_steps = [
            step.strip() for step in content.split(step_separator) if step.strip != ""
        ]
        return list_of_steps
