"""Turbo workflow verified on local CUDA and Hugging Face ZeroGPU large."""
import os
import time

import spaces
import torch
from diffusers import QwenImage21Pipeline
from huggingface_hub import hf_hub_download

from core import requested_seconds, run_batch, rgb_on_white
from angles_adapter import load_angles_adapter
from views import REFERENCE_RESOLUTION

MODEL_ID = "Qwen/Qwen-Image-2.1-Turbo"
MODEL_REVISION = "d65dbc9a7e8f6b5479e33dee6030eaab2a906509"
GPU_SIZE = os.getenv("GPU_SIZE", "large")
if GPU_SIZE not in ("large", "xlarge"):
    raise ValueError("GPU_SIZE must be large or xlarge")
QUOTA_MULTIPLIER = 2 if GPU_SIZE == "xlarge" else 1
MAX_GPU_SECONDS = int(os.getenv("MAX_GPU_SECONDS", "110"))
if not 45 <= MAX_GPU_SECONDS <= 110:
    raise ValueError("MAX_GPU_SECONDS must be between 45 and 110")
STARTUP_RESERVE_SECONDS = 15
USE_ANGLES = os.getenv("ENABLE_ANGLES_LORA", "0") == "1"
ANGLES_REPO = "akhaliq/Qwen-Image-2.1-Multiple-Angles-LoRA"
ANGLES_WEIGHT = ("checkpoints_v2/angles_v2_full_qwen21_multiple_angles_v2_"
                 "qwen21_multiple_angles_v2_000001500.safetensors")

# ZeroGPU requires startup CUDA placement, not moving the model for each request.
# Do not reuse the old 2511 Lightning or angles adapters on this architecture.
pipe = QwenImage21Pipeline.from_pretrained(
    MODEL_ID, revision=MODEL_REVISION, dtype=torch.bfloat16,
)
if USE_ANGLES:
    checkpoint = hf_hub_download(ANGLES_REPO, filename=ANGLES_WEIGHT)
    print(load_angles_adapter(pipe, checkpoint, adapter_name="angles", strength=0.8))
    pipe.fuse_lora(adapter_names=["angles"])
    pipe.unload_lora_weights()
pipe.vae.enable_tiling()
pipe.to("cuda")
if len(pipe.config.get("sample_sigmas", [])) != 8:
    raise RuntimeError("Turboの8-stepスケジュールを読み込めていません。Diffusersのバージョンを確認してください。")


def duration(reference, keys, guides, resolution, seed, extra, crop_close_views=False):
    return requested_seconds(keys, MAX_GPU_SECONDS)


@spaces.GPU(duration=duration, size=GPU_SIZE)
def run_gpu(reference, keys, guides, resolution, seed, extra, crop_close_views=False):
    started = time.perf_counter()
    torch.cuda.reset_peak_memory_stats()

    def generate_one(*, image, prompt, width, height, seed, check_budget):
        def callback(pipeline, step, timestep, callback_kwargs):
            check_budget()
            return callback_kwargs

        output = pipe(
            image=image,
            prompt=prompt,
            width=width,
            height=height,
            # Also bounds the area used to encode EACH reference; reducing only
            # PIL input dimensions does not avoid the pipeline's internal upscaling.
            output_resolution=REFERENCE_RESOLUTION[resolution],
            true_cfg_scale=1.0,
            use_kv_cache=True,
            num_images_per_prompt=1,
            generator=torch.Generator("cpu").manual_seed(seed),
            callback_on_step_end=callback,
        ).images[0]
        return rgb_on_white(output)

    allocation = duration(reference, keys, guides, resolution, seed, extra)
    result = run_batch(
        generate_one, reference, keys, guides, resolution, seed, extra, USE_ANGLES,
        budget_seconds=allocation - STARTUP_RESERVE_SECONDS,
        crop_close_views=crop_close_views,
    )
    metrics = {
        "model": MODEL_ID, "revision": MODEL_REVISION, "gpu": torch.cuda.get_device_name(),
        "gpu_size": GPU_SIZE, "reserved_seconds": allocation,
        "quota_budget_seconds": allocation * QUOTA_MULTIPLIER,
        "gpu_function_seconds": round(time.perf_counter() - started, 3),
        "peak_allocated_gib": round(torch.cuda.max_memory_allocated() / 1024**3, 2),
        "resolution": resolution, "pose_guides": guides is not None, "angles_lora": USE_ANGLES,
        "crop_close_views": crop_close_views,
        "seed": int(seed), "requested_views": keys, "completed_views": list(result.images),
        "timings": result.timings, "stop_reason": result.stop_reason,
    }
    return result.images, result.stop_reason, metrics