#' FLUX Transformer Quantization and Loading
#'
#' Quantize the 12B FLUX transformer to NF4 (~7 GB, GPU-resident on
#' 16 GB cards) or fp8 (~12 GB, CPU-resident and streamed per forward),
#' and load any format back into \code{\link{flux_transformer}}. Reuses
#' the LTX-2.3 quantization machinery (\code{ltx23_nf4_quantize},
#' \code{ltx23_nf4_linear}, \code{ltx23_fp8_linear}); only the cast set
#' and the diffusers directory layout are FLUX-specific.
#'
#' @name quantize_flux
NULL

.flux_dtype <- function(dtype) {
    switch(dtype, bfloat16 = torch::torch_bfloat16(),
           float16 = torch::torch_float16(), float32 = torch::torch_float32(),
           stop("Unsupported dtype: ", dtype))
}

# Capability probe: can the installed safetensors round-trip this dtype?
# CRAN safetensors (<= 0.2.1) can read bfloat16 but not write it, and has
# no float8 support; the fixes are upstream PRs. Cached per session;
# options(diffuseR.st_caps = list(bfloat16 = FALSE, ...)) overrides for
# tests.
.st_caps <- new.env(parent = emptyenv())
.st_can_write <- function(dtype = c("bfloat16", "float8_e4m3fn")) {
    dtype <- match.arg(dtype)
    override <- getOption("diffuseR.st_caps")
    if (!is.null(override) && !is.null(override[[dtype]])) {
        return(isTRUE(override[[dtype]]))
    }
    cached <- .st_caps[[dtype]]
    if (!is.null(cached)) {
        return(cached)
    }
    ok <- requireNamespace("safetensors", quietly = TRUE) && tryCatch({
        target <- switch(dtype, bfloat16 = torch::torch_bfloat16(),
                         float8_e4m3fn = torch::torch_float8_e4m3fn())
        x <- torch::torch_zeros(2L)$to(dtype = target)
        tmp <- tempfile(fileext = ".safetensors")
        on.exit(unlink(tmp), add = TRUE)
        safetensors::safe_save_file(list(w = x), tmp)
        y <- safetensors::safe_load_file(tmp, framework = "torch")
        !is.null(y$w)
    }, error = function(e) FALSE)
    .st_caps[[dtype]] <- ok
    ok
}

# Resolve precision = "auto": prefer an existing quantized artifact
# (fp8 first), else pick by float8 write capability. An fp8 artifact is
# only chosen if the installed safetensors can actually read float8 -
# otherwise a fork-built fp8 artifact on a CRAN-safetensors machine
# would be selected and then fail at read time. Write capability is a
# sound proxy for read capability (nothing writes fp8 but cannot read
# it).
.flux_resolve_precision <- function(precision, artifact_prefix = NULL) {
    if (!identical(precision, "auto")) {
        return(precision)
    }
    fp8_ok <- .st_can_write("float8_e4m3fn")
    if (!is.null(artifact_prefix)) {
        for (p in c("fp8", "nf4")) {
            if (p == "fp8" && !fp8_ok) {
                next
            }
            if (file.exists(file.path(paste0(artifact_prefix, p),
                                      "manifest.json"))) {
                return(p)
            }
        }
    }
    if (fp8_ok) {
        "fp8"
    } else {
        "nf4"
    }
}

# Locate a flux-family model's unquantized transformer in the hfhub
# cache. bf16 is not a built artifact: it IS the source the quantizers
# read, so the "bf16 tier" means loading those weights directly
# (flux_load_transformer() handles them as format "full"). Resolves from
# the cache only -- a loader must never trigger a multi-GB download as a
# side effect of a precision choice.
.flux_source_dir <- function(model) {
    if (!requireNamespace("hfhub", quietly = TRUE)) {
        stop("The hfhub package is required to locate bf16 source weights.",
             call. = FALSE)
    }
    spec <- switch(model,
                   flux1 = list(repo = .flux1_repo,
                                files = .flux1_transformer_files,
                                fn = "download_flux1"),
                   flux2 = list(repo = .flux2_repo,
                                files = .flux2_transformer_files,
                                fn = "download_flux2_klein"),
                   zimage = list(repo = .zimage_repo,
                                 files = .zimage_transformer_files,
                                 fn = "download_zimage_turbo"),
                   stop("No bf16 source known for model '", model, "'",
                        call. = FALSE))
    paths <- lapply(spec$files, function(f) {
        tryCatch(hfhub::hub_download(spec$repo, f, local_files_only = TRUE),
                 error = function(e) NULL)
    })
    if (any(vapply(paths, is.null, logical(1)))) {
        stop(sprintf(paste0("bf16 needs the unquantized %s transformer, ",
                            "which is not in the HuggingFace cache. Run ",
                            "%s(quantize = FALSE) to fetch it, or pick a ",
                            "quantized precision."),
                     spec$repo, spec$fn), call. = FALSE)
    }
    # Every file list starts with transformer/config.json, so the parent
    # of the first hit is the transformer directory.
    dirname(paths[[1]])
}

