"""One-time startup: env/GPU diagnostics, the fp8 just-in-time-upcast monkeypatch, and the
`load_transformer`/`load_pipeline` functions callers use to build the transformer + pipeline.
Importing this module runs the diagnostics/TF32 setup and leaves `device` and `dtype` as
module attributes; it does NOT construct the transformer or pipeline itself — the caller
(inference.py) does that explicitly so it's clear `pipe` is a constructed object, not a
module-level singleton.

load_transformer/load_pipeline load directly onto 'cuda' instead of CPU, per HF's documented
ZeroGPU pattern (https://huggingface.co/docs/hub/spaces-zerogpu#model-loading). This is only
safe because `spaces` is imported (and patches torch for its CUDA-emulation mode) before any
CUDA-touching code below runs."""
import os
import threading
import time
import types

import spaces  # noqa: F401  (must be imported before any CUDA-touching code below)
import torch
from diffusers.models.normalization import RMSNorm
from transformers import Qwen2_5_VLForConditionalGeneration

from logging_utils import (
    print_cuda_device_count,
    print_cuda_visible_devices,
    print_env_cuda_version,
    print_env_cudnn_version,
    print_env_gpu,
    print_env_package_version,
    print_env_package_version_unavailable,
    print_env_ram,
    print_env_ram_unavailable,
    print_heartbeat,
    print_int8_text_encoder_loaded,
    print_lm_head_dropped,
    print_loading_int8_text_encoder,
    print_loading_pipeline,
    print_loading_transformer,
    print_pipeline_loaded,
    print_tf32_enabled,
    print_torch_version,
    print_transformer_loaded,
    print_transformer_memory_footprint,
    print_transformer_memory_footprint_unavailable,
    print_transformer_patched,
    print_unpatched_fp8_param,
    print_using_device,
    print_vae_tiling,
)
from mode import Mode
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16

print_cuda_visible_devices()
print_torch_version()
print_using_device(device)
print_cuda_device_count()


def _log_env() -> None:
    import importlib.metadata as _meta
    if torch.cuda.is_available():
        print_env_gpu(torch.cuda.get_device_properties(0))
    print_env_cuda_version()
    print_env_cudnn_version()
    for pkg in ["spaces", "diffusers", "transformers", "gradio", "accelerate", "peft", "torchvision"]:
        try:
            print_env_package_version(pkg, _meta.version(pkg))
        except Exception as e:
            print_env_package_version_unavailable(pkg, e)
    try:
        mem: dict[str, str] = {}
        with open("/proc/meminfo") as f:
            for line in f:
                k, v = line.split(":", 1)
                mem[k.strip()] = v.strip()
        total_gb = int(mem["MemTotal"].split()[0]) / 1024**2
        avail_gb = int(mem["MemAvailable"].split()[0]) / 1024**2
        print_env_ram(total_gb, avail_gb)
    except Exception as e:
        print_env_ram_unavailable(e)


_log_env()

# TF32 matmul: ~10-15% free speedup on Ampere/Hopper (bfloat16 accumulation paths benefit too)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
print_tf32_enabled()


def _start_heartbeat(label: str) -> threading.Event:
    done = threading.Event()
    t0 = time.perf_counter()
    def _beat() -> None:
        while not done.wait(timeout=15):
            print_heartbeat(label, time.perf_counter() - t0)
    threading.Thread(target=_beat, daemon=True).start()
    return done


_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)


def _fp8_upcast_linear_forward(self: torch.nn.Linear, input: torch.Tensor) -> torch.Tensor:
    weight = self.weight.to(input.dtype) if self.weight.dtype in _FP8_DTYPES else self.weight
    bias = self.bias.to(input.dtype) if (self.bias is not None and self.bias.dtype in _FP8_DTYPES) else self.bias
    return torch.nn.functional.linear(input, weight, bias)


