# Copyright 2026 Google LLC
#
# 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
#
#     https://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.

%pip install --upgrade --quiet google-genai

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

import base64
import os
import re

from IPython.display import Audio, Image, Markdown, display
from google import genai
from google.genai import types

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))

LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "global")

client = genai.Client(enterprise=True, project=PROJECT_ID, location=LOCATION)

def display_output(response):
    for part in response.parts:
        if part.text:
            text = part.text

            # Remove section tags, add new lines
            text = re.sub(r"\[\[.*?\]\]", "", text)
            text = re.sub(
                r"\s*\[(\d+\.\d+)(?::?(\d+\.\d+))?:?\]:?", r"\n\n[\1\2] ", text
            )
            text = text.replace("[:]", "\n\n")
            clean_text = text.strip()

            display(Markdown(clean_text))
        if part.inline_data:
            display(Audio(data=part.inline_data.data, autoplay=False))

music_clip_model = "lyria-3-clip-preview"
music_model = "lyria-3-pro-preview"

prompt = """
Sophisticated, rhythmic, and aspirational track with crisp 808 percussion, digital plucks,
and muted electric guitar rhythmic strums. Include breathy, airy Alto female vocal textures
with melodic, minimalist oohs and aahs with heavy reverb and rhythmic delay.
"""

response = client.models.generate_content(
    model=music_model,
    contents=prompt,
    config=types.GenerateContentConfig(response_modalities=["AUDIO", "TEXT"]),
)
display_output(response)

!wget -q https://storage.googleapis.com/cloud-samples-data/generative-ai/image/flowers.png

input_image = "flowers.png"
display(Image(input_image, width=400))

with open(input_image, "rb") as f:
    image = f.read()

response = client.models.generate_content(
    model=music_clip_model,
    contents=[
        types.Part.from_bytes(
            data=image,
            mime_type="image/png",
        ),
        "Generate an instrumental clip based on this input image that starts slowly and builds in intensity.",
    ],
    config=types.GenerateContentConfig(
        response_modalities=["TEXT", "AUDIO"],
    ),
)

display_output(response)

!wget -q https://storage.googleapis.com/cloud-samples-data/generative-ai/image/dog-ad-2.png

input_image = "dog-ad-2.png"
display(Image(input_image, width=400))

with open(input_image, "rb") as f:
    image = f.read()

genre_lyrics = """
Genre: Upbeat, acoustic Folk-Pop with a warm and cuddly vibe. Bright acoustic guitars, a soft shaker rhythm, and a friendly, melodic vocal.

Lyrics:
Tail wags and a heavy head,
Time to curl up in your favorite bed.
Soft as a cloud, a dream come true,
The perfect spot for a dog like you.
"""

response = client.models.generate_content(
    model=music_clip_model,
    contents=[
        types.Part.from_bytes(
            data=image,
            mime_type="image/png",
        ),
        genre_lyrics,
    ],
    config=types.GenerateContentConfig(
        response_modalities=["TEXT", "AUDIO"],
    ),
)

display_output(response)

def display_interaction_output(interaction):
    outputs = getattr(interaction, "outputs", []) or []

    for output in outputs:
        is_dict = isinstance(output, dict)
        output_type = output.get("type") if is_dict else getattr(output, "type", None)

        if output_type == "text":
            text = output.get("text", "") if is_dict else getattr(output, "text", "")
            if not text:
                continue

            # Remove section tags, add new lines
            text = re.sub(r'\[\[.*?\]\]', '', text)
            text = re.sub(r'\s*\[(\d+\.\d+)(?::?(\d+\.\d+))?:?\]:?', r'\n\n[\1\2] ', text)
            text = text.replace('[:]', '\n\n')
            clean_text = text.strip()
            display(Markdown(clean_text))

        elif output_type == "audio":
            data = output.get("data") if is_dict else getattr(output, "data", None)
            if not data:
                continue

            audio_data = base64.b64decode(data)
            display(Audio(data=audio_data, autoplay=False))

interaction = client.interactions.create(
    model=music_model,
    input="Genera un tema pop",
)
display_interaction_output(interaction)

stream = client.interactions.create(
    model=music_model,
    input="Generate a song about spending a day in Seoul in Korean.",
    stream=True,
)

for event in stream:
    if event.event_type == "content.delta":
        delta_dict = event.delta if isinstance(event.delta, dict) else getattr(event, "delta", {})
        if "text" in delta_dict:
            text = delta_dict["text"]
            text = text.replace("[:]", "\n[:]")
            text = re.sub(r"(\[\d+\.\d+:\])", r"\n\1", text)
            display(Markdown(text))
        elif "data" in delta_dict and "audio" in delta_dict.get("mime_type", ""):
            display(Audio(data=base64.b64decode(delta_dict["data"]), autoplay=False))