from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

from cid.accelerator import wrap_torch_autocast
from cid.model import load_cid_model_from_pretrained

from .config import MODEL_ID, MODEL_REVISION


@dataclass(frozen=True)
class ModelBundle:
    adapter: Any
    tokenizer: Any
    text_encoder: Any
    forward_model: Any


def load_model_bundle() -> ModelBundle:
    model_dir = Path(
        snapshot_download(
            MODEL_ID,
            revision=MODEL_REVISION,
            allow_patterns=["*.json", "*.jinja", "*.safetensors", "*.py"],
        )
    )
    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

    device = torch.device("cuda")
    cid_model = load_cid_model_from_pretrained(str(model_dir), low_cpu_mem_usage=True)
    adapter = cid_model.cid_adapter.to(device)
    adapter.set_backbone_trainable(True)
    text_encoder = cid_model.make_text_encoder(
        tokenizer,
        device=device,
        embedding_device="cpu",
    )
    forward_model = wrap_torch_autocast(
        torch,
        adapter,
        device_type="cuda",
        dtype=torch.bfloat16,
    )
    adapter.eval()
    forward_model.eval()
    return ModelBundle(adapter, tokenizer, text_encoder, forward_model)


MODEL_BUNDLE = load_model_bundle()