def _fp8_upcast_rmsnorm_forward(self: RMSNorm, hidden_states: torch.Tensor) -> torch.Tensor:
    # Mirrors diffusers 0.39.0's RMSNorm.forward (CUDA path, models/normalization.py), extended
    # so an fp8-resident weight/bias gets upcast to the activation's dtype before use instead of
    # being silently skipped — the stock implementation only special-cases float16/bfloat16, so
    # an fp8 weight would otherwise reach `hidden_states * self.weight` unconverted and error
    # (no elementwise op supports bf16 x fp8 operands).
    input_dtype = hidden_states.dtype
    variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + self.eps)

    if self.weight is not None:
        weight = self.weight.to(input_dtype) if self.weight.dtype in _FP8_DTYPES else self.weight
        if weight.dtype in (torch.float16, torch.bfloat16):
            hidden_states = hidden_states.to(weight.dtype)
        hidden_states = hidden_states * weight
        if self.bias is not None:
            bias = self.bias.to(hidden_states.dtype) if self.bias.dtype in _FP8_DTYPES else self.bias
            hidden_states = hidden_states + bias
    else:
        hidden_states = hidden_states.to(input_dtype)

    return hidden_states


def _patch_fp8_modules(model: torch.nn.Module) -> int:
    # This checkpoint ships its weights natively in fp8 (torch_dtype below preserves that
    # instead of upcasting to bf16 at load time, halving resident memory: ~19GB vs ~38GB).
    # Neither nn.Linear nor RMSNorm (the two module types in this model that own their own
    # weight/bias, per the checkpoint's safetensors headers — every tensor is fp8, including
    # norm gains) have an fp8 compute kernel on this GPU, so each patched instance upcasts its
    # own weight to the input's dtype just-in-time for the op — mathematically identical to the
    # old load-time-upcast-everything approach (same values, same target dtype), just deferred
    # so only one layer's weight is transiently bf16 at a time instead of all of them.
    count = 0
    for module in model.modules():
        if isinstance(module, torch.nn.Linear) and module.weight.dtype in _FP8_DTYPES:
            module.forward = types.MethodType(_fp8_upcast_linear_forward, module)  # type: ignore[method-assign]
            count += 1
        elif isinstance(module, RMSNorm) and module.weight is not None and module.weight.dtype in _FP8_DTYPES:
            module.forward = types.MethodType(_fp8_upcast_rmsnorm_forward, module)  # type: ignore[method-assign]
            count += 1

    # Safety net: flag any other fp8-resident parameter that wasn't patched above, so a gap in
    # this allowlist surfaces as a startup log line instead of a mid-inference crash — an
    # unpatched fp8 parameter can't participate in ops with the bf16 activations around it.
    patched_types = (torch.nn.Linear, RMSNorm)
    for name, module in model.named_modules():
        if isinstance(module, patched_types):
            continue
        for pname, param in module.named_parameters(recurse=False):
            if param.dtype in _FP8_DTYPES:
                print_unpatched_fp8_param(name, pname, type(module).__name__)
    return count


def load_transformer() -> QwenImageTransformer2DModel:
    t0 = time.perf_counter()
    print_loading_transformer()
    hb = _start_heartbeat("transformer")
    transformer: QwenImageTransformer2DModel = QwenImageTransformer2DModel.from_pretrained(
        "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V23",
        torch_dtype=torch.float8_e4m3fn,
        device_map={"": device},
    )
    hb.set()
    print_transformer_loaded(time.perf_counter() - t0)
    n_patched = _patch_fp8_modules(transformer)
    print_transformer_patched(n_patched)
    try:
        print_transformer_memory_footprint(transformer.get_memory_footprint() / 1024**3)  # type: ignore[no-untyped-call]
    except Exception as e:
        print_transformer_memory_footprint_unavailable(e)
    return transformer


