"""FLAN-T5-Large prompt encoder used by MiniT2I."""
from __future__ import annotations

from typing import Optional

import torch

_T5_NAME = "google/flan-t5-large"
_SINGLETON: dict[tuple, "T5PromptEncoder"] = {}


class T5PromptEncoder:
    def __init__(
        self,
        model_name: str = _T5_NAME,
        device: str | torch.device = "cuda",
        dtype: torch.dtype = torch.float32,
        max_len: int = 256,
        cache_dir: Optional[str] = None,
    ):
        from transformers import AutoTokenizer, T5EncoderModel

        self.max_len = max_len
        self.device = torch.device(device)
        self.dtype = dtype
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
        self.model = (
            T5EncoderModel.from_pretrained(model_name, cache_dir=cache_dir)
            .to(self.device, dtype=dtype)
            .eval()
        )
        for p in self.model.parameters():
            p.requires_grad_(False)

    @torch.no_grad()
    def encode(self, prompts: list[str], max_len: Optional[int] = None):
        max_len = int(max_len or self.max_len)
        tok = self.tokenizer(
            list(prompts),
            padding="max_length",
            truncation=True,
            max_length=max_len,
            return_tensors="pt",
        )
        input_ids = tok.input_ids.to(self.device)
        attn_mask = tok.attention_mask.to(self.device)
        out = self.model(input_ids=input_ids, attention_mask=attn_mask)
        context = out.last_hidden_state.to(self.dtype)
        return context, attn_mask.to(self.dtype)


def get_prompt_encoder(
    model_name: str = _T5_NAME,
    device: str | torch.device = "cuda",
    dtype: torch.dtype = torch.float32,
    max_len: int = 256,
    cache_dir: Optional[str] = None,
) -> T5PromptEncoder:
    """Cached accessor so repeated calls reuse a single loaded T5 encoder."""
    key = (model_name, str(device), dtype, max_len, cache_dir)
    if key not in _SINGLETON:
        _SINGLETON[key] = T5PromptEncoder(model_name, device, dtype, max_len, cache_dir)
    return _SINGLETON[key]