# The Cloud Functions for Firebase SDK to create Cloud Functions and set up triggers.
from datetime import datetime, timedelta
from typing import Any, Dict, List

import google.cloud.firestore
import requests
from google import genai
from google.genai import types
import json

# The Firebase Admin SDK to access Cloud Firestore.
from firebase_admin import firestore, initialize_app
from firebase_functions import firestore_fn, https_fn, identity_fn, params, scheduler_fn

# Initialize Firebase Admin with project ID
# For emulator, no credentials needed - just specify project ID
app = initialize_app(options={"projectId": "suno-eden"})

# Suno API Configuration
SUNO_BASE_URL = "https://studio-api.staging.suno.com"

# Define the secret parameter
suno_session_token = params.SecretParam("SUNO_SESSION_TOKEN")

# Hardcoded list of allowed signup emails (replace with real emails)
ALLOWED_SIGNUP_EMAILS: List[str] = [
    "avi.naim@suno.com",
    "samuel@suno.com",
    "jack@suno.com",
    "georg@suno.com",
    "mikey@suno.com",
    "yamill@suno.com",
    "kaveh@suno.com",
    "martin@suno.com",
    "charlie@suno.com",
    "brad@suno.com",
    "paul@suno.com",
    "henry@suno.com",
    "philip@suno.com",
]

# Block account creation with any non-acme email address.
@identity_fn.before_user_created()
def validatenewuser(
    event: identity_fn.AuthBlockingEvent,
) -> identity_fn.BeforeCreateResponse | None:
    # User data passed in from the CloudEvent.
    user = event.data

    # Only users of a specific domain can sign up.
    if user.email is None or user.email not in ALLOWED_SIGNUP_EMAILS:
        # Return None so that Firebase Auth rejects the account creation.
        raise https_fn.HttpsError(
            code=https_fn.FunctionsErrorCode.INVALID_ARGUMENT,
            message="Unauthorized email",
        )

GENERATE_TAGS_SYSTEM_PROMPT = """
You are a grammy award winning record producer. Your task is to help a music-making novice create incredible covers/remixes of music.

The user will share the name of the song they'd like to cover/remix and some instructions on how they'd like to cover/remix it.

If the user references the name of an artist in their prompt, you must NEVER include the artist's name in either the cover name or the tags.
"""

def generate_tags(original_song: str, prompt: str) -> str:
    client = genai.Client(vertexai=True, project="suno-eden", location="us-central1")

    function = types.FunctionDeclaration(
        name='produce_song_cover',
        description='Produce a cover of a song',
        parameters_json_schema={
            'type': 'object',
            'properties': {
                'remix_name': {
                    'type': 'string',
                    'description': 'Name of the cover (must include the original song name and 1-2 words describing the cover/remix style in parenthesis, e.g. "Vogue (Acoustic Version)")',
                },
                'tags': {
                    'type': 'string',
                    'description': 'Concise comma-separated list of 4-5 genres and production elements that describe the overall musical style and structure of the song (e.g. Intimate acoustic folk, tense strings, driving indie rock bridge, contemplative piano outro)',
                }
            },
            'required': ['remix_name', 'tags'],
        },
    )

    tool = types.Tool(function_declarations=[function])

    resp = client.models.generate_content(
        model="gemini-2.5-flash",
        contents="Original song: " + original_song + "\n\nPrompt: " + prompt,
        config=types.GenerateContentConfig(
            system_instruction=GENERATE_TAGS_SYSTEM_PROMPT,
            tools=[tool],
            tool_config=types.ToolConfig(
                function_calling_config=types.FunctionCallingConfig(mode='ANY')
            ),
            thinking_config=types.ThinkingConfig(
                thinking_budget=1024,  # Use `0` to turn off thinking
            )
        )
    )

    return resp.function_calls[0].args

