"""Helpers shared by the Kimodo blocks."""

from dataclasses import replace

import torch
from diffusers.modular_pipelines import ModularPipelineBlocks


class ResolvableBlock(ModularPipelineBlocks):
    """Leaf block whose generic component specs can carry the model repository's resolved classes and location.

    `KimodoBlocks(components_repo=...)` fills `component_overrides`, so a sub-block popped out of the pipeline still
    loads the right custom classes from the right repository.
    """

    component_overrides = {}

    def resolve(self, specs):
        return [
            replace(spec, **self.component_overrides[spec.name]) if spec.name in self.component_overrides else spec
            for spec in specs
        ]


class Cancelled(Exception):
    pass


def checkpoint(cancel):
    if cancel is not None and cancel.is_set():
        raise Cancelled("Generation cancelled")


@torch.no_grad()
def encode_prompt(text_encoder, tokenizer, prompt, device):
    """Encode one prompt into LLM2Vec features shaped (1, 1, hidden)."""
    from .prompt_utils import tokenize_prompt
    from .sanitize import sanitize_texts

    inputs = tokenize_prompt(tokenizer, sanitize_texts([prompt])[0])
    features = text_encoder(**{name: value.to(text_encoder.device) for name, value in inputs.items()})
    return features.to(device)