# Where a loader should read weights for a resolved precision: the
# built artifact for nf4/fp8, the cached source for bf16.
.flux_model_dir <- function(model, precision, prefix) {
    if (identical(precision, "bf16")) {
        .flux_source_dir(model)
    } else {
        paste0(prefix, precision)
    }
}

# Transformer constructor arguments from a diffusers config.json
.flux_transformer_args <- function(config) {
    if (is.null(config)) {
        return(list())
    }
    if (isTRUE(config$guidance_embeds)) {
        stop("This checkpoint uses guidance embeddings (FLUX.1-dev); ",
             "only FLUX.1-schnell (guidance_embeds = false) is supported.")
    }
    args <- list(
                 in_channels = config$in_channels,
                 num_layers = config$num_layers,
                 num_single_layers = config$num_single_layers,
                 attention_head_dim = config$attention_head_dim,
                 num_attention_heads = config$num_attention_heads,
                 joint_attention_dim = config$joint_attention_dim,
                 pooled_projection_dim = config$pooled_projection_dim,
                 axes_dims_rope = config$axes_dims_rope,
                 out_channels = config$out_channels
    )
    # JSON roundtrips turn null into empty lists; drop both
    args <- Filter(function(x) !is.null(x) && length(x) > 0L, args)
    lapply(args, function(x) if (is.numeric(x)) as.integer(x) else x)
}

.flux2_transformer_args <- function(config) {
    if (is.null(config)) {
        return(list())
    }
    if (isTRUE(config$guidance_embeds)) {
        stop("This checkpoint uses guidance embeddings (FLUX.2-dev); ",
             "only the klein variants (guidance_embeds = false) are supported.")
    }
    args <- list(
                 in_channels = config$in_channels,
                 num_layers = config$num_layers,
                 num_single_layers = config$num_single_layers,
                 attention_head_dim = config$attention_head_dim,
                 num_attention_heads = config$num_attention_heads,
                 joint_attention_dim = config$joint_attention_dim,
                 timestep_guidance_channels = config$timestep_guidance_channels,
                 axes_dims_rope = config$axes_dims_rope,
                 out_channels = config$out_channels
    )
    args <- Filter(function(x) !is.null(x) && length(x) > 0L, args)
    args <- lapply(args, function(x) if (is.numeric(x)) as.integer(x) else x)
    for (field in c("mlp_ratio", "rope_theta", "eps")) {
        v <- config[[field]]
        if (!is.null(v) && length(v) == 1L) {
            args[[field]] <- as.numeric(v)
        }
    }
    args
}

.zimage_transformer_args <- function(config) {
    if (is.null(config)) {
        return(list())
    }
    args <- list(in_channels = config$in_channels, dim = config$dim,
                 n_layers = config$n_layers,
                 n_refiner_layers = config$n_refiner_layers,
                 n_heads = config$n_heads, cap_feat_dim = config$cap_feat_dim,
                 axes_dims = config$axes_dims,
                 patch_size = config$all_patch_size[[1]],
                 f_patch_size = config$all_f_patch_size[[1]])
    args <- Filter(function(x) !is.null(x) && length(x) > 0L, args)
    args <- lapply(args, function(x) if (is.numeric(x)) as.integer(x) else x)
    for (field in c("norm_eps", "rope_theta", "t_scale")) {
        v <- config[[field]]
        if (!is.null(v) && length(v) == 1L) {
            args[[field]] <- as.numeric(v)
        }
    }
    args
}

