"""Per-request inference: input validation/prep, GPU placement, the diffusion call itself,
and timing/memory instrumentation around it."""
import gc
import os
import random
import threading
import time
import traceback

from typing import Any

import gradio as gr
import numpy as np
import spaces
import torch
from PIL.Image import Image as PILImage

from dimensions import compute_output_dimensions
from image_codec import b64_to_pil_list
from logging_utils import (
    LogUploader,
    print_attn_processor_set,
    print_calling_pipe,
    print_cuda_sync_after_error,
    print_first_call_into_module,
    print_gpu_mem_status,
    print_images_predecoded,
    print_infer_end,
    print_infer_error,
    print_infer_exception,
    print_infer_gpu_properties,
    print_infer_params,
    print_infer_prompt,
    print_infer_start_header,
    print_infer_traceback,
    print_pre_vae_decode,
    print_setting_attn_processor,
    print_step_done,
    print_text_encoder_offload_skipped_int8,
    print_vae_decode_done,
    print_vae_tiling_activation,
)
from mode import Mode
from model_loading import device, load_int8_text_encoder, load_pipeline, load_transformer
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
from timing import Timer

MAX_SEED = np.iinfo(np.int32).max

_transformer = load_transformer()
pipe = load_pipeline(_transformer)

# int8 (bitsandbytes) quantized text_encoder, produced offline from the FireRed
# text_encoder's bf16 weights (~8.75GB vs ~15.4GB for the bf16 original). Repo id comes
# from a secret rather than being hardcoded here. Loaded at startup (module level) and
# .to(device)-captured by ZeroGPU the same way transformer/vae are, instead of being
# reloaded on every @spaces.GPU call — see load_int8_text_encoder's docstring/comment in
# model_loading.py for why this is unverified on a real GPU as of this commit.
_TEXT_ENCODER_INT8_REPO = os.environ.get("TEXT_ENCODER_INT8_REPO")
if not _TEXT_ENCODER_INT8_REPO:
    raise RuntimeError("TEXT_ENCODER_INT8_REPO is not set; the int8 text_encoder is required.")
pipe.text_encoder = load_int8_text_encoder(_TEXT_ENCODER_INT8_REPO)

print_setting_attn_processor()
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
print_attn_processor_set()

with open("static/negative_prompt.txt") as _f:
    negative_prompt = _f.read().strip()

_log_uploader = LogUploader(
    token=os.environ.get("HF_TOKEN"),
    repo_id=os.environ.get("LOG_DATASET_REPO"),
    max_files=int(os.environ.get("LOG_MAX_FILES", "5000")),
    batch_interval=int(os.environ.get("LOG_BATCH_INTERVAL", "60")),
)


# (label, start_mark, end_mark) rows for Timer.print_report(), matching the marks
# _infer_gpu/_make_step_callback record: pipe_start, first_step, last_step, pipe_end.
_TIMING_ROWS: list[tuple[str, str, str]] = [
    ("preprocess", "pipe_start", "first_step"),
    ("inference",  "first_step", "last_step"),
    ("vae_decode", "last_step",  "pipe_end"),
]


def _gpu_mem_str(cuda_ok: bool, sync: bool = False) -> str:
    if not cuda_ok:
        return "CUDA not available"
    if sync:
        try:
            torch.cuda.synchronize()
        except Exception as se:
            return f"CUDA sync failed: {se}"
    alloc    = torch.cuda.memory_allocated()    / 1024**3
    reserved = torch.cuda.memory_reserved()     / 1024**3
    peak     = torch.cuda.max_memory_allocated() / 1024**3
    return f"alloc={alloc:.2f}GB reserved={reserved:.2f}GB peak={peak:.2f}GB"


def _validate_infer_inputs(pil_images: list[PILImage], prompt: str) -> None:
    if not pil_images:
        raise gr.Error("Please upload at least one image to edit.")
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter an edit prompt.")


def _resolve_seed(seed: int, randomize_seed: bool) -> int:
    return random.randint(0, MAX_SEED) if randomize_seed else seed


def _spawn_log(pil_images: list[PILImage], result_image: PILImage | None, prompt: str, seed: int,
               steps: int, guidance_scale: float, width: int, height: int, duration: float,
               success: bool, error: str = "") -> None:
    threading.Thread(
        target=_log_uploader.log_inference,
        args=(pil_images, result_image, prompt, seed, steps, guidance_scale,
              width, height, duration, success, error),
        daemon=True,
    ).start()


def update_dimensions_on_upload(image: PILImage | None, max_dim: int) -> tuple[int, int]:
    if image is None:
        return max_dim, max_dim
    w, h = image.size
    return compute_output_dimensions(w, h, max_dim)