@https_fn.on_request()
def gemini_http_vertex(req: https_fn.Request) -> https_fn.Response:
    if req.method == "OPTIONS":
        resp = https_fn.Response("", status=204)
        resp.headers["Access-Control-Allow-Origin"] = "*"
        resp.headers["Access-Control-Allow-Methods"] = "POST, OPTIONS"
        resp.headers["Access-Control-Allow-Headers"] = "Content-Type, Authorization"
        return resp

    try:
        if req.method != "POST":
            return https_fn.Response(
                json.dumps({"error": "Use POST with JSON body."}),
                status=405,
                content_type="application/json",
                headers={"Access-Control-Allow-Origin": "*"},
            )

        body = req.get_json(silent=True) or {}
        original_song = body.get("original_song")
        prompt = body.get("prompt")
        model = body.get("model", "gemini-2.5-flash")
        location = body.get("location", "us-central1")
        project = "suno-eden"

        if not prompt:
            return https_fn.Response(
                json.dumps({"error": "Missing 'prompt'."}),
                status=400,
                content_type="application/json",
                headers={"Access-Control-Allow-Origin": "*"},
            )

        tags = generate_tags(original_song, prompt)

        return https_fn.Response(
            json.dumps({"text": tags, "model": model, "location": location}),
            status=200,
            content_type="application/json",
            headers={"Access-Control-Allow-Origin": "*"},
        )

    except Exception as e:
        return https_fn.Response(
            json.dumps({"error": str(e)}),
            status=500,
            content_type="application/json",
            headers={"Access-Control-Allow-Origin": "*"},
        )

def call_suno_generate_api(
    prompt: str,
    session_token: str,
    remix_of_id: str | None = None,
    tags: str | None = None,
    title: str | None = None,
) -> List[Dict[str, Any]]:
    """
    Call the Suno API to generate songs.

    Args:
        prompt: The generation prompt (or lyrics for cover)
        session_token: Suno API session token
        remix_of_id: Optional Suno clip ID to remix from
        tags: Optional tags for the song (used when remixing)
        title: Optional title for the generated song

    Returns:
        List of clip dictionaries from the API response
    """
    body: Dict[str, Any] = {
        "prompt": prompt,
        "mv": "chirp-crow",
        "make_instrumental": False,
        "wait_audio": False,
    }

    # Add title if provided
    if title:
        body["title"] = title

    # Add tags if provided
    if tags:
        body["tags"] = tags

    # Add cover_clip_id if remixing
    if remix_of_id:
        body["cover_clip_id"] = remix_of_id
        body["task"] = "cover"

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {session_token}",
    }

    response = requests.post(
        f"{SUNO_BASE_URL}/api/generate/v2-web/", json=body, headers=headers
    )

    if not response.ok:
        raise https_fn.HttpsError(
            code=https_fn.FunctionsErrorCode.INTERNAL,
            message=f"Failed to generate songs: {response.status_code} {response.text}",
        )

    data = response.json()
    return data.get("clips", [])