# Move plain-field fp8 weights (and their scales) onto a device; used
# for resident fp8 where the whole quantized model fits on the GPU
.flux_fp8_to_device <- function(module, device) {
    for (name in names(module$children)) {
        child <- module$children[[name]]
        if (!is.null(child$weight_fp8)) {
            child$weight_fp8 <- child$weight_fp8$to(device = device)
            child$weight_scale <- child$weight_scale$to(device = device)
        }
        .flux_fp8_to_device(child, device)
    }
    invisible(module)
}

# Collect an fp8 transformer's plain-field weight tensors. weight_fp8
# and weight_scale live outside parameters/buffers, so .pin_component's
# module walk misses them; this mirrors .flux_fp8_to_device's recursion
# and returns the tensors so they can be pinned and staged alongside the
# module. set_data (in .pin_component) mutates each in place, so the
# field reference stays valid and staged onload/offload covers them -
# replacing the per-generation .flux_fp8_to_device reassignment.
.flux_fp8_collect <- function(module, out = list()) {
    for (name in names(module$children)) {
        child <- module$children[[name]]
        if (!is.null(child$weight_fp8)) {
            out[[length(out) + 1L]] <- child$weight_fp8
            out[[length(out) + 1L]] <- child$weight_scale
        }
        out <- .flux_fp8_collect(child, out)
    }
    out
}

# Pin the phase-swapped components of a flux-family pipeline so the
# per-generation CPU<->GPU moves run at DMA rate (see staging.R).
# Returns a named list of staging pairs keyed by component, or NULL when
# staging does not apply (opted out, no phase offload, or CPU device).
# Only the listed components are pinned; text encoders resident on the
# host (text_device == "cpu") are omitted by the caller since they
# compute in place. Resident-fp8 transformers (flux2/zimage) also pin
# their plain weight_fp8/weight_scale fields so staged transfer covers
# them; nf4 packed weights are buffers and ride along automatically.
.flux_build_staging <- function(pipe, pin, phase_offload, device, components,
                                verbose = TRUE) {
    if (!(isTRUE(pin) && isTRUE(phase_offload) &&
            identical(device, "cuda") && torch::cuda_is_available())) {
        return(NULL)
    }
    if (verbose) {
        message("Pinning host staging buffers...")
    }
    staging <- list()
    for (nm in components) {
        module <- pipe[[nm]]
        if (is.null(module)) {
            next
        }
        extra <- if (identical(nm, "transformer") &&
            isTRUE(pipe$fp8_resident)) {
            .flux_fp8_collect(module)
        } else {
            NULL
        }
        st <- .pin_component(module, extra = extra)
        if (!is.null(st)) {
            staging[[nm]] <- st
        }
    }
    if (length(staging)) {
        staging
    } else {
        NULL
    }
}

# Family-specific hooks for quantization and loading
.flux_family_hooks <- function(config) {
    family <- .flux_family(config)
    if (family == "flux2") {
        list(model_fn = flux2_transformer, args_fn = .flux2_transformer_args,
             is_quant_key = flux2_is_quant_key, shard_prefix = "flux2-klein")
    } else if (family == "zimage") {
        list(model_fn = zimage_transformer,
             args_fn = .zimage_transformer_args,
             is_quant_key = zimage_is_quant_key,
             shard_prefix = "zimage-turbo")
    } else {
        list(model_fn = flux_transformer,
             args_fn = .flux_transformer_args,
             is_quant_key = flux_is_quant_key,
             shard_prefix = "flux1")
    }
}

