# 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 -qqq google-genai tqdm

# Use the environment variable if the user doesn't provide Project ID.
import os

from google import genai

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
LOCATION = "global"  # @param {type: "string"}
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", "global")

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

import sys

# Additional authentication is required for Google Colab
if "google.colab" in sys.modules:
    # Authenticate user to Google Cloud
    from google.colab import auth

    auth.authenticate_user()

! gcloud config set project {PROJECT_ID}
! gcloud auth application-default set-quota-project {PROJECT_ID}
! gcloud auth application-default login -q

import io
import random
import re
import wave
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from IPython.display import Audio
from google.genai.types import (
    GenerateContentConfig,
    PrebuiltVoiceConfig,
    SpeechConfig,
    VoiceConfig,
)
from pydantic import BaseModel
from tqdm import tqdm

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

LANGUAGE_CODE = "en-gb"

# Gemini TTS prebuilt voices, grouped so characters can be cast by gender.
# Gemini TTS prebuilt voices, grouped so characters can be cast by gender.
FEMALE_VOICES = [
    "Achernar",
    "Aoede",
    "Autonoe",
    "Callirrhoe",
    "Despina",
    "Erinome",
    "Gacrux",
    "Kore",
    "Laomedeia",
    "Leda",
    "Pulcherrima",
    "Sulafat",
    "Vindemiatrix",
    "Zephyr",
]
MALE_VOICES = [
    "Achird",
    "Algenib",
    "Algieba",
    "Alnilam",
    "Charon",
    "Enceladus",
    "Fenrir",
    "Iapetus",
    "Orus",
    "Puck",
    "Rasalgethi",
    "Sadachbia",
    "Sadaltager",
    "Schedar",
]

NARRATOR_VOICE = "Zubenelgenubi"
DEFAULT_VOICE = "Umbriel"

# Gemini TTS returns 24 kHz 16-bit mono PCM.
SAMPLE_RATE = 24000
SAMPLE_WIDTH = 2
SILENCE_LENGTH = 200  # In Milliseconds

# Number of lines to synthesize concurrently.
MAX_WORKERS = 8

SYSTEM_INSTRUCTION = """You are a creative and ambitious play writer. Your goal is to write a play for Text-to-Speech audio performance. Include a narrator character to describe the setting, scenes and actions occurring."""


class Character(BaseModel):
    name: str
    gender: str


class DialogueLine(BaseModel):
    speaker: str
    line: str


class Scene(BaseModel):
    setting: str
    dialogue: list[DialogueLine]


class Story(BaseModel):
    title: str
    characters: list[Character]
    scenes: list[Scene]

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(SAMPLE_WIDTH)
        wf.setframerate(SAMPLE_RATE)
        wf.writeframes(pcm)
    return buffer.getvalue()


def create_character_map(characters: list[Character]) -> dict[str, str]:
    """Maps characters to voices based on gender identified by Gemini."""
    # Create shuffled copies so we can simply .pop() from them
    male_voices = [v for v in MALE_VOICES if v != NARRATOR_VOICE]
    female_voices = list(FEMALE_VOICES)
    random.shuffle(male_voices)
    random.shuffle(female_voices)

    max_voices = len(male_voices) + len(female_voices)
    # Subtract 1 if Narrator is in characters since Narrator uses NARRATOR_VOICE
    if len(characters) > max_voices + 1:
        print(f"Too many characters {len(characters)}. Max {max_voices + 1}")

    character_map: dict[str, str] = {}

    for character in characters:
        if character.name == "Narrator":
            character_map[character.name] = NARRATOR_VOICE
            continue

        gender = character.gender.lower()

        # Try gender-matched pool first, fallback to whichever pool has voices remaining
        if gender == "female" and female_voices:
            pool = female_voices
        elif gender == "male" and male_voices:
            pool = male_voices
        else:
            # Picks male_voices if non-empty, otherwise female_voices
            pool = male_voices or female_voices

        if not pool:
            raise ValueError("Not enough voices to assign to all characters.")

        character_map[character.name] = pool.pop()

    return character_map


def synthesize_text(text: str, voice_name: str) -> bytes:
    """Synthesizes a single line with Gemini TTS and returns raw PCM audio."""
    response = client.models.generate_content(
        model=TTS_MODEL_ID,
        contents=f"TTS the following text with as much emotion and humor as possible\n\n{text}",
        config=GenerateContentConfig(
            speech_config=SpeechConfig(
                language_code=LANGUAGE_CODE,
                voice_config=VoiceConfig(
                    prebuilt_voice_config=PrebuiltVoiceConfig(voice_name=voice_name)
                ),
            ),
        ),
    )
    return response.candidates[0].content.parts[0].inline_data.data


def combine_audio_clips(clips: list[bytes]) -> bytes:
    """Joins PCM clips together, separated by a short silence."""
    silence = b"\x00" * int(SAMPLE_RATE * SAMPLE_WIDTH * SILENCE_LENGTH / 1000)
    return silence + silence.join(clips) + silence


def generate_audio_clips(story: Story, character_map: dict[str, str]) -> list[bytes]:
    lines: list[dict] = [
        {
            "line": story.title,
            "voice": character_map.get("Narrator", NARRATOR_VOICE),
        }
    ]

    # Process each scene in the play
    for scene in story.scenes:
        # Add the scene setting with the Narrator's voice
        lines.append(
            {
                "line": "Setting... " + scene.setting,
                "voice": character_map.get("Narrator", NARRATOR_VOICE),
            }
        )

        # Process each dialogue in the scene
        lines.extend(
            {
                "line": dialogue.line,
                "voice": character_map.get(dialogue.speaker, DEFAULT_VOICE),
            }
            for dialogue in scene.dialogue
        )

    # Synthesize the lines in parallel, keeping them in script order.
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        return list(
            tqdm(
                executor.map(
                    lambda line: synthesize_text(line["line"], line["voice"]), lines
                ),
                desc="Generating audio clips",
                total=len(lines),
            )
        )

PROMPT = """Write an interesting and humorous version of the full play Macbeth by William Shakespeare."""

response = client.models.generate_content(
    model=MODEL_ID,
    contents=PROMPT,
    config=GenerateContentConfig(
        system_instruction=SYSTEM_INSTRUCTION,
        response_mime_type="application/json",
        response_schema=Story,
    ),
)
story = response.parsed

with Path("macbeth_the_sitcom.json").open("r") as f:
    story = Story.model_validate_json(f.read())

story

character_to_voice = create_character_map(story.characters)
character_to_voice

audio_clips = generate_audio_clips(story, character_to_voice)

audio_data = wave_bytes(combine_audio_clips(audio_clips))

file_prefix = re.sub(r"[^\w.-]", "_", story.title).lower()
outfile_name = f"{file_prefix}-complete.wav"

with Path(outfile_name).open("wb") as f:
    f.write(audio_data)

print(f"Audio content written to file {outfile_name}")

Audio(audio_data)