@https_fn.on_call(secrets=[suno_session_token])
def generate_song(req: https_fn.CallableRequest) -> Dict[str, Any]:
    """
    Generate songs via Suno API and store them in Firestore.
    Requires authentication - user must be signed in.

    Expected parameters:
        - prompt (string, required): The generation prompt
        - remixOfId (string, optional): Firestore song ID to remix from

    Returns:
        Dictionary with the created Song IDs
    """
    # Check if user is authenticated
    if req.auth is None:
        raise https_fn.HttpsError(
            code=https_fn.FunctionsErrorCode.UNAUTHENTICATED,
            message="User must be authenticated to generate songs",
        )

    # Get user ID from Firebase Auth
    user_id = req.auth.uid

    # Extract parameters from request data
    prompt = req.data.get("prompt")
    remix_of_song_id = req.data.get("remixOfId")

    # Validate required parameters
    if not prompt:
        raise https_fn.HttpsError(
            code=https_fn.FunctionsErrorCode.INVALID_ARGUMENT,
            message="Missing required parameter: prompt",
        )

    try:
        # Get Firestore client
        firestore_client: google.cloud.firestore.Client = firestore.client()

        # If remixing, look up the source song to get the Suno clip ID and lyrics
        suno_clip_id = None
        source_song_ref = None
        source_lyrics = None
        tags = None
        title = None
        root_remix_ref = None
        pre_song_refs: List[google.cloud.firestore.DocumentReference] = []

        if remix_of_song_id:
            source_song_doc = (
                firestore_client.collection("song").document(remix_of_song_id).get()
            )

            if not source_song_doc.exists:
                raise https_fn.HttpsError(
                    code=https_fn.FunctionsErrorCode.NOT_FOUND,
                    message=f"Source song not found: {remix_of_song_id}",
                )

            source_song_data = source_song_doc.to_dict()

            source_artist_doc = (
                firestore_client.collection("user").document(source_song_data.get("createdBy", "").id).get()
            )

            source_artist_data = source_artist_doc.to_dict()

            suno_clip_id = source_song_data.get("sunoClipId")

            if not suno_clip_id:
                raise https_fn.HttpsError(
                    code=https_fn.FunctionsErrorCode.INVALID_ARGUMENT,
                    message="Source song does not have a Suno clip ID",
                )

            # Get the lyrics from the source song
            source_lyrics = source_song_data.get("lyrics", "")

            # Get the original song name and append " (Remix)"
            source_song_name = source_song_data.get("name", "Untitled")

            # Establish source and root remix references early (before tag generation)
            source_song_ref = source_song_doc.reference
            root_remix_ref = source_song_data.get("rootRemixSong") or source_song_ref

            # Pre-create two placeholder song documents so they appear immediately
            for _ in range(2):
                pre_song_data = {
                    "name": f"{source_song_name} (Remix)",
                    "createdBy": firestore_client.collection("user").document(user_id),
                    "createdAt": datetime.now(),
                    "originalPrompt": prompt,
                    "rewrittenPrompt": "",
                    "lyrics": "",
                    "imageURL": "",
                    "videoURL": None,
                    "audioURL": None,
                    "isPublished": False,
                    "playCount": 0,
                    "isUpload": False,
                    "status": "queued",
                }

                # Add rootRemixSong reference to placeholder if remixing
                if root_remix_ref:
                    pre_song_data["rootRemixSong"] = root_remix_ref

                _, pre_ref = firestore_client.collection("song").add(pre_song_data)
                pre_song_refs.append(pre_ref)

            generated_tags = generate_tags(source_song_name + " - " + source_artist_data.get("displayName", ""), prompt)

            title = generated_tags.get("remix_name", source_song_name + " (Remix)")
            tags = generated_tags.get("tags", prompt)

            # Update placeholder names to the generated remix name
            if title:
                for ref in pre_song_refs:
                    ref.update({
                        "name": title,
                        "rewrittenPrompt": tags
                    })

        # Determine the prompt and tags for the API call
        # When remixing (cover), use source lyrics as prompt and user prompt as tags
        # Otherwise, use user prompt as prompt with no tags
        api_prompt = source_lyrics if source_lyrics else prompt

        # Call Suno API to generate songs
        clips = call_suno_generate_api(
            api_prompt, suno_session_token.value, suno_clip_id, tags, title
        )

        # Create Song documents in Firestore
        song_ids = []
        # If we pre-created placeholders (remix flow)
        if pre_song_refs:
            # If no clips returned, clean up placeholders and error
            if not clips:
                for ref in pre_song_refs:
                    try:
                        ref.update({"status": "error"})
                    except Exception:
                        pass
                raise https_fn.HttpsError(
                    code=https_fn.FunctionsErrorCode.INTERNAL,
                    message="No clips returned from Suno API",
                )

            # Populate up to two placeholders with first clips
            for i, clip in enumerate(clips):
                if i < len(pre_song_refs):
                    update_data = {
                        "name": clip.get("title", "Untitled"),
                        "sunoClipId": clip.get("id"),
                        "lyrics": clip.get("metadata", {}).get("prompt", ""),
                        "imageURL": clip.get("image_url", ""),
                        "videoURL": clip.get("video_url"),
                        "originalPrompt": prompt,
                        "rewrittenPrompt": clip.get("metadata", {}).get("tags", ""),
                        "status": "queued",
                        "audioURL": clip.get("audio_url"),
                    }
                    try:
                        pre_song_refs[i].update(update_data)
                        song_ids.append(pre_song_refs[i].id)
                    except Exception:
                        # If update fails, fall back to creating a new doc
                        song_data = {
                            "name": clip.get("title", "Untitled"),
                            "sunoClipId": clip.get("id"),
                            "createdBy": firestore_client.collection("user").document(user_id),
                            "lyrics": clip.get("metadata", {}).get("prompt", ""),
                            "imageURL": clip.get("image_url", ""),
                            "videoURL": clip.get("video_url"),
                            "createdAt": datetime.now(),
                            "originalPrompt": prompt,
                            "rewrittenPrompt": clip.get("metadata", {}).get("tags", ""),
                            "isPublished": False,
                            "playCount": 0,
                            "isUpload": False,
                            "status": "queued",
                            "audioURL": clip.get("audio_url"),
                        }
                        if source_song_ref and root_remix_ref:
                            song_data["rootRemixSong"] = root_remix_ref
                        _, song_ref = firestore_client.collection("song").add(song_data)
                        song_ids.append(song_ref.id)
                else:
                    # More than two clips: create additional documents as needed
                    clip = clips[i]
                    song_data = {
                        "name": clip.get("title", "Untitled"),
                        "sunoClipId": clip.get("id"),
                        "createdBy": firestore_client.collection("user").document(user_id),
                        "lyrics": clip.get("metadata", {}).get("prompt", ""),
                        "imageURL": clip.get("image_url", ""),
                        "videoURL": clip.get("video_url"),
                        "createdAt": datetime.now(),
                        "originalPrompt": prompt,
                        "rewrittenPrompt": clip.get("metadata", {}).get("tags", ""),
                        "isPublished": False,
                        "playCount": 0,
                        "isUpload": False,
                        "status": "queued",
                        "audioURL": clip.get("audio_url"),
                    }
                    if source_song_ref and root_remix_ref:
                        song_data["rootRemixSong"] = root_remix_ref
                    _, song_ref = firestore_client.collection("song").add(song_data)
                    song_ids.append(song_ref.id)

            # If fewer than two clips, delete the extra pre-generated placeholders
            if len(clips) < len(pre_song_refs):
                for j in range(len(clips), len(pre_song_refs)):
                    try:
                        pre_song_refs[j].delete()
                    except Exception:
                        pass
        else:

            if not clips:
                raise https_fn.HttpsError(
                    code=https_fn.FunctionsErrorCode.INTERNAL,
                    message="No clips returned from Suno API",
                )

            # Non-remix path (no placeholders) - create documents for all clips as before
            for clip in clips:
                song_data = {
                    "name": clip.get("title", "Untitled"),
                    "sunoClipId": clip.get("id"),
                    "createdBy": firestore_client.collection("user").document(user_id),
                    "lyrics": clip.get("metadata", {}).get("prompt", ""),
                    "imageURL": clip.get("image_url", ""),
                    "videoURL": clip.get("video_url"),
                    "createdAt": datetime.now(),
                    "originalPrompt": prompt,
                    "rewrittenPrompt": clip.get("metadata", {}).get("tags", ""),
                    "isPublished": False,
                    "playCount": 0,
                    "isUpload": False,
                    "status": "queued",
                    "audioURL": clip.get("audio_url"),
                }
                if source_song_ref and root_remix_ref:
                    song_data["rootRemixSong"] = root_remix_ref
                _, song_ref = firestore_client.collection("song").add(song_data)
                song_ids.append(song_ref.id)

        # Mark only the true root song as remixed
        if root_remix_ref:
            try:
                root_snap = root_remix_ref.get()
                if root_snap.exists:
                    root_remix_ref.update({"hasBeenRemixed": True})
            except Exception:
                # If the root doc is missing or update fails, skip silently
                pass

        # Return the created Song IDs
        return {"songIds": song_ids, "message": "Songs created successfully"}

    except https_fn.HttpsError:
        # Re-raise HttpsError as-is
        raise
    except Exception as e:
        # traceback
        import traceback

        traceback.print_exc()
        raise https_fn.HttpsError(
            code=https_fn.FunctionsErrorCode.INTERNAL,
            message=f"Internal error: {str(e)}",
        )


