src/main.py

import os
import json
import base64
import logging
import re
import subprocess
import tempfile
import time
import http.client
import urllib.request
import concurrent.futures
from urllib.parse import urlencode
from typing import Optional, Dict, Any, List, Callable

import replicate
from openai import OpenAI
from pydantic import BaseModel

from configurations import cnf

client = OpenAI(api_key=cnf.OPENAI_API_KEY)

if getattr(cnf, "REPLICATE_API_TOKEN", None):
    os.environ["REPLICATE_API_TOKEN"] = cnf.REPLICATE_API_TOKEN

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s"
)

logger = logging.getLogger("ad_video_pipeline")


class SceneDirection(BaseModel):
    scene_number: int
    time_range: str
    visual: str
    avatar_action: str
    product_placement: str
    camera_direction: str
    on_screen_text: str


class ScriptVisualDirectionResult(BaseModel):
    duration_seconds: int
    hook: str
    voiceover_script: str
    cta: str
    visual_style: str
    scene_breakdown: List[SceneDirection]
    seedance_prompt: str
    negative_prompt: str


class VoiceoverScriptRewrite(BaseModel):
    voiceover_script: str


class StoryboardShot(BaseModel):
    shot_number: int
    time_range: str
    title: str
    camera: str
    description: str
    caption: str


class Storyboard(BaseModel):
    ad_title: str
    duration_seconds: int
    mood: str
    aspect_ratio: str
    shots: List[StoryboardShot]


def as_url_list(value) -> List[str]:
    if value is None:
        return []
    if isinstance(value, str):
        return [value]
    return list(value)


def to_public_uri(url: Optional[str]) -> Optional[str]:
    if not url:
        return None

    if url.startswith(("http://", "https://")):
        return url

    raise ValueError(
        "This pipeline requires public http/https URLs. "
        f"Got invalid value: {url}"
    )


def extract_voice_id(voice_clone_output: Any) -> str:
    logger.info("[VOICE CLONE] Raw output type: %s", type(voice_clone_output))

    if isinstance(voice_clone_output, str):
        return voice_clone_output

    if isinstance(voice_clone_output, dict):
        for key in ["voice_id", "id", "output"]:
            value = voice_clone_output.get(key)

            if isinstance(value, str) and value:
                return value

            if isinstance(value, dict):
                nested_voice_id = value.get("voice_id")
                if isinstance(nested_voice_id, str) and nested_voice_id:
                    return nested_voice_id

    if isinstance(voice_clone_output, list) and voice_clone_output:
        return extract_voice_id(voice_clone_output[0])

    raise ValueError(
        f"Could not extract voice_id from MiniMax voice clone output: {voice_clone_output}"
    )


def extract_output_url(output: Any) -> str:
    final_output = output[0] if isinstance(output, list) else output

    if isinstance(final_output, str):
        return final_output

    output_url = getattr(final_output, "url", None)

    if output_url:
        return output_url

    raise ValueError(f"Could not extract output URL from Replicate output: {output}")


def generate_script_and_visual_direction(
    product_image_url: str,
    angle: Optional[str] = None,
    concept: Optional[str] = None,
    avatar_description: Optional[str] = None,
    voice_description: Optional[str] = None,
    duration_seconds: int = 15,
    model: str = "gpt-5.5",
    manual_script: Optional[str] = None,
) -> Dict[str, Any]:

    logger.info("[STEP 1] Script + visual direction generation started")

    product_image_urls = as_url_list(product_image_url)

    hook_end = round(duration_seconds * 0.2)
    product_end = round(duration_seconds * 0.47)
    benefit_end = round(duration_seconds * 0.73)
    payoff_end = round(duration_seconds * 0.93)

    # ~2.5 words/second is a natural ad-voiceover pace. Without a target word
    # count, the model tends to write a script that's merely "short enough"
    # for the duration rather than one that fills it — leaving the back half
    # of longer (chained) videos silent.
    target_words = round(duration_seconds * 2.5)
    min_words = round(target_words * 0.85)
    max_words = round(target_words * 1.05)

    manual_script_block = (
        "Voiceover Script (use this exact script verbatim — do not rewrite, shorten, or "
        "paraphrase it; build the hook, CTA, and scene-by-scene visual direction around it):\n"
        f'"""\n{manual_script}\n"""\n'
    ) if manual_script else ""

    prompt = f"""
You are an expert performance marketing creative strategist and product image analyst.

Analyze the product image(s) and create a short video ad script and visual direction for Seedance video generation.

Marketing Angle:
{angle or "Not provided"}

Creative Concept:
{concept or "Not provided"}

Avatar Description:
{avatar_description or "Not provided"}

Voice Description:
{voice_description or "Not provided"}

Video Duration:
{duration_seconds} seconds

{manual_script_block}
First infer from the product image:
- product category
- product type
- visible features
- likely use case
- likely target audience
- main benefit

Then generate:
- hook
- voiceover script{" (reuse the exact script given above verbatim)" if manual_script else ""}
- CTA
- visual style
- scene-by-scene visual direction
- final Seedance prompt
- negative prompt

Seedance prompt rules:
- Refer to the avatar/person and the product generically (e.g. "the avatar", "the product") — do not assume specific [ImageN] index numbers, those are assigned separately based on the actual reference images supplied.
- Use [Audio1] as the final generated voiceover timing and pacing guide.
- Keep the avatar face, outfit, and body proportions consistent.
- Keep the product shape, packaging, color, and branding consistent.
- Show the product clearly.
- Make the ad realistic and social-media ready.
- Do not change the product design.
- Do not change the avatar face or outfit/clothing.
- When describing the CTA text overlay, quote it verbatim in double quotes (e.g. text overlay reads "SHOP NOW") and reuse that exact spelling everywhere it appears. Keep it short (under 5 words) so it renders legibly and spelled correctly.
- Do not caption or subtitle the voiceover on screen — the voiceover script is heard only, never shown as text. The CTA is the only on-screen text, shown once near the end.

Ad structure for {duration_seconds} seconds:
0-{hook_end}s: scroll-stopping problem hook.
{hook_end}-{product_end}s: product appears as the solution.
{product_end}-{benefit_end}s: demonstrate the key benefit.
{benefit_end}-{payoff_end}s: emotional payoff or trust point.
{payoff_end}-{duration_seconds}s: clear CTA, shown in the last shot of the video, not before.

Rules:
{'- Use the voiceover script given above exactly as written — do not alter it.' if manual_script else f'''- Write the voiceover script to fill the full {duration_seconds} seconds at a
  natural speaking pace — that's approximately {min_words}-{max_words} words.
  Do not write a noticeably shorter script that leaves the back half of the
  video silent; do not run over {max_words} words either.'''}
- Use simple, direct language.
- Do not make unverifiable claims.
- Do not mention discounts unless provided.
- Avoid medical, legal, or guaranteed claims.
- Keep visuals simple enough for a video generation model.
"""

    try:
        response = client.responses.parse(
            model=model,
            input=[
                {
                    "role": "system",
                    "content": (
                        "You generate structured direct-response ad scripts and video prompts "
                        "for Bytedance Seedance 2.0. Return only output matching the schema."
                    )
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "input_text",
                            "text": prompt
                        },
                        *[
                            {"type": "input_image", "image_url": url}
                            for url in product_image_urls
                        ]
                    ]
                }
            ],
            text_format=ScriptVisualDirectionResult,
        )

        creative = response.output_parsed.model_dump()

        if manual_script:
            creative["voiceover_script"] = manual_script

        logger.info("[STEP 1] Script + visual direction generation completed")
        logger.info("Hook: %s", creative.get("hook"))
        logger.info("Voiceover script: %s", creative.get("voiceover_script"))
        logger.info(
            "Seedance prompt length: %s",
            len(creative.get("seedance_prompt", ""))
        )

        return creative

    except Exception:
        logger.exception("[STEP 1] Script + visual direction generation failed")
        raise


