# Copyright 2024 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 os

from google import genai

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
LOCATION = "global" # @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"))
if not LOCATION:
    LOCATION = os.environ.get("GOOGLE_CLOUD_REGION")

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

import io
import wave

from IPython.display import Audio, Markdown, display
from google.genai.types import (
    GenerateContentConfig,
    MultiSpeakerVoiceConfig,
    Part,
    PrebuiltVoiceConfig,
    SpeakerVoiceConfig,
    SpeechConfig,
    VoiceConfig,
)

MODEL_ID = "gemini-3.6-flash"  # @param {type: "string"}
TTS_MODEL_ID = "gemini-3.1-flash-tts-preview"  # @param {type: "string"}

def wave_bytes(pcm: bytes) -> bytes:
    """Wrap 24 kHz 16-bit mono PCM in a WAV container for in-notebook playback."""
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(24000)
        wf.writeframes(pcm)
    return buffer.getvalue()

PDF_URL = "https://arxiv.org/pdf/1706.03762"  # @param {type: "string"}

prompt = """
The dialogue should be engaging and natural, with each speaker contributing roughly equal amounts.
Return the dialogue as a freeform script in this format:

R: [dialogue]
S: [dialogue]

Use the following information to create the content for the podcast dialogue:
"""

response = client.models.generate_content(
    model=MODEL_ID,
    contents=[
        prompt,
        Part.from_uri(file_uri=PDF_URL, mime_type="application/pdf"),
    ],
    config=GenerateContentConfig(
        system_instruction="""You are a podcast writer. Your task is to generate a fun podcast-style dialogue between two speakers, Speaker R and Speaker S for Text-to-Speech.""",
    ),
)

dialogue = response.text
print("Generated Dialogue:")
display(Markdown(dialogue))

tts_prompt = f"""TTS the following conversation between speakers R & S: {dialogue}"""

response = client.models.generate_content(
    model="gemini-3.1-flash-tts-preview",
    contents=tts_prompt,
    config=GenerateContentConfig(
        speech_config=SpeechConfig(
            language_code="en-us",
            multi_speaker_voice_config=MultiSpeakerVoiceConfig(
                speaker_voice_configs=[
                    SpeakerVoiceConfig(
                        speaker="R",
                        voice_config=VoiceConfig(
                            prebuilt_voice_config=PrebuiltVoiceConfig(
                                voice_name="Kore",
                            )
                        ),
                    ),
                    SpeakerVoiceConfig(
                        speaker="S",
                        voice_config=VoiceConfig(
                            prebuilt_voice_config=PrebuiltVoiceConfig(
                                voice_name="Achird",
                            )
                        ),
                    ),
                ]
            ),
        ),
    ),
)

audio_data = wave_bytes(response.candidates[0].content.parts[0].inline_data.data)

Audio(audio_data)