def get_song_status_from_suno(
    clip_ids: List[str], session_token: str
) -> List[Dict[str, Any]]:
    """
    Fetch song status from Suno API using feed v2 endpoint.

    Args:
        clip_ids: List of Suno clip IDs to check
        session_token: Suno API session token

    Returns:
        List of clip metadata dictionaries
    """
    if not clip_ids:
        return []

    ids_param = ",".join(clip_ids)

    headers = {
        "Authorization": f"Bearer {session_token}",
    }

    response = requests.get(
        f"{SUNO_BASE_URL}/api/feed/v2/?ids={ids_param}", headers=headers
    )

    if not response.ok:
        raise Exception(
            f"Failed to fetch song status: {response.status_code} {response.text}"
        )

    data = response.json()
    return data.get("clips", [])


def poll_and_update_songs(session_token: str) -> Dict[str, Any]:
    """
    Poll Suno API for song status updates and update Firestore.
    This is the core polling logic that can be called from both scheduled and callable functions.

    Args:
        session_token: Suno API session token

    Returns:
        Dictionary with polling results (songs updated, errors, etc.)
    """
    firestore_client: google.cloud.firestore.Client = firestore.client()

    # Calculate the cutoff time (10 minutes ago)
    cutoff_time = datetime.now() - timedelta(minutes=10)

    # Query for songs that need status updates
    # - status is not "complete" or "error"
    # - created within the last 10 minutes
    songs_query = (
        firestore_client.collection("song")
        .where("status", "not-in", ["complete", "error"])
        .where("createdAt", ">=", cutoff_time)
        .limit(100)  # Process up to 100 songs at a time
    )

    songs = songs_query.get()

    if not songs:
        print("  No songs to poll")
        return {"updated": 0, "message": "No songs to poll"}

    # Build a map of sunoClipId -> song document
    clip_id_to_song = {}
    clip_ids = []

    for song in songs:
        song_data = song.to_dict()
        clip_id = song_data.get("sunoClipId")
        if clip_id:
            clip_id_to_song[clip_id] = (song.reference, song_data)
            clip_ids.append(clip_id)

    if not clip_ids:
        print("  No valid clip IDs to poll")
        return {"updated": 0, "message": "No valid clip IDs to poll"}

    print(f"  Polling status for {len(clip_ids)} songs")

    try:
        # Fetch status from Suno API
        clips = get_song_status_from_suno(clip_ids, session_token)

        updated_count = 0

        # Update Firestore with the latest data
        for clip in clips:
            clip_id = clip.get("id")
            if clip_id not in clip_id_to_song:
                continue

            song_ref, current_data = clip_id_to_song[clip_id]

            # Determine status
            clip_status = clip.get("status", "").lower()
            if clip_status == "complete":
                status = "complete"
            elif clip_status == "streaming":
                status = "streaming"
            elif clip_status == "error" or clip_status == "failed":
                status = "error"
            else:
                status = "queued"  # Still processing

            # Prepare update data
            update_data = {
                "status": status,
                "audioURL": clip.get("audio_url") or current_data.get("audioURL"),
                "imageURL": clip.get("image_url") or current_data.get("imageURL"),
                "videoURL": clip.get("video_url") or current_data.get("videoURL"),
            }

            # Update the name if it changed
            if clip.get("title"):
                update_data["name"] = clip.get("title")

            # Update lyrics if available
            if clip.get("metadata", {}).get("prompt"):
                update_data["lyrics"] = clip.get("metadata", {}).get("prompt")

            # Update the song document
            song_ref.update(update_data)
            print(f"  Updated song {clip_id}: status={status}")
            updated_count += 1

        return {
            "updated": updated_count,
            "total": len(clip_ids),
            "message": f"Updated {updated_count} songs",
        }

    except Exception as e:
        error_msg = f"Error polling song status: {str(e)}"
        print(f"  {error_msg}")
        return {"updated": 0, "error": error_msg}


