#' FLUX Checkpoint Readers
#'
#' FLUX transformers ship in the diffusers layout: a directory with
#' \code{config.json}, one or more \code{diffusion_pytorch_model*.safetensors}
#' shards, and (when sharded) a
#' \code{diffusion_pytorch_model.safetensors.index.json} weight map.
#' These helpers open that layout behind the same checkpoint interface as
#' \code{\link{ltx23_open_checkpoint}}, so the LTX group loaders and
#' quantization machinery work unchanged. FLUX module names mirror the
#' checkpoint keys 1:1 - no key mapping is needed.
#'
#' @name checkpoint_flux
NULL

# Quantization cast set: every large linear weight in the FLUX blocks,
# including the adaLN modulation linears (norm*.linear are 3.2B params
# across the model; leaving them bf16 would not fit resident on 16 GB).
# Everything else (embedders, q/k norms, final norm, biases) stays bf16.
# Full-size census: 19 double blocks x 14 + 38 single blocks x 6 = 494
# cast weights, ~11.8B of the 12B parameters.
.flux_quant_cast_pattern <- paste0(
                                   "^(",
                                   "transformer_blocks\\.[0-9]+\\.(",
                                   "attn\\.(to_q|to_k|to_v|add_q_proj|add_k_proj|add_v_proj|to_out\\.0|to_add_out)",
                                   "|ff\\.net\\.(0\\.proj|2)|ff_context\\.net\\.(0\\.proj|2)",
                                   "|norm1\\.linear|norm1_context\\.linear",
                                   ")",
                                   "|single_transformer_blocks\\.[0-9]+\\.(",
                                   "attn\\.(to_q|to_k|to_v)|proj_mlp|proj_out|norm\\.linear",
                                   ")",
                                   ")\\.weight$"
)

#' Test whether a FLUX key is in the quantization cast set
#'
#' @param key Character vector of parameter names (diffusers-style).
#'
#' @return Logical vector.
#'
#' @export
flux_is_quant_key <- function(key) {
    grepl(.flux_quant_cast_pattern, key)
}

# FLUX.2 cast set: block linears, the three shared modulation
# projections, and the context embedder (7680 x 3072). Everything else
# (x_embedder, timestep MLP, norm_out, q/k norms) stays in the resident
# dtype. Full klein-4B census: 5 double x 12 + 20 single x 2 + 3
# modulations + context_embedder = 104 cast weights, ~3.9B of 4B params.
.flux2_quant_cast_pattern <- paste0(
                                    "^(",
                                    "transformer_blocks\\.[0-9]+\\.(",
                                    "attn\\.(to_q|to_k|to_v|add_q_proj|add_k_proj|add_v_proj|to_out\\.0|to_add_out)",
                                    "|ff\\.(linear_in|linear_out)|ff_context\\.(linear_in|linear_out)",
                                    ")",
                                    "|single_transformer_blocks\\.[0-9]+\\.attn\\.(to_qkv_mlp_proj|to_out)",
                                    "|double_stream_modulation_img\\.linear",
                                    "|double_stream_modulation_txt\\.linear",
                                    "|single_stream_modulation\\.linear",
                                    "|context_embedder",
                                    ")\\.weight$"
)

# Z-Image cast set: the seven core linears (attention q/k/v/out and the
# three SwiGLU weights) in all three block stacks. The per-block adaLN
# modulation linears, embedders, pad tokens and norms stay in the
# resident dtype. Full Turbo census: (30 layers + 2 noise_refiner + 2
# context_refiner) x 7 = 238 cast weights, ~6.0B of the 6B parameters.
.zimage_quant_cast_pattern <- paste0(
                                     "^(noise_refiner|context_refiner|layers)\\.[0-9]+\\.(",
                                     "attention\\.(to_q|to_k|to_v|to_out\\.0)",
                                     "|feed_forward\\.(w1|w2|w3)",
                                     ")\\.weight$"
)

#' Test whether a Z-Image key is in the quantization cast set
#'
#' @param key Character vector of parameter names (diffusers-style).
#'
#' @return Logical vector.
#'
#' @export
zimage_is_quant_key <- function(key) {
    grepl(.zimage_quant_cast_pattern, key)
}

#' Test whether a FLUX.2 key is in the quantization cast set
#'
#' @param key Character vector of parameter names (diffusers-style).
#'
#' @return Logical vector.
#'
#' @export
flux2_is_quant_key <- function(key) {
    grepl(.flux2_quant_cast_pattern, key)
}

