# Wan UMT5-XXL text encoder wrapper for training-time per-sample text encoding.
# Loads the Wan UMT5 model + HuggingFace tokenizer and presents the
# `encode_prompts(prompts, max_length, return_mask)` API expected by a4d_vae_wan.py.

from typing import List, Tuple, Union

import torch

from custom.wan.build_default_text_wan import WanUMT5Encoder
from custom.utils.utils import get_checkpoint
from imaginaire.utils import log


MAX_LENGTH = 512


class WanTextEncoder:
    """Wraps the Wan UMT5-XXL encoder + HF tokenizer for use during training.

    Provides `encode_prompts(prompts, max_length, return_mask)` matching the
    interface expected by `a4d_vae_wan.VAE.encode_text`.
    """

    def __init__(self, weights_path: str, tokenizer_path: str,
                 device: str = 'cuda', dtype: torch.dtype = torch.bfloat16):
        self.device = device
        self.dtype = dtype

        log.info(f'[WanTextEncoder] Loading UMT5 tokenizer from {tokenizer_path}')
        from transformers import AutoTokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(get_checkpoint(tokenizer_path))

        log.info(f'[WanTextEncoder] Loading UMT5 weights from {weights_path} ...')
        self.model = WanUMT5Encoder().to(dtype=dtype).eval()
        sd = torch.load(get_checkpoint(weights_path), map_location='cpu', weights_only=True)
        miss, surp = self.model.load_state_dict(sd, strict=False)
        if miss:
            log.warning(f'[WanTextEncoder] Missing keys: {miss[:5]}')
        if surp:
            log.warning(f'[WanTextEncoder] Surplus keys: {surp[:5]}')
        del sd
        self.model = self.model.to(device)
        log.info(f'[WanTextEncoder] Loaded UMT5-XXL ({sum(p.numel() for p in self.model.parameters()):,} params)')

    @torch.no_grad()
    def encode_prompts(
        self,
        prompts: Union[str, List[str]],
        max_length: int = MAX_LENGTH,
        return_mask: bool = False,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """Encode a batch of text prompts into UMT5-XXL embeddings.

        Matches DiffSynth's WanPrompter.encode_prompt: tokenize, encode,
        then zero positions beyond each prompt's real sequence length.

        :param prompts: single string or list of strings.
        :param max_length: max token length (default 512).
        :param return_mask: if True, return (embeddings, mask) tuple.
        :return: (B, max_length, 4096) embeddings, optionally with (B, max_length) mask.
        """
        import ftfy
        import html
        import re

        if isinstance(prompts, str):
            prompts = [prompts]

        cleaned = []
        for p in prompts:
            text = ftfy.fix_text(p)
            text = html.unescape(html.unescape(text))
            text = re.sub(r'\s+', ' ', text).strip()
            cleaned.append(text)

        enc = self.tokenizer(
            cleaned, padding='max_length', truncation=True,
            max_length=max_length, return_tensors='pt', add_special_tokens=True)
        ids = enc['input_ids'].to(self.device)
        mask = enc['attention_mask'].to(self.device)

        emb = self.model(ids, mask=mask)  # (B, L, 4096)

        seq_lens = mask.gt(0).sum(dim=1).long()
        for i, v in enumerate(seq_lens):
            emb[i, v:] = 0

        emb = emb.to(dtype=self.dtype)
        mask = mask.to(dtype=self.dtype)

        if return_mask:
            return emb, mask
        return emb