def get_hooks_start_point_for_clip(
    suno_clip_id: str,
    session_token: str,
    lyrics: str | None = None,
    max_attempts: int = 15,
    delay_seconds: float = 2.0,
) -> int | None:
    """
    Poll the Suno aligned lyrics API until the alignment is complete,
    then return the rounded start time (in seconds) of the first chorus.

    Args:
        suno_clip_id: Suno clip ID for the song.
        session_token: Suno API session token.
        lyrics: The song's lyrics text from Firestore, used to skip analysis
            when no [Chorus] section is present.
        max_attempts: Maximum number of polling attempts (when needed).
        delay_seconds: Delay between polling attempts.

    Returns:
        Integer start time in seconds for the first chorus, or None if
        the chorus could not be determined.
    """
    import time

    # If the lyrics do not contain a [Chorus] section, there is no point in
    # calling the aligned lyrics API; the marker will not appear there either.
    # In that case we fast-fail and treat the hook as starting at 0s.
    if lyrics is not None and "[Chorus]" not in lyrics and "[Verse" not in lyrics:
        print(
            f"No [Chorus] or [Verse] section in lyrics; "
            f"setting hooksStartPoint=0 without aligned lyrics lookup "
            f"for clip {suno_clip_id}."
        )
        return 0

    if not suno_clip_id:
        return None

    url = f"{SUNO_BASE_URL}/api/gen/{suno_clip_id}/aligned_lyrics/v3"
    headers = {
        "Authorization": f"Bearer {session_token}",
    }

    for attempt in range(1, max_attempts + 1):
        try:
            response = requests.get(url, headers=headers, timeout=15)
        except Exception as e:
            print(f"Error calling aligned lyrics API for {suno_clip_id}: {e}")
            return None

        if not response.ok:
            print(
                f"Aligned lyrics API returned {response.status_code} "
                f"for {suno_clip_id}: {response.text}"
            )
            return None

        try:
            data = response.json()
        except Exception as e:
            print(f"Failed to parse aligned lyrics response for {suno_clip_id}: {e}")
            return None

        state = data.get("state")
        alignment = data.get("alignment") or []

        if state == "complete":
            for item in alignment:
                word = str(item.get("word", ""))
                if "[Chorus]" in word:
                    start_s = item.get("start_s")
                    if isinstance(start_s, (int, float)):
                        hooks_start = float(start_s)
                        print(
                            f"Found chorus for {suno_clip_id} at {hooks_start}s "
                            f"(start_s={start_s})"
                        )
                        return hooks_start

            # Completed but no chorus marker found
            print(
                f"No [Chorus] marker found in alignment for {suno_clip_id} "
                f"after completion, trying [Verse] marker."
            )

            for item in alignment:
                word = str(item.get("word", ""))
                if "[Verse" in word:
                    start_s = item.get("start_s")
                    if isinstance(start_s, (int, float)):
                        hooks_start = float(start_s)
                        print(
                            f"Found verse for {suno_clip_id} at {hooks_start}s "
                            f"(start_s={start_s})"
                        )
                        return hooks_start

            # Completed but no verse marker found
            print(
                f"No [Verse] marker found in alignment for {suno_clip_id} "
                f"after completion."
            )

            return None

        # If the API indicates a failure, stop polling
        if state in {"failed", "error"}:
            print(
                f"Aligned lyrics API returned failure state '{state}' "
                f"for {suno_clip_id}"
            )
            return None

        # Not complete yet; wait and try again
        if attempt < max_attempts:
            time.sleep(delay_seconds)
        else:
            print(
                f"Aligned lyrics polling for {suno_clip_id} did not "
                f"reach 'complete' after {max_attempts} attempts."
            )

    return None