# Model family from a diffusers transformer config
.flux_family <- function(config) {
    cls <- config$`_class_name` %||% "FluxTransformer2DModel"
    if (identical(cls, "Flux2Transformer2DModel")) {
        "flux2"
    } else if (identical(cls, "ZImageTransformer2DModel")) {
        "zimage"
    } else {
        "flux1"
    }
}

#' Open a FLUX transformer checkpoint directory
#'
#' Opens a diffusers-layout transformer directory lazily (headers only).
#' Sharded checkpoints are resolved through the index.json weight map;
#' single-file checkpoints are opened directly. The transformer
#' \code{config.json} is attached as \code{$config}.
#'
#' @param transformer_dir Directory containing \code{config.json} and the
#'   \code{diffusion_pytorch_model*.safetensors} file(s).
#'
#' @return An object of class \code{ltx23_checkpoint} (shared checkpoint
#'   interface): list with \code{handle$get_tensor}, \code{keys},
#'   \code{config}, and \code{path}.
#'
#' @export
flux_open_checkpoint <- function(transformer_dir) {
    if (!requireNamespace("safetensors", quietly = TRUE)) {
        stop("The safetensors package is required to read FLUX checkpoints.")
    }
    transformer_dir <- path.expand(transformer_dir)
    if (!dir.exists(transformer_dir)) {
        stop("Checkpoint directory not found: ", transformer_dir)
    }

    config <- NULL
    config_path <- file.path(transformer_dir, "config.json")
    if (file.exists(config_path)) {
        config <- jsonlite::fromJSON(config_path, simplifyVector = TRUE)
    }

    opened <- .flux_open_sharded_dir(transformer_dir, "diffusion_pytorch_model")

    structure(
              list(handle = opened$handle, keys = opened$keys, version = NULL,
                   config = config, path = transformer_dir),
              class = "ltx23_checkpoint"
    )
}

# Open a HF-layout safetensors directory (sharded via <base>.safetensors
# .index.json, or a single <base>.safetensors) lazily; returns the
# handle/keys pair shared by all checkpoint objects.
.flux_open_sharded_dir <- function(dir, base) {
    index_path <- file.path(dir, paste0(base, ".safetensors.index.json"))
    if (file.exists(index_path)) {
        index <- jsonlite::fromJSON(index_path, simplifyVector = TRUE)
        weight_map <- unlist(index$weight_map)
        shard_files <- unique(weight_map)
        missing <- shard_files[!file.exists(file.path(dir, shard_files))]
        if (length(missing)) {
            stop("Missing checkpoint shards: ", paste(missing, collapse = ", "))
        }
        handles <- lapply(file.path(dir, shard_files), function(p) {
            safetensors::safetensors$new(p, framework = "torch")
        })
        names(handles) <- shard_files
        keys <- names(weight_map)
        handle <- list(
                       get_tensor = function(key) {
            shard <- weight_map[[key]]
            if (is.null(shard) || is.na(shard)) {
                stop("Key not found in checkpoint index: ", key)
            }
            .st_read_or_breadcrumb(function() handles[[shard]]$get_tensor(key),
                                   file.path(dir, shard))
        }
        )
    } else {
        single_path <- file.path(dir, paste0(base, ".safetensors"))
        if (!file.exists(single_path)) {
            stop("No ", base, " safetensors (or index) in ", dir)
        }
        h <- safetensors::safetensors$new(single_path, framework = "torch")
        keys <- setdiff(h$keys(), "__metadata__")
        handle <- list(get_tensor = function(key) {
            .st_read_or_breadcrumb(function() h$get_tensor(key), single_path)
        })
    }
    list(handle = handle, keys = keys)
}

#' Open a quantized FLUX artifact directory
#'
#' Opens the sharded NF4/fp8 artifact written by
#' \code{\link{flux_quantize}} through the shared checkpoint interface.
#' The manifest's embedded transformer config and \code{format} ride
#' along, so \code{\link{flux_load_transformer}} needs nothing else.
#'
#' @param dir The quantized artifact directory (with manifest.json).
#'
#' @return An \code{ltx23_checkpoint} with \code{$format} set.
#'
#' @export
flux_open_quantized <- function(dir) {
    manifest_path <- file.path(dir, "manifest.json")
    if (!file.exists(manifest_path)) {
        stop("No manifest.json in ", dir, "; run flux_quantize() first.")
    }
    ltx23_open_fp8_checkpoint(dir)
}
