import torch

from acestep.constants import DEFAULT_DIT_INSTRUCTION, SFT_GEN_PROMPT

from .models import AudioSample
from .preprocess_utils import build_metas_str


def build_text_prompt(sample: AudioSample, tag_position: str, use_genre: bool) -> str:
    """Build the text prompt for the text encoder."""
    caption = sample.get_training_prompt(tag_position, use_genre=use_genre)
    metas_str = build_metas_str(sample)
    return SFT_GEN_PROMPT.format(DEFAULT_DIT_INSTRUCTION, caption, metas_str)


def encode_text(text_encoder, text_tokenizer, text_prompt: str, device, dtype):
    """Encode caption/genre prompt into text hidden states."""
    text_inputs = text_tokenizer(
        text_prompt,
        padding="max_length",
        max_length=256,
        truncation=True,
        return_tensors="pt",
    )
    text_input_ids = text_inputs.input_ids.to(device)
    text_attention_mask = text_inputs.attention_mask.to(device).to(dtype)

    text_dev = next(text_encoder.parameters()).device
    if text_input_ids.device != text_dev:
        text_input_ids = text_input_ids.to(text_dev)
        text_attention_mask = text_attention_mask.to(text_dev)

    with torch.no_grad():
        text_outputs = text_encoder(text_input_ids)
        text_hidden_states = text_outputs.last_hidden_state.to(dtype)

    return text_hidden_states, text_attention_mask