@https_fn.on_call(secrets=[suno_session_token])
def update_song_status(req: https_fn.CallableRequest) -> Dict[str, Any]:
    """
    Callable function to manually trigger a song status poll and update.
    This can be called from the client to immediately check for song updates.

    Returns:
        Dictionary with polling results (songs updated, errors, etc.)
    """
    print("Manual song status update triggered")
    result = poll_and_update_songs(suno_session_token.value)
    return result


@scheduler_fn.on_schedule(schedule="* * * * *", secrets=[suno_session_token])
def poll_song_status(event: scheduler_fn.ScheduledEvent) -> None:
    """
    Scheduled function that polls Suno API for song status updates.
    Runs every minute via Cloud Scheduler, then polls every 3 seconds within that minute.
    """
    import time

    start_time = datetime.now()
    print(f"Poll cycle started at {start_time.isoformat()}")

    iteration = 0

    # Run polling loop for up to 1 minute
    while True:
        iteration += 1
        current_time = datetime.now()
        elapsed = (current_time - start_time).total_seconds()

        # Check if we've exceeded 1 minute
        if elapsed > 53:
            print(f"Poll cycle ending after {elapsed:.1f}s ({iteration} iterations)")
            break

        print(f"[Iteration {iteration} at {elapsed:.1f}s]")

        # Call the extracted polling logic
        poll_and_update_songs(suno_session_token.value)

        # Sleep for 3 seconds before next iteration
        time.sleep(3)


