import os
from dataclasses import dataclass
from pathlib import Path
import torch
from einops import rearrange
from PIL import Image
from safetensors.torch import load_file as load_sft
from huggingface_hub import hf_hub_download
from flux.modules.autoencoder import AutoEncoder, AutoEncoderParams

CHECKPOINTS_DIR = Path("checkpoints")
CHECKPOINTS_DIR.mkdir(exist_ok=True)


@dataclass
class VAEConfig:
    repo_id: str
    repo_ae: str
    ae_params: AutoEncoderParams


flux_vae_config = VAEConfig(
    repo_id="black-forest-labs/FLUX.1-dev",
    repo_ae="ae.safetensors",
    ae_params=AutoEncoderParams(
        resolution=64,
        in_channels=3,
        ch=128,
        out_ch=3,
        ch_mult=[1, 2, 4, 4],
        num_res_blocks=2,
        z_channels=16,
        scale_factor=0.3611,
        shift_factor=0.1159,
    ),
)


def get_checkpoint_path(filename: str, repo_id: str) -> Path:
    local_path = CHECKPOINTS_DIR / filename
    if not local_path.exists():
        print(f"Downloading {filename} from {repo_id}...")
        try:
            downloaded_path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=CHECKPOINTS_DIR)
            return Path(downloaded_path)
        except Exception as e:
            print(f"Failed to download: {e}")
    return local_path


def load_ae(device: str | torch.device = "cuda") -> AutoEncoder:
    print("Init AE")
    ckpt_path = get_checkpoint_path(flux_vae_config.repo_ae, flux_vae_config.repo_id)

    if ckpt_path.exists():
        print(f"Loading AE checkpoint: {ckpt_path}")
        with torch.device("meta"):
            ae = AutoEncoder(flux_vae_config.ae_params)
        sd = load_sft(str(ckpt_path), device=str(device))
        ae.load_state_dict(sd, strict=False, assign=True)
    else:
        raise RuntimeError("AE Checkpoint is missing! Model cannot train with random VAE.")

    return ae


def save_image(x: torch.Tensor, path: str):
    x = x.clamp(-1, 1)
    x = rearrange(x[0], "c h w -> h w c")
    img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
    img.save(path, quality=95, subsampling=0)


def print_load_warning(missing: list[str], unexpected: list[str]):
    if len(missing) > 0 and len(unexpected) > 0:
        print(f"Parameters missing: {len(missing)}. Unexpected: {len(unexpected)}")