# extract_sdxl_embeddings.py
import argparse
from pathlib import Path
from typing import List

import torch
from safetensors.torch import save_file
from diffusers import StableDiffusionXLPipeline


def read_prompts(txt_path: str) -> List[str]:
    with open(txt_path, "r", encoding="utf-8") as f:
        return [line.rstrip("\n") for line in f]


def load_sdxl(checkpoint_path: str, precision: str):
    precision = precision.lower()
    if precision == "bf16":
        dtype = torch.bfloat16
    else:
        dtype = torch.float16  # T4 suele ir mejor así para SDXL

    path = Path(checkpoint_path)
    if path.is_dir():
        pipe = StableDiffusionXLPipeline.from_pretrained(
            checkpoint_path,
            torch_dtype=dtype,
            use_safetensors=True,
        )
    else:
        # Útil para .safetensors / .ckpt de un solo archivo.
        pipe = StableDiffusionXLPipeline.from_single_file(
            checkpoint_path,
            torch_dtype=dtype,
        )

    pipe.to("cuda" if torch.cuda.is_available() else "cpu")
    pipe.set_progress_bar_config(disable=True)
    pipe.eval()
    return pipe


@torch.no_grad()
def encode_batch(pipe: StableDiffusionXLPipeline, batch_prompts: List[str]):
    device = pipe._execution_device if hasattr(pipe, "_execution_device") else next(pipe.text_encoder.parameters()).device

    # Diffusers soporta prompt_embeds y pooled_prompt_embeds en SDXL. :contentReference[oaicite:3]{index=3}
    prompt_embeds, pooled_prompt_embeds = pipe.encode_prompt(
        prompt=batch_prompts,
        prompt_2=batch_prompts,
        device=device,
        num_images_per_prompt=1,
        do_classifier_free_guidance=False,
    )[:2]

    return prompt_embeds.detach().cpu(), pooled_prompt_embeds.detach().cpu()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--sdxl_checkpoint", type=str, required=True,
                        help="Ruta al .safetensors / .ckpt o directorio Diffusers de SDXL.")
    parser.add_argument("--prompts_txt", type=str, required=True)
    parser.add_argument("--out_dir", type=str, default="output_embeddings")
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--precision", type=str, default="fp16", choices=["fp16", "bf16"])
    parser.add_argument("--pad_width", type=int, default=5)
    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    prompts = read_prompts(args.prompts_txt)
    pipe = load_sdxl(args.sdxl_checkpoint, args.precision)

    n = len(prompts)
    print(f"Procesando {n} prompts...")

    for i in range(0, n, args.batch_size):
        batch_prompts = prompts[i:i + args.batch_size]
        text_embeds, pooled_text_embeds = encode_batch(pipe, batch_prompts)

        for b in range(text_embeds.shape[0]):
            file_idx = i + b
            file_name = f"{file_idx:0{args.pad_width}d}.safetensors"
            save_path = out_dir / file_name

            sample = {
                "text_embeds": text_embeds[b:b+1].contiguous().to(torch.float16),
                "pooled_text_embeds": pooled_text_embeds[b:b+1].contiguous().to(torch.float16),
            }
            save_file(sample, str(save_path))

        print(f"Guardado hasta {min(i + args.batch_size - 1, n - 1):0{args.pad_width}d}")

    print(f"Listo. Salida en: {out_dir}")


if __name__ == "__main__":
    main()