@firestore_fn.on_document_updated(
    document="song/{songId}", secrets=[suno_session_token]
)
def ensure_hooks_start_point_on_like(
    event: firestore_fn.Event[
        firestore_fn.Change[firestore_fn.DocumentSnapshot]
    ],
) -> None:
    """
    Firestore trigger that ensures certain songs have a hooksStartPoint.

    When a song document is updated such that either:
      - isLikedByCreator is set to True, or
      - isUpload is set to True,
    and hooksStartPoint is not already present, this function will call the Suno
    aligned lyrics API to determine the start of the first chorus and store it
    under hooksStartPoint.
    """
    before_snap = event.data.before
    after_snap = event.data.after

    if before_snap is None or after_snap is None:
        return

    before_data = before_snap.to_dict() or {}
    after_data = after_snap.to_dict() or {}

    before_liked = before_data.get("isLikedByCreator")
    after_liked = after_data.get("isLikedByCreator")
    before_upload = before_data.get("isUpload")
    after_upload = after_data.get("isUpload")

    # Only act when either:
    # - isLikedByCreator has just been set to True, or
    # - isUpload has just been set to True
    liked_just_set_true = bool(after_liked) and not bool(before_liked)
    upload_just_set_true = bool(after_upload) and not bool(before_upload)

    if not (liked_just_set_true or upload_just_set_true):
        return

    # If hooksStartPoint is already present on the document, do nothing
    if "hooksStartPoint" in after_data:
        return

    suno_clip_id = after_data.get("sunoClipId")
    if not suno_clip_id:
        print(
            f"Song {after_snap.reference.path} liked by creator but has no sunoClipId; "
            f"skipping hooksStartPoint computation."
        )
        return

    hooks_start = get_hooks_start_point_for_clip(
        suno_clip_id, suno_session_token.value, after_data.get("lyrics")
    )

    if hooks_start is None:
        print(
            f"Could not determine hooksStartPoint for "
            f"{after_snap.reference.path} (clip {suno_clip_id})."
        )
        return

    try:
        after_snap.reference.update({"hooksStartPoint": hooks_start})
        print(
            f"Set hooksStartPoint={hooks_start} on "
            f"{after_snap.reference.path} (clip {suno_clip_id})."
        )
    except Exception as e:
        print(
            f"Failed to update hooksStartPoint for "
            f"{after_snap.reference.path} (clip {suno_clip_id}): {e}"
        )


@https_fn.on_request()
def migrate_has_been_remixed(req: https_fn.Request) -> https_fn.Response:
    # CORS preflight
    if req.method == "OPTIONS":
        resp = https_fn.Response("", status=204)
        resp.headers["Access-Control-Allow-Origin"] = "*"
        resp.headers["Access-Control-Allow-Methods"] = "POST, OPTIONS"
        resp.headers["Access-Control-Allow-Headers"] = "Content-Type, Authorization"
        return resp

    if req.method != "POST":
        return https_fn.Response(
            json.dumps({"error": "Use POST to run this migration."}),
            status=405,
            content_type="application/json",
            headers={"Access-Control-Allow-Origin": "*"},
        )

    try:
        firestore_client: google.cloud.firestore.Client = firestore.client()

        # Stream all generated songs (not uploads)
        generated_songs_iter = (
            firestore_client.collection("song").where("isUpload", "==", False).stream()
        )

        # Collect unique root remix song references
        root_path_to_ref: Dict[str, google.cloud.firestore.DocumentReference] = {}
        total_generated: int = 0
        with_root_field: int = 0
        for song_snap in generated_songs_iter:
            total_generated += 1
            data = song_snap.to_dict()
            root_ref = data.get("rootRemixSong")
            if root_ref:
                with_root_field += 1
                root_path_to_ref[root_ref.path] = root_ref

        # Batch update roots to set hasBeenRemixed = True (skip deleted roots)
        updated_count = 0
        pending_refs = list(root_path_to_ref.values())
        BATCH_SIZE = 500

        while pending_refs:
            chunk = pending_refs[:BATCH_SIZE]
            pending_refs = pending_refs[BATCH_SIZE:]

            # Fetch snapshots for this chunk to ensure they exist
            snapshots = list(firestore_client.get_all(chunk))

            batch = firestore_client.batch()
            writes_in_batch = 0
            for snap in snapshots:
                if snap.exists:
                    batch.update(snap.reference, {"hasBeenRemixed": True})
                    writes_in_batch += 1
            if writes_in_batch > 0:
                batch.commit()
                updated_count += writes_in_batch

        return https_fn.Response(
            json.dumps(
                {
                    "message": "Migration completed",
                    "totalGeneratedSongs": total_generated,
                    "generatedWithRootRemixSong": with_root_field,
                    "uniqueRootsUpdated": updated_count,
                }
            ),
            status=200,
            content_type="application/json",
            headers={"Access-Control-Allow-Origin": "*"},
        )
    except Exception as e:
        return https_fn.Response(
            json.dumps({"error": str(e)}),
            status=500,
            content_type="application/json",
            headers={"Access-Control-Allow-Origin": "*"},
        )