#' Quantize a FLUX transformer to NF4 or fp8 shards
#'
#' Streams the bf16 diffusers checkpoint tensor by tensor. Cast-set
#' weights (see \code{\link{flux_is_quant_key}}) are stored as NF4
#' (packed uint8 + \code{<key>_absmax} float32 blocks) or as
#' float8_e4m3fn with an absmax/448 per-tensor \code{<key>_scale};
#' everything else is copied through unchanged. The manifest embeds the
#' transformer config, so the source checkpoint is not needed again
#' after quantization.
#'
#' @param transformer_dir Source diffusers transformer directory.
#' @param output_dir Output directory for shards + manifest (default:
#'   the per-format location under \code{tools::R_user_dir}).
#' @param format "nf4" or "fp8".
#' @param shard_bytes Numeric. Target shard size in bytes. The default
#'   1.9e9 keeps every shard under the 2^31-byte (~2.15 GB) ceiling that
#'   stock CRAN safetensors can read, so the artifact loads fork-free.
#'   Pass a larger value (e.g. 4e9) only for local builds you will read
#'   back with a fork-patched safetensors.
#' @param force Logical. Re-quantize even if a valid manifest exists.
#' @param verbose Logical.
#'
#' @return Invisibly, the manifest list.
#'
#' @export
flux_quantize <- function(transformer_dir, output_dir = NULL,
                          format = c("nf4", "fp8"), shard_bytes = 1.9e9,
                          force = FALSE, verbose = TRUE) {
    format <- match.arg(format)
    if (format == "fp8" && !.st_can_write("float8_e4m3fn")) {
        stop("The installed safetensors package cannot write float8 ",
             "tensors (needs the float8 support from ",
             "mlverse/safetensors#13, which is merged upstream but not ",
             "yet on CRAN; install the development version of ",
             "safetensors from GitHub, or use format = \"nf4\").",
             call. = FALSE)
    }
    # Residents load into the compute dtype either way; bf16 halves the
    # artifact but the CRAN build of safetensors 0.2.1 cannot write it
    # (the fix is merged upstream; the probe decides, not the version)
    resident_dtype <- if (.st_can_write("bfloat16")) {
        torch::torch_bfloat16()
    } else {
        message("safetensors cannot write bfloat16; storing resident ",
                "tensors as float32 (larger artifact, same results)")
        torch::torch_float32()
    }
    if (is.null(output_dir)) {
        output_dir <- file.path(tools::R_user_dir("diffuseR", "data"),
                                paste0("flux1-schnell-", format))
    }

    manifest_path <- file.path(output_dir, "manifest.json")
    if (!force && file.exists(manifest_path)) {
        manifest <- jsonlite::fromJSON(manifest_path)
        if (identical(manifest$format, format) &&
            all(file.exists(file.path(output_dir, manifest$shards)))) {
            if (verbose) {
                message(toupper(format), " artifact already present: ",
                        output_dir)
            }
            return(invisible(manifest))
        }
    }
    dir.create(output_dir, recursive = TRUE, showWarnings = FALSE)

    ckpt <- flux_open_checkpoint(transformer_dir)
    hooks <- .flux_family_hooks(ckpt$config)
    if (format == "fp8") {
        fp8 <- torch::torch_float8_e4m3fn()
    }

    shard <- list()
    shard_size <- 0
    shard_files <- character(0)
    n_cast <- 0L

    flush_shard <- function() {
        if (!length(shard)) {
            return()
        }
        fname <- sprintf("%s-%s-%05d.safetensors", hooks$shard_prefix,
                         format, length(shard_files) + 1L)
        safetensors::safe_save_file(shard, file.path(output_dir, fname))
        shard_files[[length(shard_files) + 1L]] <<- fname
        if (verbose) {
            message(sprintf("  wrote %s (%.2f GB, %d tensors)", fname,
                            shard_size / 1e9, length(shard)))
        }
        shard <<- list()
        shard_size <<- 0
        gc(verbose = FALSE)
    }

    keys <- ckpt$keys
    for (i in seq_along(keys)) {
        key <- keys[[i]]
        tensor <- ckpt$handle$get_tensor(key)

        if (hooks$is_quant_key(key)) {
            torch::with_no_grad({
                if (format == "nf4") {
                    q <- ltx23_nf4_quantize(tensor)
                    shard[[key]] <- q$packed
                    shard[[paste0(key, "_absmax")]] <- q$absmax
                    shard_size <- shard_size + prod(tensor$shape) * 0.5625
                } else {
                    scale <- tensor$abs()$max()$to(dtype = torch::torch_float32())$
                    clamp(min = 1e-12) / 448
                    shard[[key]] <- (tensor$to(dtype = torch::torch_float32()) / scale)$to(dtype = fp8)
                    shard[[paste0(key, "_scale")]] <- scale
                    shard_size <- shard_size + prod(tensor$shape)
                }
            })
            n_cast <- n_cast + 1L
        } else {
            # Residents load into the compute dtype anyway; storing bf16
            # keeps fp32-shipped checkpoints (Z-Image) from doubling the
            # artifact. Identity for bf16 sources.
            shard[[key]] <- tensor$to(dtype = resident_dtype)
            shard_size <- shard_size + prod(tensor$shape) * 2
        }
        rm(tensor)

        if (shard_size >= shard_bytes) {
            flush_shard()
        }
        if (i %% 200L == 0L) {
            gc(verbose = FALSE)
            if (verbose) {
                message(sprintf("  quantizing %d/%d tensors", i, length(keys)))
            }
        }
    }
    flush_shard()

    manifest <- list(
                     source = basename(transformer_dir),
                     format = format,
                     shards = shard_files,
                     tensors = length(keys),
                     cast = n_cast,
                     config = ckpt$config
    )
    jsonlite::write_json(manifest, manifest_path, auto_unbox = TRUE,
                         pretty = TRUE)
    if (verbose) {
        message(sprintf("Quantized %d/%d tensors to %s across %d shards: %s",
                        n_cast, length(keys), format, length(shard_files),
                        output_dir))
    }
    invisible(manifest)
}