def infer(images_b64_json: str, prompt: str, seed: int, randomize_seed: bool, guidance_scale: float,
          steps: int, mode: object, gpu_duration: int = 30,
          progress: gr.Progress = gr.Progress(track_tqdm=True)) -> tuple[PILImage, int]:
    # CPU-only preprocessing — GPU not yet allocated
    gc.collect()
    mode = Mode.from_value(mode)
    pil_images = b64_to_pil_list(images_b64_json)
    _validate_infer_inputs(pil_images, prompt)
    seed = _resolve_seed(seed, randomize_seed)
    width, height = update_dimensions_on_upload(pil_images[0], mode.max_dim)
    t0 = time.perf_counter()
    try:
        result_image, seed, duration = _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, mode, int(gpu_duration))
        # _spawn_log is called here (main process) so the thread survives after _infer_gpu's
        # @spaces.GPU subprocess exits — previously the daemon thread was killed on subprocess exit.
        _spawn_log(pil_images, result_image, prompt, seed, steps, guidance_scale, width, height, duration, True)
        return result_image, seed
    except Exception as e:
        duration = time.perf_counter() - t0
        # Diagnosing "Could not parse server response. Syntax error '<'" client-side errors —
        # that means the browser got HTML instead of JSON from the SSE stream, which points to
        # Gradio failing to serialize this exception rather than the exception itself. Logging
        # the concrete type/module here (not just str(e)) so we can tell whether it's a plain
        # Exception, a gr.Error, or something from the `spaces` package with non-standard attrs.
        print_infer_exception(e)
        traceback.print_exc()
        _spawn_log(pil_images, None, prompt, seed, steps, guidance_scale, width, height, duration, False, str(e))
        raise


def _log_infer_start(prompt: str, steps: int, guidance_scale: float, seed: int, gpu_duration: int, mode: Mode) -> None:
    print_infer_start_header()
    print_infer_params(steps, guidance_scale, seed, gpu_duration, mode)
    print_infer_prompt(prompt)


def _log_gpu_properties(cuda_ok: bool) -> torch.cuda._CudaDeviceProperties | None:
    if not cuda_ok:
        return None
    p = torch.cuda.get_device_properties(0)
    print_infer_gpu_properties(p)
    torch.cuda.reset_peak_memory_stats()
    return p


def _instrument_first_touch(modules_with_names: list[tuple[torch.nn.Module, str]], t0: float) -> None:
    """Install self-removing forward-pre-hooks that log the moment each module is first entered."""
    def _make_hook(name: str, handle_box: dict[str, torch.utils.hooks.RemovableHandle]) -> Any:
        def _hook(mod: torch.nn.Module, inputs: tuple[Any, ...]) -> None:
            print_first_call_into_module(name, _gpu_mem_str(True, sync=True), time.perf_counter() - t0)
            handle_box["h"].remove()
        return _hook
    for module, name in modules_with_names:
        handle_box: dict[str, torch.utils.hooks.RemovableHandle] = {}
        handle_box["h"] = module.register_forward_pre_hook(_make_hook(name, handle_box))


def _make_step_callback(steps: int, timer: Timer, t0: float, mode: Mode, cuda_ok: bool = False) -> Any:
    """Build the diffusers step callback that logs per-step timing and marks timer checkpoints."""
    step_times: list[float] = []
    def _step_cb(pipeline: Any, step_idx: int, timestep: Any, cb_kwargs: dict[str, Any]) -> dict[str, Any]:
        now = time.perf_counter()
        step_times.append(now)
        if step_idx == 0:
            timer.mark("first_step")
        timer.mark("last_step")  # overwritten each step; final value = end of last step
        delta_ms = (now - (step_times[-2] if len(step_times) > 1 else t0)) * 1000
        tag = " ← includes cold-start (offload hook install + first weight transfer)" if step_idx == 0 else ""
        print_step_done(step_idx, steps, delta_ms, tag, now - t0)
        # text_encoder is always int8 (bitsandbytes) quantized now. It's never evicted ahead of
        # VAE decode's fp32-upcast memory spike: bitsandbytes>=0.48 (pinned: 0.49.2) does
        # support .to()-moving an already-quantized Int8Params (see load_int8_text_encoder in
        # model_loading.py), but this codepath isn't wired up to use that — its ~8.75GB
        # footprint needs the safety net less than the old ~15GB bf16 text_encoder did anyway.
        # Also skipped when accelerate hooks are managing placement (offload-fallback path), to
        # avoid fighting their own device bookkeeping.
        if step_idx == steps - 1 and getattr(pipeline.text_encoder, "_hf_hook", None) is None:
            print_text_encoder_offload_skipped_int8()
        if step_idx == steps - 1 and cuda_ok:
            print_pre_vae_decode(_gpu_mem_str(True, sync=True), time.perf_counter() - t0)
            torch.cuda.reset_peak_memory_stats()
        return cb_kwargs
    return _step_cb


def _log_infer_error(e: Exception, t0: float, timer: Timer) -> None:
    print_infer_error(e, time.perf_counter() - t0)
    print_infer_traceback()
    try:
        torch.cuda.synchronize()
    except Exception as cuda_err:
        print_cuda_sync_after_error(cuda_err)
    timer.print_report(_TIMING_ROWS, total=("pipe_start", "pipe_end"))