@https_fn.on_request(secrets=[suno_session_token], timeout_sec=300)
def migrate_hooks_start_point(req: https_fn.Request) -> https_fn.Response:
    """
    HTTP function to backfill hooksStartPoint for existing songs.

    This scans for songs where isLikedByCreator == True and hooksStartPoint is
    not yet present, then computes and stores hooksStartPoint using the Suno
    aligned lyrics API.
    """
    # CORS preflight
    if req.method == "OPTIONS":
        resp = https_fn.Response("", status=204)
        resp.headers["Access-Control-Allow-Origin"] = "*"
        resp.headers["Access-Control-Allow-Methods"] = "POST, OPTIONS"
        resp.headers["Access-Control-Allow-Headers"] = "Content-Type, Authorization"
        return resp

    if req.method != "POST":
        return https_fn.Response(
            json.dumps({"error": "Use POST to run this migration."}),
            status=405,
            content_type="application/json",
            headers={"Access-Control-Allow-Origin": "*"},
        )

    try:
        firestore_client: google.cloud.firestore.Client = firestore.client()

        # Limit the number of documents processed per invocation to avoid
        # timeouts. Run the migration multiple times if needed. The limit is
        # applied per query (liked songs, uploads), so an invocation may
        # process up to roughly 2 * BATCH_LIMIT unique documents.
        BATCH_LIMIT = 500

        # Collect songs that either:
        #   - are liked by the creator, or
        #   - are uploads,
        # and deduplicate by document path.
        song_snaps: Dict[str, google.cloud.firestore.DocumentSnapshot] = {}

        liked_iter = (
            firestore_client.collection("song")
            .where("isLikedByCreator", "==", True)
            .limit(BATCH_LIMIT)
            .stream()
        )
        for snap in liked_iter:
            song_snaps[snap.reference.path] = snap

        upload_iter = (
            firestore_client.collection("song")
            .where("isUpload", "==", True)
            .limit(BATCH_LIMIT)
            .stream()
        )
        for snap in upload_iter:
            song_snaps.setdefault(snap.reference.path, snap)

        processed = 0
        updated = 0
        skipped_already_set = 0
        skipped_missing_clip = 0
        skipped_no_chorus = 0
        errors = 0

        for song_snap in song_snaps.values():
            processed += 1
            data = song_snap.to_dict() or {}

            # If hooksStartPoint is already present, assume this doc has
            # already been processed.
            if "hooksStartPoint" in data and isinstance(data.get("hooksStartPoint"), float):
                skipped_already_set += 1
                continue

            suno_clip_id = data.get("sunoClipId")

            try:
                hooks_start = get_hooks_start_point_for_clip(
                    suno_clip_id, suno_session_token.value, data.get("lyrics")
                )
            except Exception as e:
                print(
                    f"Error computing hooksStartPoint for "
                    f"{song_snap.reference.path} (clip {suno_clip_id}): {e}"
                )
                errors += 1
                continue

            # If we had no clip ID and the helper still returned None, we can't
            # proceed (this should only happen when we actually needed to call
            # the aligned lyrics API).
            if hooks_start is None and not suno_clip_id:
                skipped_missing_clip += 1
                continue

            if hooks_start is None:
                skipped_no_chorus += 1
                continue

            try:
                song_snap.reference.update({"hooksStartPoint": hooks_start})
                updated += 1
            except Exception as e:
                print(
                    f"Failed to update hooksStartPoint for "
                    f"{song_snap.reference.path} (clip {suno_clip_id}): {e}"
                )
                errors += 1

        return https_fn.Response(
            json.dumps(
                {
                    "message": "hooksStartPoint migration completed for this batch",
                    "batchLimit": BATCH_LIMIT,
                    "processed": processed,
                    "updated": updated,
                    "skippedAlreadySet": skipped_already_set,
                    "skippedMissingClipId": skipped_missing_clip,
                    "skippedNoChorusFound": skipped_no_chorus,
                    "errors": errors,
                }
            ),
            status=200,
            content_type="application/json",
            headers={"Access-Control-Allow-Origin": "*"},
        )
    except Exception as e:
        return https_fn.Response(
            json.dumps({"error": str(e)}),
            status=500,
            content_type="application/json",
            headers={"Access-Control-Allow-Origin": "*"},
        )