comic_generator.py

"""
Comic Generator with Character Consistency Integration
Safe wrapper that integrates character consistency system with automatic fallback to legacy system
"""

import asyncio
import os
from typing import Dict, List, Optional, Union
from PIL import Image
from config import Config

# Try to import character consistency system (optional dependency)
try:
    from character_consistency_manager import CharacterConsistencyManager
    CHARACTER_CONSISTENCY_AVAILABLE = True
    print("✓ Character Consistency system available")
except ImportError as e:
    CHARACTER_CONSISTENCY_AVAILABLE = False
    print(f"⚠️  Character Consistency not available - using legacy system")
    print(f"   Import error: {e}")


class ComicGeneratorWithConsistency:
    """
    Unified comic generator with character consistency support and legacy fallback

    This class provides a safe integration layer between the new character consistency
    system and the existing comic generation logic in app_comic.py

    Behavior:
    - If ENABLE_NANO_BANANA=true AND system available: Use character consistency
    - If character consistency fails OR disabled: Automatic fallback to legacy system
    - Legacy system always available as safety net
    """

    def __init__(self):
        """Initialize comic generator with optional character consistency"""
        self.consistency_enabled = (
            CHARACTER_CONSISTENCY_AVAILABLE and
            Config.ENABLE_NANO_BANANA
        )

        self.consistency_manager = None

        if self.consistency_enabled:
            try:
                self.consistency_manager = CharacterConsistencyManager()
                print("✓ Character Consistency Manager initialized")
            except Exception as e:
                print(f"⚠️  Failed to initialize Character Consistency: {e}")
                print("🔄 Will use legacy system")
                self.consistency_enabled = False

        # Stats tracking
        self.consistency_success_count = 0
        self.consistency_failure_count = 0
        self.fallback_count = 0

    async def generate_comic_from_prompt(
        self,
        story_prompt: str,
        num_panels: int = 3,
        art_style: str = "manga",
        include_surprise_character: bool = True,
        force_legacy: bool = False,
        output_dir: Optional[str] = None
    ) -> Dict:
        """
        Generate comic with automatic character consistency or legacy fallback

        Args:
            story_prompt: User's story prompt
            num_panels: Number of panels to generate
            art_style: Art style for the comic
            include_surprise_character: Add AI-generated surprise character
            force_legacy: Force use of legacy system (for testing)
            output_dir: Output directory for generated files

        Returns:
            Dictionary with generation results:
            {
                "success": bool,
                "method": str,  # "character_consistency" or "legacy"
                "panels": List[Image],
                "characters": Dict[str, str],  # Only if character_consistency used
                "output_dir": str,
                "panel_info": List[Dict],  # Panel metadata
                "error": Optional[str]
            }
        """

        # Determine which system to use
        use_consistency = (
            self.consistency_enabled and
            not force_legacy and
            self.consistency_manager is not None
        )

        if use_consistency:
            print(f"\n🎨 Generating comic with CHARACTER CONSISTENCY system...")
            print(f"   Story: {story_prompt[:100]}...")
            print(f"   Panels: {num_panels}, Style: {art_style}")

            try:
                result = await self._generate_with_consistency(
                    story_prompt=story_prompt,
                    num_panels=num_panels,
                    art_style=art_style,
                    include_surprise_character=include_surprise_character,
                    output_dir=output_dir
                )

                if result["success"]:
                    self.consistency_success_count += 1
                    print(f"✓ Character consistency generation successful!")
                    return result
                else:
                    print(f"⚠️  Character consistency failed: {result.get('error')}")
                    print(f"🔄 Falling back to legacy system...")
                    self.consistency_failure_count += 1

            except Exception as e:
                print(f"⚠️  Character consistency exception: {e}")
                print(f"🔄 Falling back to legacy system...")
                self.consistency_failure_count += 1

        # Fallback to legacy system
        self.fallback_count += 1
        print(f"\n🎨 Generating comic with LEGACY system...")

        return await self._generate_with_legacy(
            story_prompt=story_prompt,
            num_panels=num_panels,
            art_style=art_style,
            output_dir=output_dir
        )

    async def _generate_with_consistency(
        self,
        story_prompt: str,
        num_panels: int,
        art_style: str,
        include_surprise_character: bool,
        output_dir: Optional[str]
    ) -> Dict:
        """
        Generate comic using character consistency system

        This wraps the CharacterConsistencyManager.generate_comic_from_prompt method
        in a safe async context with proper error handling
        """

        try:
            # Call character consistency system (runs in executor for async compatibility)
            loop = asyncio.get_event_loop()
            result = await loop.run_in_executor(
                None,
                self.consistency_manager.generate_comic_from_prompt,
                story_prompt,
                num_panels,
                art_style,
                include_surprise_character,
                output_dir or "./auto_generated_comics"
            )

            # Add success flag and method
            result["success"] = True
            result["method"] = "character_consistency"

            return result

        except Exception as e:
            return {
                "success": False,
                "method": "character_consistency",
                "error": str(e),
                "panels": [],
                "characters": {},
                "panel_info": []
            }

    async def _generate_with_legacy(
        self,
        story_prompt: str,
        num_panels: int,
        art_style: str,
        output_dir: Optional[str]
    ) -> Dict:
        """
        Generate comic using legacy system (placeholder for app_comic.py integration)

        NOTE: This is a placeholder. In actual integration, this would call the
        existing comic generation logic from app_comic.py

        For now, returns a structure indicating legacy mode was used
        """

        return {
            "success": True,
            "method": "legacy",
            "panels": [],
            "characters": {},
            "panel_info": [
                {
                    "panel_number": i + 1,
                    "characters": [],
                    "filename": f"legacy_panel_{i+1}.png"
                }
                for i in range(num_panels)
            ],
            "output_dir": output_dir or "./generated_comics",
            "note": "Legacy system placeholder - integrate with app_comic.py logic"
        }

    def get_stats(self) -> Dict:
        """Get generation statistics"""
        total = self.consistency_success_count + self.consistency_failure_count + self.fallback_count

        return {
            "total_generations": total,
            "consistency_enabled": self.consistency_enabled,
            "consistency_available": CHARACTER_CONSISTENCY_AVAILABLE,
            "consistency_successes": self.consistency_success_count,
            "consistency_failures": self.consistency_failure_count,
            "legacy_fallbacks": self.fallback_count,
            "consistency_rate": (
                f"{(self.consistency_success_count / total * 100):.1f}%"
                if total > 0 else "N/A"
            )
        }

    def print_stats(self):
        """Print generation statistics"""
        stats = self.get_stats()

        print("\n" + "="*70)
        print("  COMIC GENERATOR STATISTICS")
        print("="*70)
        print(f"Character Consistency Available: {stats['consistency_available']}")
        print(f"Character Consistency Enabled:   {stats['consistency_enabled']}")
        print(f"\nTotal Generations:               {stats['total_generations']}")
        print(f"  ✓ Consistency Successes:       {stats['consistency_successes']}")
        print(f"  ⚠ Consistency Failures:        {stats['consistency_failures']}")
        print(f"  🔄 Legacy Fallbacks:            {stats['legacy_fallbacks']}")
        print(f"\nConsistency Success Rate:        {stats['consistency_rate']}")
        print("="*70)


