# Multimodal Text & Vision Response Generator
from PIL import Image
import torch
from config import MOCK_MODE

# Import spaces wrapper conditionally for Hugging Face ZeroGPU compatibility
try:
    import spaces

    has_spaces = True
except ImportError:
    has_spaces = False


def gpu_decorator(fn):
    """Decorator to conditionally apply spaces.GPU on Hugging Face Spaces."""
    if has_spaces:
        return spaces.GPU(fn)
    return fn


@gpu_decorator
def generate_response(
    prompt: str,
    image: Image.Image | None = None,
    max_new_tokens: int = 1024,
    temperature: float = 0.4,
) -> str:
    """Generates a text completion from Gemma 4 given a prompt and optional image."""
    # Under MOCK_MODE, return a default mock structured copy output instantly
    if MOCK_MODE:
        return """
---SECTION 1: STRUCTURED ATTRIBUTES---
Category: Ceramic Mug
Materials: stoneware clay, seafoam drip glaze
Colors: seafoam green, cream, sandy brown
Style: rustic, cozy

---SECTION 2: SEO LISTING COPY---
Title: Hand-Thrown Seafoam Green Stoneware Clay Mug - Cozy 12oz Rustic Coffee Cup
Description: Savor your favorite warm beverages in this beautifully rustic, hand-thrown stoneware mug. Each piece is individually crafted on the pottery wheel and finished in a seafoam green drip glaze that mimics ocean foam rolling onto sandy shores. Its sturdy clay body retains heat, keeping your tea or coffee warm as you start your morning. Ideal for cozy cabin kitchens or as a unique gift for stoneware pottery collectors.
Tags: ceramic mug, hand thrown cup, stoneware pottery, seafoam drip glaze, wheel thrown mug, handmade clay mug, rustic coffee cup, artisan teacup, earth clay mugs, green ceramic mug, cabin kitchenware, unique ceramic gift, cozy coffee mugs

---SECTION 3: SOCIAL MARKETING---
Instagram: Morning rituals just got cozier. 🌊☕ Our new wheel-thrown stoneware mug features a seafoam drip glaze that makes every sip feel like a seaside breeze. Handcrafted, unique, and built to last. Tap the link in our bio to shop the kitchen collection! #potterylife #handcraftedstoneware #ceramicmug #oceanvibes #cozymorning
Pinterest: Wheel-thrown green ceramic stoneware mug with seafoam drip glaze. Handmade pottery rustic coffee cups and tea mugs. Click through to view details.
"""

    # Lazy-load transformers model and processor
    from .loader import load_gemma_model, get_device

    model, processor = load_gemma_model()
    device = get_device()

    # Build prompt structure
    content = []
    if image is not None:
        content.append({"type": "image", "image": image})
    content.append({"type": "text", "text": prompt})
    messages = [{"role": "user", "content": content}]

    # Format chats template syntax
    text_prompt = processor.apply_chat_template(messages, add_generation_prompt=True)

    # Tokenize input assets
    if image is not None:
        inputs = processor(text=text_prompt, images=image, return_tensors="pt")
    else:
        inputs = processor(text=text_prompt, return_tensors="pt")

    # Move tensors to hardware
    inputs = {k: v.to(device) for k, v in inputs.items()}
    if "pixel_values" in inputs:
        inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
    if "audio_values" in inputs:
        inputs["audio_values"] = inputs["audio_values"].to(model.dtype)

    # Run inference under non-gradient checks
    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=True if temperature > 0.0 else False,
        )

    # Extract model output tokens and decode
    input_len = inputs["input_ids"].shape[1]
    response_ids = generated_ids[0][input_len:]
    response_text = processor.decode(response_ids, skip_special_tokens=True)
    return response_text.strip()