"""Shared CPU policy for the H3 generator and its optional efficient conditioner.

No model loads, networking, CUDA initialization or monkey-patches at import.
The setup extension uses the pinned Diffusers validation and media normalization.
"""
from __future__ import annotations

import math
from PIL import Image

PROTOCOL = "h3-reference-budget-v1"
RESIZE_MODES = ("legacy", "match")


def validate_resize_mode(mode):
    if mode not in RESIZE_MODES:
        raise ValueError("Unsupported H3 reference resize policy.")
    return mode


def scheduler_points(evaluations):
    n = float(evaluations)
    if not math.isfinite(n) or not n.is_integer() or not 4 <= n <= 40:
        raise ValueError("H3 steps must be an integer from 4 to 40.")
    # The pinned MiniMaxH3Scheduler includes terminal zero in the point count.
    return int(n) + 1


def reference_size(size, canvas, mode="match", multiple=32):
    validate_resize_mode(mode)
    width, height = map(int, size)
    target_width, target_height = map(int, canvas)
    if min(width, height, target_width, target_height) <= 0:
        raise ValueError("Image and canvas dimensions must be positive.")
    if max(width / height, height / width) > 4:
        raise ValueError("H3 image references must have an aspect ratio between 1:4 and 4:1.")
    if mode == "legacy":
        scale = 2048 / min(width, height)
        return tuple(max(multiple, round(edge * scale / multiple) * multiple) for edge in (width, height))
    scale = min(1.0, math.sqrt(target_width * target_height / (width * height)))
    # Floor rather than round: the policy never exceeds its source or area cap,
    # except the mandatory minimum cell for tiny (<32 px) images.
    return tuple(max(multiple, int(edge * scale) // multiple * multiple) for edge in (width, height))


def resize_reference(image, canvas, mode="match"):
    size = reference_size(image.size, canvas, mode)
    return image if image.size == size else image.resize(size, Image.Resampling.LANCZOS)


def efficient_blocks(original_blocks_type):
    """Create a local block class; never mutate the upstream or global pipeline."""
    from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
    from diffusers.modular_pipelines.modular_pipeline_utils import InputParam
    from diffusers.modular_pipelines.minimax_h3 import MiniMaxH3ImageReference

    class BudgetReferenceSetup(MiniMaxH3Ref2VASetupStep):
        @property
        def inputs(self):
            return [*super().inputs, InputParam(
                name="reference_resize_mode", type_hint=str, default="legacy",
                description="Negotiated reference policy; must match the conditioning encoder.")]

        def __call__(self, components, state):
            # Preserve all upstream validation, audio handling and frame-grid rules.
            # Its image enlargement is CPU-only; the replacement below occurs
            # before either the Qwen encoder or the video VAE consumes the images.
            mode = validate_resize_mode(self.get_block_state(state).reference_resize_mode)
            components, state = super().__call__(components, state)
            if mode == "match":
                # get_block_state() only projects declared inputs. Normalized
                # references are an OUTPUT, available on the full PipelineState.
                # Read the resolved canvas there too, then replace only this
                # output; all upstream audio/video/frame results stay intact.
                canvas = (state.get("width"), state.get("height"))
                normalized_references = [
                    MiniMaxH3ImageReference(image=resize_reference(original.image, canvas, mode))
                    if original.kind == "image" else normalized
                    for original, normalized in zip(state.get("references"), state.get("normalized_references"))
                ]
                state.set("normalized_references", normalized_references)
            return components, state

    class EfficientBlocks(original_blocks_type):
        block_classes = [BudgetReferenceSetup if cls is MiniMaxH3Ref2VASetupStep else cls
                         for cls in original_blocks_type.block_classes]

    return EfficientBlocks


def validate_canvas(height, width):
    raw = (float(height), float(width))
    if any(not math.isfinite(x) or not x.is_integer() for x in raw):
        raise ValueError("Canvas dimensions must be whole numbers.")
    height, width = map(int, raw)
    if min(height, width) < 256 or max(height, width) > 1536 or height % 32 or width % 32:
        raise ValueError("H3 canvas must use multiples of 32, between 256 and 1536 pixels.")
    if height * width > 1_200_000 or max(height / width, width / height) > 4:
        raise ValueError("Canvas exceeds the supported area or aspect ratio.")
    return height, width