# Example usage and testing
async def test_comic_generator():
    """Test the comic generator with both systems"""

    print("\n" + "="*70)
    print("  COMIC GENERATOR TEST")
    print("="*70)

    generator = ComicGeneratorWithConsistency()

    test_prompt = "A yogi teaching a young aspirant the rules of meditation and mindfulness leading to levitation by his own example in a beautiful vibrant colorful jungle"

    # Test 1: Character consistency (if available)
    print("\n--- Test 1: Character Consistency Mode ---")
    result1 = await generator.generate_comic_from_prompt(
        story_prompt=test_prompt,
        num_panels=3,
        art_style="manga",
        include_surprise_character=True
    )

    print(f"\nResult: {result1['method']}")
    print(f"Success: {result1['success']}")
    if result1.get('characters'):
        print(f"Characters: {list(result1['characters'].keys())}")

    # Test 2: Force legacy mode
    print("\n--- Test 2: Forced Legacy Mode ---")
    result2 = await generator.generate_comic_from_prompt(
        story_prompt=test_prompt,
        num_panels=3,
        art_style="manga",
        force_legacy=True
    )

    print(f"\nResult: {result2['method']}")
    print(f"Success: {result2['success']}")

    # Print stats
    generator.print_stats()


if __name__ == "__main__":
    # Run test
    asyncio.run(test_comic_generator())