app.py

import os
import uuid
import time
from dotenv import load_dotenv
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from google import genai
from google.genai import types
from google.cloud import texttospeech
import json
import asyncio
import re

# Load environment variables from .env file
load_dotenv()

PROJECT_ID = os.getenv("PROJECT_ID", "")
LOCATION = os.getenv("LOCATION", "")

client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Create directories
os.makedirs("static", exist_ok=True)
os.makedirs("dist", exist_ok=True)

# Mount static directory for generated images/audio
app.mount("/static", StaticFiles(directory="static"), name="static")

# Mount dist directory for Vue build output
if os.path.exists("dist/assets"):
    app.mount("/assets", StaticFiles(directory="dist/assets"), name="assets")


async def run_creative_agent_stream(user_prompt: str):
    tts_client = texttospeech.TextToSpeechAsyncClient()

    event_queue = asyncio.Queue()
    tts_text_queue = asyncio.Queue()
    enqueued_texts = set()
    scene_state_lock = asyncio.Lock()
    current_scene_index = -1
    pending_tts_count = {}
    scene_closed = {}
    emitted_scene_complete = set()

    async def maybe_emit_scene_complete(scene_index: int):
        if scene_index < 0:
            return

        should_emit = False
        async with scene_state_lock:
            if scene_index in emitted_scene_complete:
                return

            if (
                scene_closed.get(scene_index, False)
                and pending_tts_count.get(scene_index, 0) == 0
            ):
                emitted_scene_complete.add(scene_index)
                should_emit = True

        if should_emit:
            await event_queue.put({"type": "scene_complete", "sceneIndex": scene_index})

    async def enqueue_tts_text(text: str):
        nonlocal current_scene_index
        scene_index = -1

        async with scene_state_lock:
            scene_index = current_scene_index
            if scene_index >= 0:
                pending_tts_count[scene_index] = (
                    pending_tts_count.get(scene_index, 0) + 1
                )

        if scene_index >= 0:
            await tts_text_queue.put((scene_index, text))

    async def generate_image(prompt: str) -> str:
        nonlocal current_scene_index
        print(f"[Tool] Generating image: {prompt[:30]}...")
        time.sleep(2)
        filename = f"img_{uuid.uuid4().hex[:8]}.jpg"
        filepath = os.path.join("static", filename)

        try:
            result = await client.aio.models.generate_images(
                model="imagen-3.0-generate-001",
                prompt=prompt,
                config=types.GenerateImagesConfig(
                    number_of_images=1,
                    output_mime_type="image/jpeg",
                    aspect_ratio="16:9",
                ),
            )
            generated_images = result.generated_images or []
            if not generated_images or generated_images[0].image is None:
                raise RuntimeError("Image generation returned no image")
            generated_images[0].image.save(filepath)

            previous_scene_index = -1
            new_scene_index = -1
            async with scene_state_lock:
                previous_scene_index = current_scene_index
                current_scene_index += 1
                new_scene_index = current_scene_index
                pending_tts_count.setdefault(new_scene_index, 0)
                scene_closed[new_scene_index] = False
                if previous_scene_index >= 0:
                    scene_closed[previous_scene_index] = True

            print(f"[Tool] Image saved: {filepath}")
            await event_queue.put(
                {
                    "type": "image",
                    "src": f"/static/{filename}",
                    "sceneIndex": new_scene_index,
                }
            )
            await maybe_emit_scene_complete(previous_scene_index)
            return (
                "Image generated successfully. Now write the story text for this scene."
            )
        except Exception as e:
            print(f"[Tool Error] Image failed: {e}")
            return "Image generation failed. Continue writing."

    async def generate_speech(text: str) -> str:
        """Gemini tool to generate speech audio for given text."""
        print(f"[Tool] generate_speech called: {text[:30]}...")

        # Avoid duplicate audio generation
        if text in enqueued_texts:
            print(f"[Tool] Speech already enqueued, skipping: {text[:20]}...")
            return "Speech generation already in progress for this text."

        enqueued_texts.add(text)

        # Enqueue for TTS worker
        await enqueue_tts_text(text)

        return f"Speech generation queued for: {text[:30]}..."

    async def tts_worker():
        while True:
            item = await tts_text_queue.get()
            if item is None:
                break

            scene_index, text = item

            print(f"[TTS] Generating audio for: {text[:20]}...")
            filename = f"aud_{uuid.uuid4().hex[:8]}.mp3"
            filepath = os.path.join("static", filename)

            try:
                request = texttospeech.SynthesizeSpeechRequest(
                    input=texttospeech.SynthesisInput(text=text),
                    voice=texttospeech.VoiceSelectionParams(
                        language_code="en-US",
                        name="Charon",
                        model_name="gemini-2.5-flash-tts",
                    ),
                    audio_config=texttospeech.AudioConfig(
                        audio_encoding=texttospeech.AudioEncoding.MP3,
                        pitch=0.0,
                        speaking_rate=1.0,
                    ),
                )

                response = await tts_client.synthesize_speech(request=request)

                def save_audio_to_disk():
                    with open(filepath, "wb") as out:
                        out.write(response.audio_content)

                await asyncio.to_thread(save_audio_to_disk)

                print(f"[TTS] Audio saved: {filepath}")
                await event_queue.put(
                    {
                        "type": "audio",
                        "src": f"/static/{filename}",
                        "sceneIndex": scene_index,
                    }
                )
            except Exception as e:
                print(f"[TTS Error] Audio failed: {e}")
            finally:
                async with scene_state_lock:
                    if scene_index >= 0:
                        pending_tts_count[scene_index] = max(
                            0, pending_tts_count.get(scene_index, 0) - 1
                        )

                await maybe_emit_scene_complete(scene_index)
                tts_text_queue.task_done()

    system_instruction = """
    You are a storyteller. 
    Rule 1: Write the story in plain text ONLY. Absolutely NO HTML tags and NO Markdown formatting.
    Rule 2: For each scene, you MUST follow this exact sequence:
            Step A: Call `generate_image` to create the visual context.
            Step B: Write the narrative text for that scene, the text should be 3-4 sentences long.
    Rule 3: Repeat this sequence for all scenes in the story.
    """

    async def llm_producer():
        try:
            chat = client.aio.chats.create(
                model="gemini-2.5-pro",
                config=types.GenerateContentConfig(
                    system_instruction=system_instruction,
                    tools=[generate_image],
                    temperature=0.7,
                ),
            )

            response_stream = await chat.send_message_stream(user_prompt)
            text_buffer = ""

            async for chunk in response_stream:
                if chunk.text:
                    async with scene_state_lock:
                        scene_index_for_chunk = current_scene_index

                    if scene_index_for_chunk >= 0:
                        await event_queue.put(
                            {
                                "type": "text",
                                "chunk": chunk.text,
                                "sceneIndex": scene_index_for_chunk,
                            }
                        )

                    text_buffer += chunk.text

                    while True:
                        match = re.search(r"([.?!。？！\n]+[\s]*)", text_buffer)
                        if not match:
                            break

                        split_idx = match.end()
                        sentence = text_buffer[:split_idx].strip()
                        text_buffer = text_buffer[split_idx:]

                        if sentence:
                            await enqueue_tts_text(sentence)

            if text_buffer.strip():
                await enqueue_tts_text(text_buffer.strip())

        except Exception as e:
            print(f"LLM Producer Error: {e}")
        finally:
            async with scene_state_lock:
                final_scene_index = current_scene_index
                if final_scene_index >= 0:
                    scene_closed[final_scene_index] = True

            await maybe_emit_scene_complete(final_scene_index)
            await tts_text_queue.put(None)

    tts_task = asyncio.create_task(tts_worker())
    llm_task = asyncio.create_task(llm_producer())

    async def wait_all():
        await llm_task
        await tts_task
        await event_queue.put(None)

    asyncio.create_task(wait_all())

    while True:
        event = await event_queue.get()
        if event is None:
            break
        yield f"data: {json.dumps(event)}\n\n"


@app.get("/")
def serve_frontend():
    if os.path.exists("dist/index.html"):
        return FileResponse("dist/index.html")
    elif os.path.exists("demo.html"):
        return FileResponse("demo.html")
    else:
        return {
            "message": "TaleSpark API is running. Build the frontend with 'cd frontend && npm run build'"
        }


@app.post("/api/generate")
async def api_generate(request: Request):
    data = await request.json()
    user_prompt = data.get("prompt", "")

    return StreamingResponse(
        run_creative_agent_stream(user_prompt), media_type="text/event-stream"
    )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=8000)