#' Load a FLUX transformer from any checkpoint format
#'
#' Builds \code{\link{flux_transformer}} from the checkpoint's embedded
#' config and loads the weights. Dispatches on the checkpoint format:
#'
#' \itemize{
#'   \item full precision (\code{\link{flux_open_checkpoint}}): weights
#'     stream into the model in \code{dtype} on \code{device}.
#'   \item \code{"nf4"} (\code{\link{flux_open_quantized}}): cast-set
#'     linears become \code{ltx23_nf4_linear}; the whole model (packed
#'     weights included) moves to \code{device} and stays resident.
#'   \item \code{"fp8"}: cast-set linears become
#'     \code{ltx23_fp8_linear}; fp8 weights stay CPU-resident (optionally
#'     pinned) and stream to \code{device} inside each forward.
#' }
#'
#' @param ckpt A checkpoint from \code{\link{flux_open_checkpoint}} or
#'   \code{\link{flux_open_quantized}}.
#' @param device Character. Compute device.
#' @param dtype Character. Model dtype ("bfloat16" or "float32"). For
#'   quantized formats this sets the resident (non-quantized) tensors
#'   and must match the compute dtype: bfloat16 for GPU compute,
#'   float32 for CPU compute.
#' @param pin Logical. Pin fp8 host memory for faster transfers
#'   (streamed fp8 only).
#' @param fp8_resident Logical. Keep the fp8 weights on \code{device}
#'   instead of streaming from the CPU - right for models whose whole
#'   quantized footprint fits in VRAM (FLUX.2 klein-4B: ~4 GB).
#' @param verbose Logical.
#' @param ... Overrides for \code{\link{flux_transformer}} arguments
#'   (tiny test configs).
#'
#' @return The loaded \code{flux_transformer} in eval mode.
#'
#' @export
flux_load_transformer <- function(ckpt, device = "cuda", dtype = "bfloat16",
                                  pin = TRUE, fp8_resident = FALSE,
                                  verbose = TRUE, ...) {
    stopifnot(inherits(ckpt, "ltx23_checkpoint"))
    format <- ckpt$format %||% "full"
    if (identical(format, "fp8") && !.st_can_write("float8_e4m3fn")) {
        stop("This fp8 artifact needs float8 support the installed ",
             "safetensors lacks (mlverse/safetensors#13, merged upstream ",
             "but not yet on CRAN). Install the development version of ",
             "safetensors from GitHub, or rebuild the artifact as nf4.",
             call. = FALSE)
    }
    hooks <- .flux_family_hooks(ckpt$config)

    args <- utils::modifyList(hooks$args_fn(ckpt$config), list(...))
    model <- do.call(hooks$model_fn, args)

    if (format == "full") {
        model$to(dtype = .flux_dtype(dtype))
        res <- ltx23_load_group(ckpt, ckpt$keys, model, verbose = verbose)
        if (length(res$unmapped) || length(res$unfilled)) {
            stop("FLUX transformer load: ", length(res$unmapped),
                 " unmapped keys, ", length(res$unfilled), " unfilled params")
        }
        model$to(device = device)
        model$eval()
        return(model)
    }

    if (!format %in% c("nf4", "fp8")) {
        stop("Unknown checkpoint format: ", format)
    }
    # Residents (embedders, norms, biases) in the compute dtype: the
    # quantized linears dequantize into the input's dtype at forward, so
    # the two must agree (bfloat16 on GPU, float32 for CPU compute)
    model$to(dtype = .flux_dtype(dtype))

    if (format == "nf4") {
        sib_suffix <- "_absmax"
    } else {
        sib_suffix <- "_scale"
    }
    sib_keys <- ckpt$keys[endsWith(ckpt$keys, paste0(".weight", sib_suffix))]
    main_keys <- setdiff(ckpt$keys, sib_keys)

    dests <- c(model$named_parameters(), model$named_buffers())
    filled <- character(0)
    unmapped <- character(0)

    torch::with_no_grad({
        for (i in seq_along(main_keys)) {
            key <- main_keys[[i]]

            if (hooks$is_quant_key(key) &&
                        paste0(key, sib_suffix) %in% sib_keys) {
                segments <- strsplit(key, ".", fixed = TRUE)[[1]]
                parent <- .ltx23_walk_module(model, utils::head(segments, -2L))
                leaf <- segments[length(segments) - 1L]
                old <- .ltx23_walk_module(parent, leaf)
                if (is.null(old)) {
                    unmapped <- c(unmapped, key)
                    next
                }
                w_shape <- old$weight$shape
                has_bias <- !is.null(old$bias)
                quant_mod <- if (format == "nf4") {
                    ltx23_nf4_linear(w_shape[1], w_shape[2], bias = has_bias)
                } else {
                    ltx23_fp8_linear(w_shape[1], w_shape[2], bias = has_bias)
                }
                if (has_bias) {
                    # Adopt the original bias parameter; its checkpoint key
                    # loads through the pre-swap destination map
                    quant_mod$bias <- old$bias
                }
                if (format == "nf4") {
                    quant_mod$set_nf4_weight(
                        ckpt$handle$get_tensor(key),
                        ckpt$handle$get_tensor(paste0(key, sib_suffix))
                    )
                } else {
                    quant_mod$set_fp8_weight(
                        ckpt$handle$get_tensor(key),
                        ckpt$handle$get_tensor(paste0(key, sib_suffix)),
                        pin = pin
                    )
                }
                do.call(`$<-`, list(parent, leaf, quant_mod))
                filled <- c(filled, key)
            } else {
                dest <- dests[[key]]
                if (is.null(dest)) {
                    unmapped <- c(unmapped, key)
                    next
                }
                dest$copy_(ckpt$handle$get_tensor(key))
                filled <- c(filled, key)
            }

            if (i %% 100L == 0L) {
                gc(verbose = FALSE)
                if (verbose && i %% 500L == 0L) {
                    message(sprintf("  loaded %d/%d transformer tensors", i,
                                    length(main_keys)))
                }
            }
        }
    })
    gc(verbose = FALSE)

    if (length(unmapped)) {
        stop("FLUX ", format, " load: ", length(unmapped),
             " unmapped keys, e.g. ",
             paste(utils::head(unmapped, 3), collapse = ", "))
    }
    # Weight params replaced by quantized modules won't be "filled"
    expected_missing <- hooks$is_quant_key(names(dests))
    unfilled <- setdiff(names(dests)[!expected_missing], filled)
    if (length(unfilled)) {
        stop("FLUX ", format, " load: ", length(unfilled),
             " unfilled params, e.g. ",
             paste(utils::head(unfilled, 3), collapse = ", "))
    }

    # NF4: everything (packed buffers included) onto the device.
    # FP8: residents move; the plain-field fp8 weights stay CPU-side
    # unless fp8_resident moves them (small models that fit on the GPU).
    model$to(device = device)
    if (format == "fp8" && fp8_resident) {
        .flux_fp8_to_device(model, device)
    }
    model$eval()
    # Block intermediates are large at image resolutions; per-block gc
    # keeps the quantized-linear temporaries bounded
    options(diffuseR.block_gc = TRUE)
    if (verbose) {
        message("FLUX transformer ready (", format, ") on ", device)
    }
    model
}
