#
# Text pipeline ported from Apple ml-starflow utils/common.py::preprocess_text
# and utils/common.py::encode_text (commit 20984834c1a68ee679faa8ed9cda2842f02b042d),
# restricted to the single text encoder used by STARFlow-V (google/flan-t5-xl,
# a plain HF encoder-only model, no chat-template / Gemma / label-embedding path).
# For licensing see accompanying LICENSE file.
# Copyright (C) 2025 Apple Inc. All Rights Reserved.
#
"""
Prompt wrapping and FLAN-T5-XL text encoding for STARFlow-V.

Exact upstream wrapping order (verified against utils/common.py):
    1. fps wrap:        "A video with {fps} fps:\n{prompt}\n"   (only if fps is not None)
    2. noise prepend:   "Noise Level {int(noise_std*1000)}:\n{wrapped}"
    3. aspect append:   "{wrapped}\n in a {aspect_ratio} aspect ratio.\n"

The unconditional ("") caption used for classifier-free guidance goes through
the exact same wrapping (it receives the same fps/noise_std/aspect_ratio
metadata as the conditional caption), matching upstream sample.py which
appends "" to the caption list *before* calling encode_text.
"""
import os
from typing import List, Optional, Union

import torch


def preprocess_text(
    text: List[str],
    aspect_ratio: Optional[str] = None,
    fps: Optional[Union[int, torch.Tensor]] = None,
    noise_std: Optional[Union[float, torch.Tensor]] = None,
) -> List[str]:
    if fps is not None:
        if isinstance(fps, torch.Tensor):
            fps = [int(f) for f in fps.tolist()]
        elif isinstance(fps, int):
            fps = [fps] * len(text)
        text = [f"A video with {f} fps:\n{txt}\n" if f > 0 else f"An image:\n{txt}\n"
                for txt, f in zip(text, fps)]

    if noise_std is not None:
        if isinstance(noise_std, torch.Tensor):
            noise_std = [int(n * 1000) for n in noise_std.view(-1).tolist()]
        elif isinstance(noise_std, float):
            noise_std = [int(noise_std * 1000)] * len(text)
        text = [f'Noise Level {n}:\n{txt}' for n, txt in zip(noise_std, text)]

    if aspect_ratio is not None:
        text = [f"{txt}\n in a {aspect_ratio} aspect ratio.\n" for txt in text]

    return text


@torch.no_grad()
def encode_text(text_encoder, tokenizer, text: List[str], max_length: int, device,
                 return_tokens: bool = False, **kwargs):
    text = preprocess_text(text, **kwargs)
    tokenized_outputs = tokenizer(
        text, padding="max_length", truncation=True, return_tensors="pt", max_length=max_length)
    tokenized_outputs = {key: val.to(device) for key, val in tokenized_outputs.items()}
    y = text_encoder(**tokenized_outputs).last_hidden_state
    y = y * tokenized_outputs['attention_mask'].unsqueeze(-1)  # mask out padding
    if return_tokens:
        return y, tokenized_outputs
    return y


class STARFlowVTextEncoder:
    """
    Loads a local FLAN-T5-XL directory (tokenizer + encoder-only T5 weights)
    with no network access. Kept dependency-light: only transformers is
    required (no sentencepiece network fetch, no accelerate device_map
    auto-download).
    """

    hidden_size = 2048  # google/flan-t5-xl d_model

    def __init__(self, model_dir: str, dtype: torch.dtype = torch.bfloat16):
        if not os.path.isdir(model_dir):
            raise FileNotFoundError(
                f"STARFlow-V text encoder directory not found: {model_dir}. "
                f"Provide a local FLAN-T5-XL checkout (config.json, tokenizer files, "
                f"and model weights); this loader never fetches from the network."
            )
        from transformers import AutoTokenizer, T5EncoderModel, T5Config

        self.model_dir = model_dir
        self.dtype = dtype

        config = T5Config.from_pretrained(model_dir, local_files_only=True)
        if config.d_model != self.hidden_size:
            raise ValueError(
                f"Expected FLAN-T5-XL (d_model={self.hidden_size}), got d_model={config.d_model} "
                f"from {model_dir}. STARFlow-V's checkpoint proj_txt layers require this exact size."
            )

        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
        self.model = T5EncoderModel.from_pretrained(
            model_dir, config=config, torch_dtype=dtype, local_files_only=True)
        self.model.requires_grad_(False)
        self.model.eval()

        self.device = torch.device("cpu")
        self.load_device = torch.device("cpu")

    def to(self, device):
        self.model.to(device)
        self.device = torch.device(device)
        return self

    def model_size(self) -> int:
        return sum(p.numel() * p.element_size() for p in self.model.parameters())

    @torch.no_grad()
    def encode(self, prompts: List[str], txt_size: int, device,
               aspect_ratio: Optional[str] = None,
               fps: Optional[int] = None,
               noise_std: Optional[float] = None) -> torch.Tensor:
        self.model.to(device)
        return encode_text(
            self.model, self.tokenizer, prompts, txt_size, device,
            aspect_ratio=aspect_ratio, fps=fps, noise_std=noise_std,
        )