engines/image_generation_engine.py

import logging
import os
import time
from typing import Any, Dict, List, Optional

import replicate

from engines.visual_prompt_engine import VisualPromptEngine

logger = logging.getLogger("ONYX")


class ImageGenerationEngine:
    def __init__(self):
        print("🚀 IMAGE ENGINE INIT START")

        self.api_token = os.getenv("REPLICATE_API_TOKEN")
        if not self.api_token:
            print("❌ TOKEN MANQUANT")
            raise Exception("REPLICATE_API_TOKEN manquant")

        self.client = replicate.Client(api_token=self.api_token)

        self.text_model = os.getenv(
            "ONYX_TEXT_IMAGE_MODEL",
            "black-forest-labs/flux-schnell"
        )

        self.master_model = os.getenv(
            "ONYX_MASTER_IMAGE_MODEL",
            "black-forest-labs/flux-schnell"
        )

        self.reference_model = os.getenv(
            "ONYX_REFERENCE_IMAGE_MODEL",
            "ideogram-ai/ideogram-character"
        )

        self.max_images = int(os.getenv("ONYX_MAX_IMAGES", "2"))
        self.delay_between_images = int(os.getenv("ONYX_DELAY_BETWEEN_IMAGES", "12"))
        self.max_attempts_per_image = int(os.getenv("ONYX_MAX_ATTEMPTS_PER_IMAGE", "2"))
        self.aspect_ratio = os.getenv("ONYX_IMAGE_ASPECT_RATIO", "3:4")

        self.visual_prompt_engine = VisualPromptEngine()

        print("✅ IMAGE ENGINE READY")
        print(f"📝 text_model = {self.text_model}")
        print(f"🧬 master_model = {self.master_model}")
        print(f"🎭 reference_model = {self.reference_model}")
        print(f"🪙 max_images = {self.max_images}")
        print(f"🪙 max_attempts_per_image = {self.max_attempts_per_image}")

    def _clean(self, value: Any) -> str:
        return " ".join(str(value or "").strip().split())

    def _is_rate_limit_error(self, error: Exception) -> bool:
        msg = str(error).lower()
        return (
            "429" in msg
            or "rate limit" in msg
            or "too many requests" in msg
            or "throttled" in msg
            or "throttle" in msg
        )

    def _is_credit_or_payment_error(self, error: Exception) -> bool:
        msg = str(error).lower()
        return (
            "402" in msg
            or "payment required" in msg
            or "insufficient balance" in msg
            or "insufficient credit" in msg
            or "billing" in msg
            or "quota exceeded" in msg
        )

    def _wait_before_retry(self, attempt: int):
        schedule = [10, 20]
        seconds = schedule[min(attempt, len(schedule) - 1)]
        logger.warning(f"⏳ Retry in {seconds}s...")
        time.sleep(seconds)

    def _extract_output_url(self, output) -> Optional[str]:
        if isinstance(output, list) and output:
            output = output[0]

        if hasattr(output, "url"):
            return output.url

        if isinstance(output, str):
            return output

        return None

    def _run_model_once(self, model: str, model_input: Dict[str, Any]) -> Optional[str]:
        logger.info(f"🎨 Calling Replicate model: {model}")
        logger.info(f"🧪 Input keys: {list(model_input.keys())}")

        output = self.client.run(model, input=model_input)
        return self._extract_output_url(output)

    def _run_model_safe(self, model: str, model_input: Dict[str, Any]) -> Optional[str]:
        for attempt in range(self.max_attempts_per_image):
            try:
                logger.info(f"🎨 Try {attempt + 1}/{self.max_attempts_per_image} on {model}")
                image_url = self._run_model_once(model, model_input)

                if image_url and isinstance(image_url, str) and image_url.strip():
                    logger.info("✅ Image generated")
                    return image_url

                logger.warning("⚠️ Empty image response")
                if attempt < self.max_attempts_per_image - 1:
                    self._wait_before_retry(attempt)

            except Exception as e:
                logger.error(f"❌ Image error on {model}: {e}")

                if self._is_rate_limit_error(e):
                    logger.warning("🚦 Rate limit detected")
                    if attempt < self.max_attempts_per_image - 1:
                        self._wait_before_retry(attempt)
                        continue
                    return None

                if self._is_credit_or_payment_error(e):
                    logger.warning("💳 Credit/payment issue detected -> stop immediate")
                    return None

                if attempt < self.max_attempts_per_image - 1:
                    self._wait_before_retry(attempt)

        logger.warning(f"❌ Final image failure on {model}")
        return None

    def generate_image_once(self, prompt: str) -> Optional[str]:
        return self._run_model_once(
            self.text_model,
            {"prompt": prompt}
        )

    def generate_safe(self, prompt: str) -> Optional[str]:
        return self._run_model_safe(
            self.text_model,
            {"prompt": prompt}
        )

    def generate_batch(self, prompts: List[str]) -> List[str]:
        prompts = (prompts or [])[:self.max_images]

        if not prompts:
            return []

        images: List[str] = []

        for i, prompt in enumerate(prompts):
            logger.info(f"🎬 Text-only image {i + 1}/{len(prompts)}")

            img = self.generate_safe(prompt)
            if not img:
                logger.warning("❌ One text-only image failed -> returning partial batch")
                break

            images.append(img)

            if i < len(prompts) - 1:
                logger.info(f"🕒 Anti-throttle pause: {self.delay_between_images}s")
                time.sleep(self.delay_between_images)

        return images

    def build_prompts_from_pipeline(
        self,
        scenes: List[str],
        visual_bible: Dict[str, Any],
        shot_plan: List[Dict[str, Any]],
        creative_mode: str
    ) -> List[str]:
        prompts: List[str] = []
        limited_scenes = (scenes or [])[:self.max_images]

        for i, scene_text in enumerate(limited_scenes):
            shot = shot_plan[i] if i < len(shot_plan) else {}

            prompt = self.visual_prompt_engine.build_prompt(
                scene_text=self._clean(scene_text),
                visual_bible=visual_bible or {},
                shot_plan=shot or {},
                creative_mode=creative_mode
            )
            prompts.append(prompt)

        return prompts

    def _build_master_reference_prompt(
        self,
        visual_bible: Dict[str, Any],
        creative_mode: str
    ) -> str:
        main_subject = self._clean(visual_bible.get("main_subject"))
        identity = self._clean(visual_bible.get("identity"))
        face = self._clean(visual_bible.get("face"))
        hair = self._clean(visual_bible.get("hair"))
        body = self._clean(visual_bible.get("body"))
        outfit = self._clean(visual_bible.get("outfit"))
        lighting = self._clean(visual_bible.get("lighting"))
        color_palette = self._clean(visual_bible.get("color_palette"))
        environment = self._clean(visual_bible.get("environment"))
        style_keywords = self._clean(visual_bible.get("style_keywords"))
        camera_language = self._clean(visual_bible.get("camera_language"))
        reference_master_prompt = self._clean(visual_bible.get("reference_master_prompt"))

        return f"""
character master reference image, single recurring protagonist, identity lock reference

MAIN SUBJECT:
{main_subject}

IDENTITY:
{identity}
{face}
{hair}
{body}
{outfit}

STYLE:
{style_keywords}

LIGHTING:
{lighting}

COLOR:
{color_palette}

ENVIRONMENT:
{environment}

CAMERA LANGUAGE:
{camera_language}

MASTER REFERENCE INSTRUCTION:
{reference_master_prompt}

IMPORTANT:
- one single character only
- neutral pose
- front-facing or slight three-quarter view
- face fully readable
- same outfit that will be reused later
- no action pose
- no running
- no fight move
- no collage
- no text
- no logo
- no watermark
- clean composition
- identity-first image
""".strip()

    def _build_reference_scene_input(
        self,
        prompt: str,
        reference_image_url: str,
        creative_mode: str
    ) -> Dict[str, Any]:
        style_type = "Auto"
        if creative_mode == "anime":
            style_type = "Fiction"
        elif creative_mode in ["story", "emotional", "tiktok"]:
            style_type = "Realistic"

        return {
            "prompt": prompt,
            "character_reference_image": reference_image_url,
            "style_type": style_type,
            "aspect_ratio": self.aspect_ratio
        }

    def generate_master_reference(
        self,
        visual_bible: Dict[str, Any],
        creative_mode: str
    ) -> Optional[str]:
        master_prompt = self._build_master_reference_prompt(
            visual_bible=visual_bible,
            creative_mode=creative_mode
        )

        logger.info("🧬 Generating master reference image...")
        return self._run_model_safe(
            self.master_model,
            {"prompt": master_prompt}
        )

    def generate_batch_v3(
        self,
        scenes: List[str],
        visual_bible: Dict[str, Any],
        shot_plan: List[Dict[str, Any]],
        creative_mode: str
    ) -> List[str]:
        limited_scenes = (scenes or [])[:self.max_images]
        if not limited_scenes:
            return []

        image_prompts = self.build_prompts_from_pipeline(
            scenes=limited_scenes,
            visual_bible=visual_bible,
            shot_plan=shot_plan,
            creative_mode=creative_mode
        )

        master_reference_url = self.generate_master_reference(
            visual_bible=visual_bible,
            creative_mode=creative_mode
        )

        if not master_reference_url:
            logger.warning("⚠️ Master reference failed -> stop image pipeline to save credits")
            return []

        logger.info(f"🧬 Master reference ready: {master_reference_url}")

        images: List[str] = []

        for i, prompt in enumerate(image_prompts):
            logger.info(f"🎬 Reference-based scene {i + 1}/{len(image_prompts)}")

            model_input = self._build_reference_scene_input(
                prompt=prompt,
                reference_image_url=master_reference_url,
                creative_mode=creative_mode
            )

            img = self._run_model_safe(self.reference_model, model_input)

            if not img:
                logger.warning("❌ One reference-based image failed -> stop and return partial batch")
                break

            images.append(img)

            if i < len(image_prompts) - 1:
                logger.info(f"🕒 Anti-throttle pause: {self.delay_between_images}s")
                time.sleep(self.delay_between_images)

        return images

    def generate_batch_v2(
        self,
        scenes: List[str],
        visual_bible: Dict[str, Any],
        shot_plan: List[Dict[str, Any]],
        creative_mode: str
    ) -> List[str]:
        return self.generate_batch_v3(
            scenes=scenes,
            visual_bible=visual_bible,
            shot_plan=shot_plan,
            creative_mode=creative_mode
        )