def rewrite_voiceover_script_for_length(
    creative: Dict[str, Any],
    target_words: int,
    model: str = "gpt-5.5",
) -> str:
    logger.info("[STEP 2C] Rewriting voiceover script to hit ~%s words", target_words)

    prompt = f"""
This voiceover script was written for a video ad, but the synthesized speech
runs shorter than the video's actual duration — it needs to be longer.
Rewrite it to be approximately {target_words} words (natural, conversational
ad copy — do not pad with filler, repeated lines, or slow it down artificially;
add genuine additional content: more benefit detail, a supporting sentence,
etc.), while keeping the same hook, message, tone, and CTA meaning.

Original script:
{creative.get("voiceover_script")}

Hook: {creative.get("hook")}
CTA: {creative.get("cta")}
"""

    try:
        response = client.responses.parse(
            model=model,
            input=[
                {
                    "role": "system",
                    "content": (
                        "You rewrite ad voiceover scripts to hit a target word count while "
                        "preserving meaning and tone. Return only output matching the schema."
                    )
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            text_format=VoiceoverScriptRewrite,
        )

        rewritten = response.output_parsed.voiceover_script

        logger.info("[STEP 2C] Voiceover script rewritten (%s words)", len(rewritten.split()))

        return rewritten

    except Exception:
        logger.exception("[STEP 2C] Voiceover script rewrite failed")
        raise


def ensure_voiceover_fills_duration(
    creative: Dict[str, Any],
    generated_audio_url: str,
    voice_id: str,
    duration_seconds: int,
    script_model: str = "gpt-5.5",
    max_attempts: int = 3,
    allow_rewrite: bool = True,
) -> str:
    # TTS speaking pace varies and can undershoot the script's assumed
    # words/second, leaving the tail of the video silent. A single correction
    # pass isn't guaranteed to land within tolerance either (the rewrite is
    # itself just an LLM guess at a word count) — so this re-measures the
    # actual synthesized audio after each attempt and keeps correcting
    # (reusing the same voice_id, no need to re-clone) until it's close
    # enough or attempts run out. Skipped entirely for a manually written
    # script — rewriting it would defeat the point of writing it by hand.
    if not allow_rewrite:
        return generated_audio_url

    for attempt in range(1, max_attempts + 1):
        actual_audio_seconds = _get_audio_duration_seconds(download_image_bytes(generated_audio_url))

        if actual_audio_seconds >= duration_seconds * 0.9:
            return generated_audio_url

        if attempt == max_attempts:
            logger.warning(
                "Voiceover audio still short (%.1fs vs %ss) after %s attempts — proceeding anyway",
                actual_audio_seconds, duration_seconds, max_attempts,
            )
            return generated_audio_url

        logger.info(
            "Voiceover audio is %.1fs, short of the requested %ss (attempt %s/%s) — regenerating with a longer script",
            actual_audio_seconds, duration_seconds, attempt, max_attempts,
        )

        original_words = len(creative["voiceover_script"].split())
        adjusted_target_words = round(original_words * (duration_seconds / actual_audio_seconds))

        creative["voiceover_script"] = rewrite_voiceover_script_for_length(
            creative=creative,
            target_words=adjusted_target_words,
            model=script_model,
        )
        generated_audio_url = generate_speech_with_minimax(
            text=creative["voiceover_script"],
            voice_id=voice_id,
        )

    return generated_audio_url


def generate_storyboard(
    creative: Dict[str, Any],
    duration_seconds: int,
    model: str = "gpt-5.5",
    window_start: Optional[int] = None,
    window_end: Optional[int] = None,
    prior_shots: Optional[List[Dict[str, Any]]] = None,
    aspect_ratio: str = "9:16",
) -> Dict[str, Any]:

    logger.info("[STEP 1B] Storyboard generation started")

    is_scoped = window_start is not None and window_end is not None
    shots_start = window_start if is_scoped else 0
    shots_end = window_end if is_scoped else duration_seconds
    is_final_window = shots_end >= duration_seconds

    scope_block = ""
    if is_scoped:
        prior_block = (
            "\n".join(f"- {s['title']}: {s['description']}" for s in prior_shots)
            if prior_shots else "(none — this is the first segment)"
        )
        scope_block = f"""
This call covers ONLY the {shots_start}-{shots_end} second window of the full
{duration_seconds}-second ad, not the whole ad. Give shot time_range values as
absolute seconds within the full {duration_seconds}s timeline (e.g. shots in
this window look like "{shots_start}-{shots_start + 3}s", not restarting at 0).

Shots already covered in earlier segments — do not repeat these beats,
continue the story forward from here:
{prior_block}
"""

    cta_rule = (
        f"The CTA must only appear as on-screen text in the final shot, ending "
        f"exactly at {duration_seconds}s — never earlier."
        if is_final_window else
        "Do not show the CTA in this segment — it belongs only in the ad's final segment."
    )

    orientation = {
        "9:16": "vertical", "3:4": "vertical",
        "16:9": "horizontal", "21:9": "horizontal", "4:3": "horizontal",
        "1:1": "square",
    }.get(aspect_ratio, "vertical")

    prompt = f"""
You are a commercial video storyboard artist. Break the ad script and visual
direction below into a shot-by-shot storyboard for a {duration_seconds}-second
{orientation} ({aspect_ratio}) social video, in the style of a professional shot list.
{scope_block}
Hook: {creative.get("hook")}
Voiceover script: {creative.get("voiceover_script")}
Visual style: {creative.get("visual_style")}
Scene breakdown: {json.dumps(creative.get("scene_breakdown", []))}
Seedance prompt: {creative.get("seedance_prompt")}

Produce enough shots to cover the {shots_end - shots_start} seconds of {"this window" if is_scoped else "the ad"}
(roughly one shot per 2-3 seconds), with shot time ranges that are contiguous
and cover {shots_start} to {shots_end} with no gaps or overlaps. For each shot give:
- a time_range in exact "START-ENDs" format (e.g. "0-3s"), matching the shot's
  place in the overall {duration_seconds}-second timeline
- a short punchy title (2-5 words, ALL CAPS style)
- a camera direction line (framing, movement, angle, handheld/steady) written
  like a director's note
- a one-sentence description of the action/visual in the shot
- a one-sentence caption explaining why this shot matters to the ad

{cta_rule}

Keep each shot's hand action simple and singular — the avatar's hands should
do at most one thing at a time (e.g. hold the product, OR adjust an outfit,
OR gesture). Never describe the avatar holding more than one object at once
or performing two distinct hand actions simultaneously (e.g. "holding two
necklaces while checking a watch") — these poses are hard to render correctly
and tend to produce extra or malformed hands.

Also give an overall ad_title, mood (3-4 words), and aspect_ratio ("{aspect_ratio}").
"""

    try:
        response = client.responses.parse(
            model=model,
            input=[
                {
                    "role": "system",
                    "content": (
                        "You produce structured shot-by-shot video storyboards. "
                        "Return only output matching the schema."
                    )
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            text_format=Storyboard,
        )

        storyboard = response.output_parsed.model_dump()

        log_storyboard(storyboard)

        logger.info("[STEP 1B] Storyboard generation completed")

        return storyboard

    except Exception:
        logger.exception("[STEP 1B] Storyboard generation failed")
        raise


def log_storyboard(storyboard: Dict[str, Any]) -> None:
    logger.info("[STORYBOARD] Ad title: %s", storyboard.get("ad_title"))
    logger.info(
        "[STORYBOARD] Duration: %ss | Mood: %s | Aspect ratio: %s",
        storyboard.get("duration_seconds"),
        storyboard.get("mood"),
        storyboard.get("aspect_ratio"),
    )
    logger.info("[STORYBOARD] Generated %s shots", len(storyboard.get("shots", [])))

    for shot in storyboard.get("shots", []):
        logger.info(
            "[STORYBOARD] Shot %s - %s | Camera: %s | %s | (%s)",
            shot.get("shot_number"),
            shot.get("title"),
            shot.get("camera"),
            shot.get("description"),
            shot.get("caption"),
        )


def download_image_bytes(url: str) -> bytes:
    request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(request, timeout=30) as resp:
        return resp.read()


def generate_storyboard_poster(
    storyboard: Dict[str, Any],
    avatar_image_url: str,
    product_image_url: str,
    model: str = "gpt-image-2",
) -> str:
    logger.info("[STEP 1B] Generating full storyboard poster in one request (%s)", model)

    avatar_urls = as_url_list(avatar_image_url)
    product_urls = as_url_list(product_image_url)

    avatar_bytes_list = [download_image_bytes(url) for url in avatar_urls]
    product_bytes_list = [download_image_bytes(url) for url in product_urls]

    avatar_ref_range = "reference image 1" if len(avatar_urls) == 1 else f"reference images 1-{len(avatar_urls)}"
    product_start = len(avatar_urls) + 1
    product_end = len(avatar_urls) + len(product_urls)
    product_ref_range = (
        f"reference image {product_start}" if len(product_urls) == 1
        else f"reference images {product_start}-{product_end}"
    )

    shots = storyboard["shots"]
    panels = "\n".join(
        f'Panel {shot["shot_number"]} — "{shot["title"]}": {shot["description"]} '
        f'Camera: {shot["camera"]} Caption: {shot["caption"]}'
        for shot in shots
    )

    prompt = (
        f'Design one wide storyboard reference sheet on a dark background, in the style of '
        f'a professional shot-list infographic. Header bar reads "{storyboard["ad_title"].upper()} '
        f'— {len(shots)} SHOTS". Arrange all {len(shots)} panels in a grid (3 per row), each '
        "with a bold yellow numbered badge, a short bold title, a photo depicting that panel's "
        "action, a yellow \"CAMERA:\" line with the camera direction below the photo, and a "
        "muted caption line below that. Only the pose, action, and camera framing change per "
        "panel to match its description — everything else about the person and product must "
        "stay identical to the reference images in every panel:\n"
        f"- Person ({avatar_ref_range}): same face, same facial features, same skin tone, "
        "same hair, same body proportions and height, same outfit — do not regenerate a "
        "different-looking person, even in close-up or cropped panels. If there is more than "
        "one person reference image, they show the same person from different angles.\n"
        f"- Product ({product_ref_range}): same design, shape, color, materials, and exact "
        "proportions (e.g. pendant size relative to chain, gemstone count and layout, chain "
        "length) — do not redesign, resize, or simplify the product. A macro/close-up panel may "
        "fill more of the frame with the product because the camera is closer, but the product "
        "itself, at its real-world scale relative to the person's hand or neck, must not be "
        "drawn larger, smaller, or differently designed than the reference image.\n"
        "- Anatomy: every person must have exactly two hands and two arms, five fingers per "
        "hand, in correct proportion — no extra, duplicated, merged, or malformed hands/limbs "
        "in any panel, even in busy or cropped compositions.\n\n"
        f'Footer bar reads "DURATION: ~{storyboard["duration_seconds"]} SECONDS | MOOD: '
        f'{storyboard["mood"].upper()} | ASPECT RATIO: {storyboard["aspect_ratio"]}".\n\n'
        f"Panels:\n{panels}"
    )

    try:
        response = client.images.edit(
            image=[
                (f"avatar{i}.jpg", b, "image/jpeg")
                for i, b in enumerate(avatar_bytes_list)
            ] + [
                (f"product{i}.jpg", b, "image/jpeg")
                for i, b in enumerate(product_bytes_list)
            ],
            prompt=prompt,
            model=model,
            size="1536x1024",
            n=1,
        )

        poster_image_url = f"data:image/png;base64,{response.data[0].b64_json}"

        logger.info("[STEP 1B] Storyboard poster generated")

        return poster_image_url

    except Exception:
        logger.exception("[STEP 1B] Storyboard poster generation failed")
        raise


def clone_voice_with_minimax(
    voice_public_url: str,
    need_noise_reduction: bool = False
) -> str:

    logger.info("[STEP 2A] MiniMax voice cloning started")
    logger.info("[STEP 2A] Voice file URL: %s", voice_public_url)

    voice_public_url = to_public_uri(voice_public_url)

    try:
        output = replicate.run(
            "minimax/voice-cloning",
            input={
                "voice_file": voice_public_url,
                "need_noise_reduction": need_noise_reduction
            }
        )

        voice_id = extract_voice_id(output)

        logger.info("[STEP 2A] MiniMax voice cloning completed")
        logger.info("[STEP 2A] Voice ID: %s", voice_id)

        return voice_id

    except Exception:
        logger.exception("[STEP 2A] MiniMax voice cloning failed")
        raise


def generate_speech_with_minimax(
    text: str,
    voice_id: str,
) -> str:

    logger.info("[STEP 2B] MiniMax speech generation started")
    logger.info("[STEP 2B] Voice ID: %s", voice_id)
    logger.info("[STEP 2B] Text length: %s chars", len(text))

    try:
        output = replicate.run(
            "minimax/speech-2.8-turbo",
            input={
                "text": text,
                "voice_id": voice_id
            }
        )

        generated_audio_url = extract_output_url(output)

        logger.info("[STEP 2B] MiniMax speech generation completed")
        logger.info("[STEP 2B] Generated audio URL: %s", generated_audio_url)

        return generated_audio_url

    except Exception:
        logger.exception("[STEP 2B] MiniMax speech generation failed")
        raise


def generate_audio_from_voice_url(
    voiceover_script: str,
    voice_public_url: str
) -> Dict[str, str]:

    logger.info("Full MiniMax audio generation started")
    logger.info("Voice public URL: %s", voice_public_url)

    try:
        if voice_public_url.startswith(("http://", "https://")):
            voice_id = clone_voice_with_minimax(
                voice_public_url=voice_public_url,
                need_noise_reduction=False
            )
        else:
            voice_id = voice_public_url
            logger.info("[STEP 2A] Skipping voice cloning, using existing voice ID: %s", voice_id)

        generated_audio_url = generate_speech_with_minimax(
            text=voiceover_script,
            voice_id=voice_id,
        )

        logger.info("Full MiniMax audio generation completed")
        logger.info("Voice ID: %s", voice_id)
        logger.info("Generated audio URL: %s", generated_audio_url)

        return {
            "voice_id": voice_id,
            "generated_audio_url": generated_audio_url
        }

    except Exception:
        logger.exception("Full MiniMax audio generation failed")
        raise


KIE_HOST = "api.kie.ai"
KIE_CREATE_TASK_PATH = "/api/v1/jobs/createTask"
KIE_RECORD_INFO_PATH = "/api/v1/jobs/recordInfo"

# File upload lives on a separate Kie host from the jobs API above.
KIE_UPLOAD_HOST = "kieai.redpandaai.co"
KIE_FILE_BASE64_UPLOAD_PATH = "/api/file-base64-upload"

# Burned into every generated video (bottom-right corner) — fixed brand
# requirement, not something the frontend exposes.
AMALFA_LOGO_URL = "https://creative-library-beta.static.fabfunnel.com/330/fabAds/537/56b468f6-6720-4db1-b15e-725a623ddb84.jpg"


def get_kie_api_token() -> str:
    token = (
        getattr(cnf, "KIE_API_TOKEN", None)
        or getattr(cnf, "KIE_API_KEY", None)
        or os.getenv("KIE_API_TOKEN")
        or os.getenv("KIE_API_KEY")
    )

    if not token:
        raise ValueError(
            "Missing Kie API token. Add KIE_API_TOKEN or KIE_API_KEY to your config/env."
        )

    return token


def kie_api_request(
    method: str,
    path: str,
    payload: Optional[Dict[str, Any]] = None,
    query: Optional[Dict[str, Any]] = None,
    host: str = KIE_HOST,
) -> Dict[str, Any]:

    if query:
        path = f"{path}?{urlencode(query)}"

    body = json.dumps(payload) if payload is not None else None

    headers = {
        "Authorization": f"Bearer {get_kie_api_token()}",
    }

    if payload is not None:
        headers["Content-Type"] = "application/json"

    conn = http.client.HTTPSConnection(host, timeout=60)

    try:
        conn.request(method, path, body=body, headers=headers)
        res = conn.getresponse()
        raw = res.read().decode("utf-8")
    finally:
        conn.close()

    try:
        data = json.loads(raw) if raw else {}
    except json.JSONDecodeError:
        data = {"raw": raw}

    if res.status >= 400:
        raise RuntimeError(
            f"Kie API HTTP {res.status} {res.reason}: "
            f"{json.dumps(data, ensure_ascii=False)}"
        )

    if isinstance(data, dict):
        success = data.get("success")

        # Different Kie endpoints report success differently: file-upload
        # returns a "success" bool with its own msg text, while
        # createTask/recordInfo have no "success" field and rely on
        # msg == "success". Prefer the explicit bool when present.
        if success is not None:
            if not success:
                raise RuntimeError(f"Kie API error: {json.dumps(data, ensure_ascii=False)}")
        elif data.get("msg") not in (None, "success"):
            raise RuntimeError(f"Kie API error: {json.dumps(data, ensure_ascii=False)}")

    return data


KIE_IMAGE_MIN_SIDE = 300
KIE_IMAGE_MAX_SIDE = 4096
KIE_IMAGE_MAX_BYTES = 10 * 1024 * 1024


def _get_image_dimensions(image_bytes: bytes, suffix: str) -> tuple:
    with tempfile.TemporaryDirectory() as tmpdir:
        path = os.path.join(tmpdir, f"img{suffix}")
        with open(path, "wb") as f:
            f.write(image_bytes)

        result = subprocess.run(
            [
                "ffprobe", "-v", "error", "-select_streams", "v:0",
                "-show_entries", "stream=width,height",
                "-of", "csv=p=0", path,
            ],
            check=True, capture_output=True, text=True,
        )

        w, h = result.stdout.strip().split(",")[:2]
        return int(w), int(h)


def _ffmpeg_resize_jpeg(image_bytes: bytes, suffix: str, width: int, height: int, quality: int) -> bytes:
    with tempfile.TemporaryDirectory() as tmpdir:
        in_path = os.path.join(tmpdir, f"in{suffix}")
        out_path = os.path.join(tmpdir, "out.jpg")

        with open(in_path, "wb") as f:
            f.write(image_bytes)


        subprocess.run(
            ["ffmpeg", "-y", "-i", in_path, "-vf", f"scale={width}:{height}", "-q:v", str(quality), out_path],
            check=True, capture_output=True,
        )

        with open(out_path, "rb") as f:
            return f.read()


def normalize_image_for_kie(file_bytes: bytes, file_name: str) -> bytes:

    suffix = os.path.splitext(file_name)[1].lower() or ".jpg"

    try:
        w, h = _get_image_dimensions(file_bytes, suffix)
    except Exception:
        logger.exception("[IMAGE] Could not read dimensions of %s — uploading as-is", file_name)
        return file_bytes

    long_side, short_side = max(w, h), min(w, h)
    scale = 1.0
    if long_side > KIE_IMAGE_MAX_SIDE:
        scale = KIE_IMAGE_MAX_SIDE / long_side
    if short_side * scale < KIE_IMAGE_MIN_SIDE:
        # ponytail: assumes aspect ratio isn't so extreme that hitting the 300px
        # min pushes the long side back over 6000 (>20:1); Seedance rejects
        # those anyway. Clamp here if that ever shows up.
        scale = KIE_IMAGE_MIN_SIDE / short_side

    # ffmpeg's scale filter needs even dimensions for some encoders; round to even.
    new_w = max(2, round(w * scale / 2) * 2)
    new_h = max(2, round(h * scale / 2) * 2)

    already_ok = (
        scale == 1.0
        and len(file_bytes) <= KIE_IMAGE_MAX_BYTES
        and suffix in (".jpg", ".jpeg")
    )
    if already_ok:
        return file_bytes

    result = _ffmpeg_resize_jpeg(file_bytes, suffix, new_w, new_h, quality=3)

    # Still too big (rare — a huge, busy image at the max dimension): step the
    # JPEG quality down until it fits rather than shrink the image further.
    for quality in (7, 15, 25):
        if len(result) <= KIE_IMAGE_MAX_BYTES:
            break
        result = _ffmpeg_resize_jpeg(file_bytes, suffix, new_w, new_h, quality=quality)

    logger.info(
        "[IMAGE] Normalized %s: %sx%s (%.1fMB) -> %sx%s (%.1fMB)",
        file_name, w, h, len(file_bytes) / 1e6, new_w, new_h, len(result) / 1e6,
    )

    return result


def upload_image_to_kie(file_bytes: bytes, file_name: str, mime_type: str = "image/jpeg") -> str:
    return upload_file_to_kie(normalize_image_for_kie(file_bytes, file_name), file_name, "image/jpeg")


def upload_file_to_kie(file_bytes: bytes, file_name: str, mime_type: str = "image/jpeg") -> str:
    base64_data_uri = f"data:{mime_type};base64,{base64.b64encode(file_bytes).decode('ascii')}"

    response = kie_api_request(
        "POST",
        KIE_FILE_BASE64_UPLOAD_PATH,
        payload={
            "base64Data": base64_data_uri,
            "uploadPath": "avatar-ad-video",
            "fileName": file_name,
        },
        host=KIE_UPLOAD_HOST,
    )

    download_url = response.get("data", {}).get("downloadUrl")

    if not download_url:
        raise RuntimeError(f"Kie file upload did not return a downloadUrl: {response}")

    return download_url


def create_kie_seedance_task(
    prompt: str,
    avatar_image_url: Optional[str] = None,
    product_image_url: Optional[str] = None,
    generated_audio_url: Optional[str] = None,
    negative_prompt: Optional[str] = None,
    first_frame_url: Optional[str] = None,
    reference_video_url: Optional[str] = None,
    duration: int = 15,
    model: str = "bytedance/seedance-2-5",
    callback_url: Optional[str] = None,
    resolution: str = "720p",
    aspect_ratio: str = "9:16",
    return_last_frame: bool = False,
    generate_audio: bool = True,
    web_search: bool = False,
) -> str:

    input_payload: Dict[str, Any] = {
        "prompt": prompt,
        "return_last_frame": return_last_frame,
        "generate_audio": generate_audio,
        "resolution": resolution,
        "aspect_ratio": aspect_ratio,
        "duration": duration,
        "web_search": web_search,
    }

    # Kie's Multimodal Reference-to-Video (reference_image_urls) and
    # Image-to-Video (first_frame_url) modes are mutually exclusive.
    if first_frame_url:
        input_payload["first_frame_url"] = to_public_uri(first_frame_url)
    else:
        reference_image_urls = [
            to_public_uri(url) for url in as_url_list(avatar_image_url) + as_url_list(product_image_url)
        ]

        if len(reference_image_urls) > 9:
            logger.warning(
                "[STEP 3] Kie accepts at most 9 reference images, got %s — truncating to the first 9",
                len(reference_image_urls),
            )
            reference_image_urls = reference_image_urls[:9]

        input_payload["reference_image_urls"] = reference_image_urls

        if generated_audio_url:
            input_payload["reference_audio_urls"] = [
                to_public_uri(generated_audio_url)
            ]


        reference_video_urls = [to_public_uri(url) for url in as_url_list(reference_video_url)][:10]
        if reference_video_urls:
            input_payload["reference_video_urls"] = reference_video_urls

    if negative_prompt:
        input_payload["negative_prompt"] = negative_prompt

    payload: Dict[str, Any] = {
        "model": model,
        "input": input_payload,
    }

    if callback_url:
        payload["callBackUrl"] = callback_url

    logger.info("[STEP 3] Kie Seedance createTask payload: %s", payload)

    response = kie_api_request(
        "POST",
        KIE_CREATE_TASK_PATH,
        payload=payload
    )

    task_id = response.get("data", {}).get("taskId")

    if not task_id:
        raise RuntimeError(
            f"Kie createTask response did not include taskId: {response}"
        )

    return task_id


def get_kie_task_info(task_id: str) -> Dict[str, Any]:

    return kie_api_request(
        "GET",
        KIE_RECORD_INFO_PATH,
        query={"taskId": task_id},
    )


def extract_kie_result_urls(task_info: Dict[str, Any]) -> List[str]:


    data = task_info.get("data", {}) if isinstance(task_info, dict) else {}
    result_json = data.get("resultJson")

    if not result_json:
        return []

    if isinstance(result_json, str):
        try:
            parsed = json.loads(result_json)
        except json.JSONDecodeError:
            return []
    elif isinstance(result_json, dict):
        parsed = result_json
    else:
        return []

    urls: List[str] = []

    for key in ["resultUrls", "videoUrls", "urls", "outputUrls"]:
        value = parsed.get(key)

        if isinstance(value, list):
            urls.extend(
                [item for item in value if isinstance(item, str)]
            )

    for key in ["url", "videoUrl", "video_url", "output"]:
        value = parsed.get(key)

        if isinstance(value, str):
            urls.append(value)

    return urls


LAST_FRAME_KEYS = [
    "lastFrameUrl", "last_frame_url", "lastFrame",
    "endFrameUrl", "finalFrameUrl",
]


def extract_kie_last_frame_url(task_info: Dict[str, Any]) -> Optional[str]:
    data = task_info.get("data", {}) if isinstance(task_info, dict) else {}

    result_json = data.get("resultJson")
    parsed = {}

    if isinstance(result_json, str):
        try:
            parsed = json.loads(result_json)
        except json.JSONDecodeError:
            parsed = {}
    elif isinstance(result_json, dict):
        parsed = result_json

    for source in (parsed, data):
        for key in LAST_FRAME_KEYS:
            value = source.get(key)
            if isinstance(value, str) and value:
                return value

    return None


def wait_for_kie_task(
    task_id: str,
    timeout_seconds: int = 900,
    initial_interval_seconds: float = 3.0,
    max_interval_seconds: float = 15.0,
) -> Dict[str, Any]:


    deadline = time.time() + timeout_seconds
    interval = initial_interval_seconds

    while time.time() < deadline:
        task_info = get_kie_task_info(task_id)
        data = task_info.get("data", {})

        state = data.get("state") or data.get("status")
        progress = data.get("progress")

        logger.info(
            "[STEP 3] Kie Seedance task %s state=%s progress=%s",
            task_id,
            state,
            progress,
        )

        if state == "success":
            result_urls = extract_kie_result_urls(task_info)

            return {
                "task_id": task_id,
                "state": state,
                "video_url": result_urls[0] if result_urls else None,
                "result_urls": result_urls,
                "raw": task_info,
            }

        if state in {"fail", "failed", "error"}:
            fail_code = data.get("failCode", "")
            fail_msg = data.get("failMsg", "")

            raise RuntimeError(
                "Kie Seedance task failed. "
                f"task_id={task_id}, "
                f"failCode={fail_code}, "
                f"failMsg={fail_msg}"
            )

        time.sleep(interval)
        interval = min(max_interval_seconds, interval * 1.5)

    raise TimeoutError(
        f"Kie Seedance task timed out after {timeout_seconds}s: {task_id}"
    )


SEEDANCE_MAX_SEGMENT_SECONDS = 15


def _get_media_duration_seconds(media_bytes: bytes, suffix: str = ".mp3") -> float:
    with tempfile.TemporaryDirectory() as tmpdir:
        media_path = os.path.join(tmpdir, f"media{suffix}")
        with open(media_path, "wb") as f:
            f.write(media_bytes)

        result = subprocess.run(
            [
                "ffprobe", "-v", "error", "-show_entries", "format=duration",
                "-of", "default=noprint_wrappers=1:nokey=1", media_path,
            ],
            check=True, capture_output=True, text=True,
        )

        return float(result.stdout.strip())


def _get_audio_duration_seconds(audio_bytes: bytes) -> float:
    return _get_media_duration_seconds(audio_bytes, suffix=".mp3")


def _trim_audio(audio_bytes: bytes, end_seconds: float) -> bytes:
    with tempfile.TemporaryDirectory() as tmpdir:
        in_path = os.path.join(tmpdir, "in.mp3")
        out_path = os.path.join(tmpdir, "out.mp3")

        with open(in_path, "wb") as f:
            f.write(audio_bytes)

        subprocess.run(
            ["ffmpeg", "-y", "-i", in_path, "-t", str(end_seconds), "-c", "copy", out_path],
            check=True, capture_output=True,
        )

        with open(out_path, "rb") as f:
            return f.read()


def _split_into_segment_lengths(duration: int, max_segment_seconds: int = SEEDANCE_MAX_SEGMENT_SECONDS) -> List[int]:
    # Even distribution (not greedy max-chunking) so no segment ends up
    # below Kie's 4s minimum, e.g. duration=46 -> [12,12,11,11], not [15,15,15,1].
    num_segments = -(-duration // max_segment_seconds)  # ceil division
    base, remainder = divmod(duration, num_segments)

    return [base + (1 if i < remainder else 0) for i in range(num_segments)]


def _concat_and_mux_segments(video_byte_chunks: List[bytes], audio_bytes: Optional[bytes]) -> bytes:
    with tempfile.TemporaryDirectory() as tmpdir:
        segment_paths = []

        for i, chunk in enumerate(video_byte_chunks):
            path = os.path.join(tmpdir, f"segment_{i}.mp4")
            with open(path, "wb") as f:
                f.write(chunk)
            segment_paths.append(path)

        input_args = []
        for path in segment_paths:
            input_args += ["-i", path]

        filter_inputs = "".join(f"[{i}:v]" for i in range(len(segment_paths)))
        filter_complex = f"{filter_inputs}concat=n={len(segment_paths)}:v=1:a=0[outv]"

        concatenated_path = os.path.join(tmpdir, "concatenated.mp4")
        subprocess.run(
            [
                "ffmpeg", "-y", *input_args,
                "-filter_complex", filter_complex,
                "-map", "[outv]",
                "-c:v", "libx264", "-pix_fmt", "yuv420p",
                concatenated_path,
            ],
            check=True, capture_output=True,
        )

        if not audio_bytes:
            with open(concatenated_path, "rb") as f:
                return f.read()

        audio_path = os.path.join(tmpdir, "audio.mp3")
        with open(audio_path, "wb") as f:
            f.write(audio_bytes)

        # The output must always match the VIDEO's own length exactly —
        # neither "-shortest" (cuts the video short if audio is shorter,
        # which we don't want) nor no limit at all (leaves a black/frozen
        # tail with only audio playing if audio is longer, which is exactly
        # what happened here). Measuring the video and capping with -t
        # handles both directions: audio trimmed if it overshoots, left to
        # end naturally/silently if it undershoots.
        with open(concatenated_path, "rb") as f:
            video_duration = _get_media_duration_seconds(f.read(), suffix=".mp4")

        final_path = os.path.join(tmpdir, "final.mp4")
        subprocess.run(
            [
                "ffmpeg", "-y", "-i", concatenated_path, "-i", audio_path,
                "-c:v", "copy", "-c:a", "aac", "-map", "0:v:0", "-map", "1:a:0",
                "-t", str(video_duration),
                final_path,
            ],
            check=True, capture_output=True,
        )

        with open(final_path, "rb") as f:
            return f.read()


def apply_logo_watermark(video_url: str, logo_url: str = AMALFA_LOGO_URL) -> str:
    video_bytes = download_image_bytes(video_url)
    logo_bytes = download_image_bytes(logo_url)

    with tempfile.TemporaryDirectory() as tmpdir:
        video_path = os.path.join(tmpdir, "in.mp4")
        logo_path = os.path.join(tmpdir, "logo.jpg")
        out_path = os.path.join(tmpdir, "out.mp4")

        with open(video_path, "wb") as f:
            f.write(video_bytes)
        with open(logo_path, "wb") as f:
            f.write(logo_bytes)

        subprocess.run(
            [
                "ffmpeg", "-y", "-i", video_path, "-i", logo_path,
                "-filter_complex",
                # Fixed 90px-tall logo, bottom-right with a 24px margin —
                # independent of the video's own resolution/aspect ratio.
                "[1:v]scale=-1:90[logo];[0:v][logo]overlay=W-w-24:H-h-24[outv]",
                "-map", "[outv]", "-map", "0:a?",
                "-c:v", "libx264", "-pix_fmt", "yuv420p", "-c:a", "copy",
                out_path,
            ],
            check=True, capture_output=True,
        )

        with open(out_path, "rb") as f:
            watermarked_bytes = f.read()

    return upload_file_to_kie(watermarked_bytes, "final_video_watermarked.mp4", "video/mp4")


def _parse_time_range_seconds(time_range: str) -> Optional[tuple]:
    numbers = re.findall(r"\d+", time_range or "")

    if len(numbers) < 2:
        return None

    return int(numbers[0]), int(numbers[1])


def _shots_in_window(shots: List[Dict[str, Any]], window_start: int, window_end: int) -> List[Dict[str, Any]]:
    matching = []

    for shot in shots:
        parsed = _parse_time_range_seconds(shot.get("time_range", ""))
        if not parsed:
            continue

        shot_start, shot_end = parsed
        if shot_start < window_end and shot_end > window_start:
            matching.append(shot)

    return matching


def _shots_prompt_block(shots: List[Dict[str, Any]]) -> str:
    return "\n".join(
        f'- {shot["title"]}: {shot["description"]} (Camera: {shot["camera"]})'
        for shot in shots
    )


def generate_seedance_video(
    generated_script: str,
    avatar_image_url: str,
    product_image_url: str,
    generated_audio_url: str,
    negative_prompt: Optional[str] = None,
    reference_video_url: Optional[str] = None,
    duration: int = 15,
    model: str = "bytedance/seedance-2-5",
    callback_url: Optional[str] = None,
    poll: bool = True,
    timeout_seconds: int = 1800,
    resolution: str = "720p",
    aspect_ratio: str = "9:16",
    return_last_frame: bool = False,
    storyboard: Optional[Dict[str, Any]] = None,
    creative: Optional[Dict[str, Any]] = None,
    script_model: str = "gpt-5.5",
    on_storyboard_ready: Optional[Callable[[Dict[str, Any]], None]] = None,
) -> Dict[str, Any]:

    logger.info("[STEP 3] Kie Seedance video generation started")
    logger.info("Avatar image URL: %s", avatar_image_url)
    logger.info("Product image URL: %s", product_image_url)
    logger.info("Generated audio URL: %s", generated_audio_url)
    logger.info("Duration: %s", duration)
    logger.info("Model: %s", model)

    avatar_urls = as_url_list(avatar_image_url)
    product_urls = as_url_list(product_image_url)

    avatar_image_refs = ", ".join(f"[Image{i + 1}]" for i in range(len(avatar_urls)))
    product_image_refs = ", ".join(
        f"[Image{i + 1}]" for i in range(len(avatar_urls), len(avatar_urls) + len(product_urls))
    )

    # The reference video is a CONCEPT reference only. Without this fence the
    # model borrows the people and products in it too, which is how a
    # different-looking avatar ends up in the video.
    reference_video_note = (
        "\nUse [Video1] ONLY as a creative concept reference: copy its idea, structure, "
        "shot progression, pacing, camera movement, framing, transitions, and editing rhythm.\n"
        "Do NOT take anything else from it. Ignore its people, faces, bodies, clothing, hair, "
        "voices, products, brands, logos, on-screen text, locations, and background objects. "
        "None of them may appear in the output.\n"
        f"The person is always {avatar_image_refs} and the product is always {product_image_refs} — "
        "never the ones from [Video1].\n"
        if reference_video_url else ""
    )

    safe_video_prompt = f"""
{generated_script}

Use {avatar_image_refs} as the avatar/person reference{'s' if len(avatar_urls) > 1 else ''} (same person, different angles if more than one).
Use {product_image_refs} as the product reference{'s' if len(product_urls) > 1 else ''} (same product, different angles if more than one).
Use [Audio1] as the generated voiceover timing and pacing reference.
{reference_video_note}Keep the avatar face, outfit, and body proportions consistent.
Keep the product shape, packaging, color, and branding consistent.
Show the product clearly in a realistic, social-media-ready ad style.
Do not change the product design, avatar face, or avatar outfit/clothing.
""".strip()

    segment_lengths = _split_into_segment_lengths(duration)
    shots = storyboard.get("shots", []) if storyboard else []

    try:
        if len(segment_lengths) == 1:
            task_id = create_kie_seedance_task(
                prompt=safe_video_prompt,
                avatar_image_url=avatar_image_url,
                product_image_url=product_image_url,
                generated_audio_url=generated_audio_url,
                negative_prompt=negative_prompt,
                reference_video_url=reference_video_url,
                duration=duration,
                model=model,
                callback_url=callback_url,
                resolution=resolution,
                aspect_ratio=aspect_ratio,
                generate_audio=True,
                return_last_frame=return_last_frame,
            )

            logger.info("[STEP 3] Kie Seedance task created: %s", task_id)

            if not poll:
                return {
                    "task_id": task_id,
                    "state": "submitted",
                    "video_url": None,
                    "result_urls": [],
                }

            result = wait_for_kie_task(
                task_id=task_id,
                timeout_seconds=timeout_seconds,
            )

            logger.info("[STEP 3] Kie Seedance video generation completed")
            logger.info("[STEP 3] Video URL: %s", result.get("video_url"))

            return result

        logger.info(
            "[STEP 3] Duration %ss exceeds the %ss single-call max — chaining %s segments",
            duration, SEEDANCE_MAX_SEGMENT_SECONDS, len(segment_lengths),
        )


        shots_per_segment = None
        combined_storyboard = storyboard

        if creative is not None:
            logger.info("[STEP 3] Generating a storyboard for each of the %s segments", len(segment_lengths))
            shots_per_segment = []
            all_shots = []
            first_segment_storyboard = None
            seg_start = 0

            for seg_len in segment_lengths:
                seg_end = seg_start + seg_len
                seg_storyboard = generate_storyboard(
                    creative=creative,
                    duration_seconds=duration,
                    model=script_model,
                    window_start=seg_start,
                    window_end=seg_end,
                    prior_shots=all_shots,
                    aspect_ratio=aspect_ratio,
                )

                if first_segment_storyboard is None:
                    first_segment_storyboard = seg_storyboard

                for shot in seg_storyboard["shots"]:
                    shot["shot_number"] = len(all_shots) + 1
                    all_shots.append(shot)

                shots_per_segment.append(seg_storyboard["shots"])
                seg_start = seg_end

            combined_storyboard = {
                **first_segment_storyboard,
                "duration_seconds": duration,
                "shots": all_shots,
            }

            if on_storyboard_ready:
                on_storyboard_ready(combined_storyboard)

        shots = combined_storyboard.get("shots", []) if combined_storyboard else []

        # Kie rejects reference_audio_urls longer than ~15.2s regardless of
        # the video's own requested duration ("audio duration must be less
        # than or equal to 15.2") — but generated_audio_url is now the FULL
        # voiceover (sized to fill the whole ad by ensure_voiceover_fills_duration).
        # Only segment 1 gets an audio reference, so trim just its own
        # window out of the full track before passing it to Kie.
        segment1_audio_url = generated_audio_url
        if generated_audio_url and segment_lengths[0] < duration:
            trimmed_bytes = _trim_audio(download_image_bytes(generated_audio_url), segment_lengths[0])
            segment1_audio_url = upload_file_to_kie(trimmed_bytes, "segment1_audio.mp3", "audio/mpeg")

        segment_video_urls = []
        video_byte_chunks = []
        segment_start = 0

        for i, segment_len in enumerate(segment_lengths):
            is_last_segment = i == len(segment_lengths) - 1
            segment_end = segment_start + segment_len
            window_shots = (
                shots_per_segment[i] if shots_per_segment is not None
                else _shots_in_window(shots, segment_start, segment_end)
            )
            shots_block = _shots_prompt_block(window_shots)

            cta_instruction = (
                "\n\nThis is the final segment of the ad — the CTA must appear as on-screen "
                "text only in the last shot above, ending at the very end of this segment, "
                "never earlier."
                if is_last_segment else ""
            )

            # No segment gets the whole-ad seedance_prompt (which narrates
            # all the way through the CTA) — given a prompt describing the
            # full ad's arc, Kie would compress hook->CTA into just this
            # segment's short duration, which is why the CTA was showing up
            # mid-video. Use only the shots scoped to this exact segment's
            # window; fall back to the whole-ad text only when no
            # per-segment storyboard exists at all (single-segment videos).
            segment_content = (
                f"Follow this shot list for this part of the video:\n{shots_block}"
                if shots_block else generated_script
            )

            # Every segment re-anchors to the ORIGINAL avatar/product
            # reference images (not the previous segment's last frame) so
            # identity can't drift across a chain of hand-offs — continuity
            # of the story comes from each segment's storyboard already
            # being written aware of what prior segments covered.
            continuity_note = (
                "" if i == 0 else
                "This continues directly from the previous segment — keep the exact same "
                "avatar, outfit, product, and setting, and continue the ad's action naturally "
                "without restarting or repeating earlier beats.\n\n"
            )

            segment_prompt = f"""
{continuity_note}{segment_content}

Use {avatar_image_refs} as the avatar/person reference{'s' if len(avatar_urls) > 1 else ''} (same person, different angles if more than one).
Use {product_image_refs} as the product reference{'s' if len(product_urls) > 1 else ''} (same product, different angles if more than one).
Use [Audio1] as the generated voiceover timing and pacing reference.
{reference_video_note}Keep the avatar face, outfit, and body proportions consistent.
Keep the product shape, packaging, color, and branding consistent.
Show the product clearly in a realistic, social-media-ready ad style.
Do not change the product design, avatar face, or avatar outfit/clothing.
""".strip() + cta_instruction

            task_id = create_kie_seedance_task(
                prompt=segment_prompt,
                avatar_image_url=avatar_image_url,
                product_image_url=product_image_url,
                generated_audio_url=segment1_audio_url if i == 0 else None,
                negative_prompt=negative_prompt,
                reference_video_url=reference_video_url,
                duration=segment_len,
                model=model,
                resolution=resolution,
                aspect_ratio=aspect_ratio,
                generate_audio=(i == 0),
            )

            segment_start = segment_end

            logger.info("[STEP 3] Segment %s/%s task created: %s", i + 1, len(segment_lengths), task_id)

            segment_result = wait_for_kie_task(task_id=task_id, timeout_seconds=timeout_seconds)

            if not segment_result.get("video_url"):
                raise RuntimeError(f"Segment {i + 1} did not return a video_url: {segment_result}")

            segment_video_urls.append(segment_result["video_url"])
            video_byte_chunks.append(download_image_bytes(segment_result["video_url"]))

        logger.info("[STEP 3] All %s segments generated, stitching with ffmpeg", len(segment_lengths))

        audio_bytes = download_image_bytes(generated_audio_url) if generated_audio_url else None
        final_video_bytes = _concat_and_mux_segments(video_byte_chunks, audio_bytes)

        final_video_url = upload_file_to_kie(final_video_bytes, "final_video.mp4", "video/mp4")

        logger.info("[STEP 3] Kie Seedance video generation completed (chained)")
        logger.info("[STEP 3] Video URL: %s", final_video_url)

        return {
            "task_id": None,
            "state": "success",
            "video_url": final_video_url,
            "result_urls": [final_video_url],
            "segment_video_urls": segment_video_urls,
            "storyboard": combined_storyboard,
            "raw": None,
        }

    except Exception:
        logger.exception("[STEP 3] Kie Seedance video generation failed")
        raise


def generate_product_ad_video(
    product_image_url: str,
    avatar_image_url: str,
    voice_public_url: str,
    angle: str,
    concept: str,
    avatar_description: Optional[str] = None,
    voice_description: Optional[str] = None,
    duration_seconds: int = 15,
    script_model: str = "gpt-5.5",
    aspect_ratio: str = "9:16",
    manual_script: Optional[str] = None,
    reference_video_url: Optional[str] = None,
    on_progress: Optional[Callable[[str], None]] = None,
) -> Dict[str, Any]:

    if on_progress is None:
        on_progress = lambda _message: None

    logger.info("Product ad video pipeline started")
    logger.info("Product image URL: %s", product_image_url)
    logger.info("Avatar image URL: %s", avatar_image_url)
    logger.info("Voice public URL: %s", voice_public_url)

    try:
        on_progress("Step 1/5: Writing script and visual direction...")

        creative = generate_script_and_visual_direction(
            product_image_url=product_image_url,
            angle=angle,
            concept=concept,
            avatar_description=avatar_description,
            voice_description=voice_description or (
                "Clone the reference voice and generate a natural performance ad voiceover."
            ),
            duration_seconds=duration_seconds,
            model=script_model,
            manual_script=manual_script,
        )

        logger.info("Step 1 completed")

        is_chained = duration_seconds > SEEDANCE_MAX_SEGMENT_SECONDS

        # For a single-call (<=15s) video, the storyboard covers the whole
        # thing and can be built right away. For a chained video, a separate
        # storyboard is generated per segment *inside* generate_seedance_video
        # (so each segment's continuation prompt gets shots actually written
        # for that window) — it isn't known until on_storyboard_ready fires.
        storyboard = None
        if not is_chained:
            storyboard = generate_storyboard(
                creative=creative,
                duration_seconds=duration_seconds,
                model=script_model,
                aspect_ratio=aspect_ratio,
            )
            logger.info("Step 1B (storyboard) completed")

        def _build_poster(sb: Dict[str, Any]) -> Optional[str]:
            try:
                return generate_storyboard_poster(
                    storyboard=sb,
                    avatar_image_url=avatar_image_url,
                    product_image_url=product_image_url
                )
            except Exception:
                logger.exception("Storyboard poster generation failed, continuing without it")
                return None

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as poster_executor:
            poster_future_holder: Dict[str, Any] = {}

            def _on_storyboard_ready(sb: Dict[str, Any]) -> None:
                poster_future_holder["future"] = poster_executor.submit(_build_poster, sb)

            if storyboard is not None:
                _on_storyboard_ready(storyboard)

            on_progress("Step 2/5: Cloning voice and generating speech...")

            audio_result = generate_audio_from_voice_url(
                voiceover_script=creative["voiceover_script"],
                voice_public_url=voice_public_url
            )

            generated_audio_url = audio_result["generated_audio_url"]
            voice_id = audio_result["voice_id"]

            logger.info("Step 2 completed")
            logger.info("Voice ID: %s", voice_id)
            logger.info("Generated audio URL: %s", generated_audio_url)

            generated_audio_url = ensure_voiceover_fills_duration(
                creative=creative,
                generated_audio_url=generated_audio_url,
                voice_id=voice_id,
                duration_seconds=duration_seconds,
                script_model=script_model,
                allow_rewrite=not manual_script,
            )

            on_progress("Step 3/5: Generating video (this can take several minutes)...")

            video_result = generate_seedance_video(
                generated_script=creative["seedance_prompt"],
                avatar_image_url=avatar_image_url,
                product_image_url=product_image_url,
                generated_audio_url=generated_audio_url,
                negative_prompt=creative.get("negative_prompt"),
                reference_video_url=reference_video_url,
                duration=duration_seconds,
                model="bytedance/seedance-2-5",
                aspect_ratio=aspect_ratio,
                storyboard=storyboard,
                creative=creative if is_chained else None,
                script_model=script_model,
                on_storyboard_ready=_on_storyboard_ready if is_chained else None,
            )

            if is_chained:
                storyboard = video_result.get("storyboard") or {}

            logger.info("Step 3 completed")

            if video_result.get("video_url"):
                on_progress("Step 4/5: Adding logo watermark...")
                video_result["video_url"] = apply_logo_watermark(video_result["video_url"])
                logger.info("Logo watermark applied: %s", video_result["video_url"])

            on_progress("Step 5/5: Finalizing storyboard poster...")

            poster_future = poster_future_holder.get("future")
            poster_image_url = poster_future.result() if poster_future else None
            if poster_image_url:
                storyboard["poster_image_url"] = poster_image_url

        on_progress("Done")
        logger.info("Product ad video pipeline completed")

        return {
            "creative": creative,
            "storyboard": storyboard,
            "voice_id": voice_id,
            "generated_audio_url": generated_audio_url,
            "video": video_result
        }

    except Exception:
        logger.exception("Product ad video pipeline failed")
        raise