# gpu_duration is declared to @spaces.GPU as-is, with no added buffer. There used to be a
# _COLD_START_BUFFER_S padding this value, sized to cover a previously-measured worst case
# ~30s int8 text_encoder reload (HF hub/disk cache miss) plus general slack for cold-start
# overhead and gpu_duration under-estimating real diffusion+decode time. #43 eliminated that
# per-request reload (text_encoder now loads once at startup, see load_int8_text_encoder in
# model_loading.py), and the gpu_duration slider's range (10-120s in the UI) already reaches
# _MAX_GPU_DURATION_S on its own — so instead of guessing a hidden buffer value from limited
# data, the UI's default gpu_duration values were raised to what the old buffered totals used
# to be (see mode_toggle.js's MODE_GPU_DURATION), and users can now dial it themselves.
# It doesn't cost extra quota either way (billing is by real usage, not the declared duration
# — see huggingface.co/docs/hub/spaces-zerogpu), but declaring too little makes ZeroGPU kill
# the call mid-run with "GPU task aborted".
_MAX_GPU_DURATION_S = 120  # matches the gpu_duration slider's max in the UI


@spaces.GPU(duration=lambda *a, **kw: min(int(a[8]), _MAX_GPU_DURATION_S) if len(a) > 8 else 60)  # type: ignore[untyped-decorator]
def _infer_gpu(pil_images: list[PILImage], prompt: str, seed: int, guidance_scale: float, steps: int,
               width: int, height: int, mode: Mode, gpu_duration: int = 30) -> tuple[PILImage, int, float]:
    # `pipe` is deliberately read here as the module-level global, NOT taken as a parameter of
    # this function: @spaces.GPU marshals every argument of the decorated function through a
    # multiprocessing.Queue to a persistent forked worker (spaces/zero/wrappers.py's
    # thread_wrapper), which means it gets pickled. pipe.transformer's fp8-upcast monkeypatch
    # (model_loading.py's _patch_fp8_modules, `module.forward = types.MethodType(fn, module)`)
    # isn't pickle-safe — the bound method's __name__ doesn't match the `forward` attribute
    # it's stored under, so reconstructing it via getattr(module, fn.__name__) raises
    # AttributeError in the worker ("GPU task aborted"). See #44.
    _cuda_ok = torch.cuda.is_available()
    timer = Timer(_cuda_ok)
    t0 = time.perf_counter()

    _log_infer_start(prompt, steps, guidance_scale, seed, gpu_duration, mode)
    _log_gpu_properties(_cuda_ok)

    print_gpu_mem_status(_gpu_mem_str(_cuda_ok), time.perf_counter() - t0)

    if _cuda_ok:
        _instrument_first_touch(
            [(pipe.text_encoder, "text_encoder"), (pipe.transformer, "transformer"), (pipe.vae, "vae")],
            t0,
        )

    print_images_predecoded(len(pil_images), width, height, seed)
    if _cuda_ok:
        _will_tile = pipe.vae.use_tiling and (
            width > pipe.vae.tile_sample_min_width or height > pipe.vae.tile_sample_min_height
        )
        print_vae_tiling_activation(_will_tile, pipe.vae.tile_sample_min_height, pipe.vae.tile_sample_min_width)

    generator = torch.Generator(device=device).manual_seed(seed)
    step_cb = _make_step_callback(steps, timer, t0, mode, _cuda_ok)

    timer.mark("pipe_start")
    print_calling_pipe(time.perf_counter() - t0)
    try:
        result_image = pipe(
            image=pil_images, prompt=prompt, negative_prompt=negative_prompt,
            height=height, width=width, num_inference_steps=steps,
            generator=generator, true_cfg_scale=guidance_scale,
            callback_on_step_end=step_cb,
            callback_on_step_end_tensor_inputs=["latents"],
        ).images[0]
        timer.mark("pipe_end")
        print_vae_decode_done(_gpu_mem_str(_cuda_ok, sync=True), time.perf_counter() - t0)
        timer.print_report(_TIMING_ROWS, total=("pipe_start", "pipe_end"))
        duration = timer.elapsed_ms("pipe_start", "pipe_end") / 1000.0
        return result_image, seed, duration
    except Exception as e:
        _log_infer_error(e, t0, timer)
        raise
    finally:
        # No manual pipe.to("cpu"): in the offload-fallback case that fights the
        # hooks' own device bookkeeping (they return each component to CPU after
        # its forward), and in the normal fast-path case the worker's GPU access
        # is reclaimed by ZeroGPU when this call returns regardless — paying for
        # a D2H transfer here would just be wasted GPU-billed time.
        gc.collect()
        torch.cuda.empty_cache()
        print_infer_end(time.perf_counter() - t0)