# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Custom modular-diffusers block: image OR text -> Ideogram 4 structured-JSON caption.

Wraps a vision-language model (Gemma-4-E4B via `Gemma4ForConditionalGeneration`) with two modes, selected by input:
  - `image` given  -> caption the image into Ideogram 4's native JSON schema (the `ideogram4-caption` Space logic).
  - `prompt` given (no image) -> enhance the short text idea into a JSON caption (Ideogram's "magic prompt",
    reusing the canonical `CAPTION_SYSTEM_MESSAGE` from `diffusers.pipelines.ideogram4.prompt_enhancer`).

Either way the produced `caption` can be fed straight into the Ideogram 4 generation pipeline as the `prompt`
(the model is trained on these JSON captions). This is the code-usable backend, minus the Gradio UI.
"""

import json
import math
import re

import torch

from diffusers.modular_pipelines.modular_pipeline import ModularPipelineBlocks, PipelineState, SequentialPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from diffusers.pipelines.ideogram4.prompt_enhancer import CAPTION_SYSTEM_MESSAGE, CAPTION_USER_TEMPLATE
from diffusers.utils import logging

from transformers import AutoProcessor, Gemma4ForConditionalGeneration


logger = logging.get_logger(__name__)  # pylint: disable=invalid-name


# Ideogram 4's native caption schema (from the official ideogram-oss/ideogram4 docs). Key ORDER matters for
# quality, bbox is [ymin, xmin, ymax, xmax] in 0-1000 normalized coords, colors are UPPERCASE #RRGGBB.
DEFAULT_INSTRUCTION = (
    "You are a captioner for the Ideogram 4 image model. Look at the image and output ONE JSON object (no prose, "
    "no markdown fences) that reconstructs it in Ideogram 4's native caption schema. Keep keys in the exact order "
    "shown.\n\n"
    "SCHEMA (keys in order):\n"
    "- high_level_description: one or two sentences summarizing the whole image (subject + medium/style).\n"
    "- style_description: an object with EXACTLY ONE of `photo` or `art_style`:\n"
    "    photo caption order:     aesthetics, lighting, photo, medium, color_palette\n"
    "    non-photo caption order: aesthetics, lighting, medium, art_style, color_palette\n"
    '    color_palette: optional, up to 16 UPPERCASE hex colors (e.g. "#1B1B2F").\n'
    "- compositional_deconstruction: {background, elements}.\n"
    "    background: the global setting/backdrop.\n"
    "    elements: list; keys in order per type:\n"
    "      obj:  type, bbox, desc, color_palette\n"
    "      text: type, bbox, text, desc, color_palette\n\n"
    "RULES:\n"
    "- bbox is [ymin, xmin, ymax, xmax] in 0-1000 normalized coordinates (origin top-left), independent of "
    "resolution. bbox and per-element color_palette are optional (<=5 colors per element).\n"
    "- Main subject first. A coherent subject (one person/animal/object) is exactly ONE element; its parts go in "
    "its `desc`. `desc` is identity-first and concrete (color, material, pose, distinguishing features).\n"
    "- Commit to one concrete value each (no 'various'/'or similar'). `text` elements carry the literal in-image "
    "characters, verbatim. Colors UPPERCASE #RRGGBB. Describe only what is visible; do not invent."
)


def _balance_json(frag: str) -> str:
    """Close unbalanced strings/brackets in a truncated JSON fragment so it can be parsed."""
    stack, in_str, esc = [], False, False
    for ch in frag:
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
        elif ch == '"':
            in_str = True
        elif ch in "{[":
            stack.append(ch)
        elif ch == "}" and stack and stack[-1] == "{":
            stack.pop()
        elif ch == "]" and stack and stack[-1] == "[":
            stack.pop()
    out = frag
    if in_str:
        out += '"'
    out = out.rstrip()
    if out.endswith(","):
        out = out[:-1]
    for ch in reversed(stack):
        out += "}" if ch == "{" else "]"
    return out


def _parse_json(text: str):
    """Parse the model output into (dict|None, repaired: bool). Robust to markdown fences, trailing prose, and
    truncated output (balances brackets). `repaired=True` means the output was incomplete and was closed to parse."""
    text = text.strip()
    text = re.sub(r"^```(?:json)?\s*", "", text)
    text = re.sub(r"\s*```$", "", text).strip()
    start = text.find("{")
    if start == -1:
        return None, False
    frag = text[start:]
    try:
        obj = json.loads(frag[: frag.rfind("}") + 1], strict=False)
        if isinstance(obj, dict):
            return obj, False
    except json.JSONDecodeError:
        pass
    try:
        obj, _ = json.JSONDecoder(strict=False).raw_decode(frag)
        if isinstance(obj, dict):
            return obj, False
    except json.JSONDecodeError:
        pass
    try:
        obj = json.loads(_balance_json(frag), strict=False)
        return (obj, True) if isinstance(obj, dict) else (None, False)
    except json.JSONDecodeError:
        return None, False


class Ideogram4CaptionStep(ModularPipelineBlocks):
    model_name = "ideogram4_caption"

    @property
    def description(self) -> str:
        return (
            "Turns an image OR a short text prompt into Ideogram 4's native structured-JSON caption using a "
            "vision-language model (Gemma-4). If `image` is given it is captioned; otherwise `prompt` is enhanced "
            "(magic-prompt) into a caption. Outputs `caption` (JSON string), `caption_json` (parsed dict, or None if "
            "unparseable), and `caption_raw` (raw decoded text)."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("caption_model", Gemma4ForConditionalGeneration, description="Vision-language captioner (Gemma-4)."),
            ComponentSpec("caption_processor", AutoProcessor, description="Processor paired with the captioner."),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam("image", description="PIL image to caption. If omitted, `prompt` is enhanced into a caption."),
            InputParam("prompt", type_hint=str, description="Short text idea to enhance into a JSON caption (used when no image is given)."),
            InputParam("instruction", default=DEFAULT_INSTRUCTION, description="Instruction for image-caption mode (ignored when enhancing a prompt)."),
            InputParam("height", type_hint=int, description="Target height; sets the aspect ratio hint for prompt enhancement."),
            InputParam("width", type_hint=int, description="Target width; sets the aspect ratio hint for prompt enhancement."),
            InputParam("max_new_tokens", type_hint=int, default=2048, description="Max tokens to generate (captions can be long)."),
            InputParam(
                "temperature",
                type_hint=float,
                default=0.0,
                description="Sampling temperature; 0 = greedy (deterministic).",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("caption", type_hint=str, description="Pretty-printed JSON caption (raw text if unparseable)."),
            OutputParam("caption_json", type_hint=dict, description="Parsed caption dict, or None if unparseable."),
            OutputParam("caption_raw", type_hint=str, description="Raw decoded model output."),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState):
        block_state = self.get_block_state(state)

        model = components.caption_model
        processor = components.caption_processor
        device = model.device  # run on wherever the model physically is (offload-friendly)

        if block_state.image is not None:
            content = [
                {"type": "image", "image": block_state.image},
                {"type": "text", "text": block_state.instruction},
            ]
        elif block_state.prompt:
            h = int(block_state.height or 1024)
            w = int(block_state.width or 1024)
            d = math.gcd(w, h) or 1
            aspect_ratio = f"{w // d}:{h // d}"
            enhance_text = (
                CAPTION_SYSTEM_MESSAGE
                + "\n\n"
                + CAPTION_USER_TEMPLATE.format(aspect_ratio=aspect_ratio, original_prompt=block_state.prompt)
            )
            content = [{"type": "text", "text": enhance_text}]
        else:
            raise ValueError("Provide either `image` (to caption) or `prompt` (to enhance into a caption).")

        messages = [{"role": "user", "content": content}]
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            enable_thinking=False,  # Gemma-4 reasoning mode off for structured-JSON output
        ).to(device)

        do_sample = float(block_state.temperature) > 0
        generated = model.generate(
            **inputs,
            max_new_tokens=int(block_state.max_new_tokens),
            do_sample=do_sample,
            temperature=float(block_state.temperature) if do_sample else None,
        )
        input_len = inputs["input_ids"].shape[1]
        text = processor.decode(generated[0][input_len:], skip_special_tokens=True)

        data, _repaired = _parse_json(text)
        block_state.caption_json = data
        block_state.caption = json.dumps(data, indent=2, ensure_ascii=False) if data is not None else text
        block_state.caption_raw = text

        self.set_block_state(state, block_state)
        return components, state


# auto_docstring
class Ideogram4CaptionBlocks(SequentialPipelineBlocks):
    """
    Standalone image/text -> Ideogram 4 structured-JSON caption pipeline (Gemma-4). Load with
    `ModularPipeline.from_pretrained(repo, trust_remote_code=True)` +
    `load_components(torch_dtype=torch.bfloat16)`, then either
    `pipe(image=img, output="caption")` (caption an image) or `pipe(prompt="a cat...", output="caption")`
    (enhance a text idea into a caption).
    """

    model_name = "ideogram4_caption"
    block_classes = [Ideogram4CaptionStep]
    block_names = ["caption"]

    @property
    def description(self) -> str:
        return "Image or text -> Ideogram 4 structured-JSON caption (Gemma-4; captions an image, or enhances a prompt)."

    @property
    def outputs(self) -> list[OutputParam]:
        return [
            OutputParam("caption", type_hint=str),
            OutputParam("caption_json", type_hint=dict),
            OutputParam("caption_raw", type_hint=str),
        ]