def drop_lm_head(text_encoder: torch.nn.Module, tag: str) -> None:
    # encode_prompt() (pipeline_qwenimage_edit_plus.py) only ever reads outputs.hidden_states,
    # never outputs.logits — but Qwen2_5_VLForConditionalGeneration.forward() unconditionally
    # runs hidden_states through lm_head regardless (its `logits_to_keep=0` default means "keep
    # all", not "skip"). That's a dead ~152064x3584 vocab projection: ~1GB resident weight in
    # bf16 (~0.5GB int8) plus a wasted full-vocab matmul on every prompt encode. Identity() keeps
    # forward()'s `self.lm_head(hidden_states)` call working (hidden_states pass through
    # unchanged) while dropping both the weight and the matmul.
    text_encoder.lm_head = torch.nn.Identity()
    print_lm_head_dropped(tag)


def load_int8_text_encoder(repo_id: str) -> Qwen2_5_VLForConditionalGeneration:
    # Loaded directly onto `device` at startup, same call (device_map={"": device}) that used
    # to run per-request inside inference.py's @spaces.GPU-decorated call. This is a bet that
    # ZeroGPU's module-level-.to('cuda') capture (see load_transformer above) also covers a
    # pre-quantized bitsandbytes int8 checkpoint the way it does plain bf16/fp8 tensors:
    # bitsandbytes>=0.48 (pinned here: 0.49.2) added .to()-device-move support for an
    # already-quantized Int8Params (a plain data+CB/SCB move, no requantize — see
    # bitsandbytes/nn/modules.py's Int8Params.to()), which is what transformers' own
    # device-move guard checks for. UNVERIFIED on a real GPU worker as of this commit — if
    # ZeroGPU's tensor-pack capture doesn't handle bnb's custom Tensor subclass the same way
    # it handles ordinary tensors, this needs to move back to a live per-request load inside
    # the @spaces.GPU call (as it was before).
    t0 = time.perf_counter()
    print_loading_int8_text_encoder(repo_id)
    hb = _start_heartbeat("int8_text_encoder")
    text_encoder: Qwen2_5_VLForConditionalGeneration = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        repo_id,
        device_map={"": device},
        dtype=torch.bfloat16,
    )
    hb.set()
    drop_lm_head(text_encoder, "startup")
    print_int8_text_encoder_loaded(time.perf_counter() - t0)
    return text_encoder


def load_pipeline(transformer: QwenImageTransformer2DModel) -> QwenImageEditPlusPipeline:
    t0 = time.perf_counter()
    print_loading_pipeline()
    hb = _start_heartbeat("pipeline")
    pipeline = QwenImageEditPlusPipeline.from_pretrained(
        "FireRedTeam/FireRed-Image-Edit-1.1",
        transformer=transformer,
        torch_dtype=dtype,
    )
    drop_lm_head(pipeline.text_encoder, "startup")
    # text_encoder deliberately NOT moved to `device` here (stays CPU-resident, unused).
    # inference.py replaces `pipe.text_encoder` outright with a separately-loaded int8
    # checkpoint (see load_int8_text_encoder above) right after load_pipeline() returns, so
    # this bf16 copy is dead weight from the moment the pipeline is constructed. ZeroGPU's
    # module-level-.to('cuda') emulation registers every moved tensor in its own internal
    # `tensor_packs` bookkeeping, which holds a permanent reference for the life of the
    # (forked) worker regardless of what our own code does with `pipe.text_encoder`
    # afterwards — so moving this bf16 copy to CUDA here would have ZeroGPU pin ~15GB of
    # weights that get replaced before a single request runs. Verified via OOM in an earlier
    # experiment that moved it anyway (see #24): alloc=43.46GB (~19GB transformer + ~15.4GB
    # stranded bf16 text_encoder + ~8.75GB int8 text_encoder) against a 47.4GB MIG slice.
    pipeline.vae.to(device)
    hb.set()
    pipeline.vae.enable_tiling(
        tile_sample_min_height=Mode.HIGH_DETAIL.max_dim, tile_sample_min_width=Mode.HIGH_DETAIL.max_dim
    )
    print_vae_tiling(pipeline.vae.tile_sample_min_height, pipeline.vae.tile_sample_min_width, pipeline.vae.use_tiling)
    print_pipeline_loaded(time.perf_counter() - t0)
    return pipeline