import os
import torch
import torch.nn as nn
from transformers import CLIPTokenizer, CLIPTextModel


class FrozenCLIPTextEncoder(nn.Module):
    def __init__(self, model_path_or_id: str, device: str = "cuda"):
        super().__init__()
        self.device = device
        
        if os.path.isdir(model_path_or_id):
            assert os.path.exists(os.path.join(model_path_or_id, "config.json")), \
                f"Error: Local path '{model_path_or_id}' does not contain 'config.json'. Check your weight path."
            
            self.tokenizer = CLIPTokenizer.from_pretrained(model_path_or_id, local_files_only=True)
            self.text_encoder =  CLIPTextModel.from_pretrained(model_path_or_id, local_files_only=True).to(self.device)
        else:
            self.tokenizer = CLIPTokenizer.from_pretrained(model_path_or_id)
            self.text_encoder =  CLIPTextModel.from_pretrained(model_path_or_id).to(self.device)
        
        self.text_encoder.eval()
        for param in self.text_encoder.parameters():
            param.requires_grad = False

    @torch.no_grad()
    def forward(self, text: str or list[str]) -> torch.Tensor:
        if isinstance(text, str):
            text = [text]

        inputs = self.tokenizer(
            text,
            padding="max_length",
            truncation=True,
            max_length=77,
            return_tensors="pt",
        ).to(self.device)

        outputs = self.text_encoder(**inputs)

        text_tokens = outputs.last_hidden_state          # [B, 77, 512]
        text_tokens = text_tokens / text_tokens.norm(dim=-1, keepdim=True).clamp_min(1e-6)

        return text_tokens