character_consistency_manager.py

"""
Character Consistency Manager
Orchestrates Nano Banana + LlamaIndex hybrid workflow for character persistence
"""

import os
import json
from typing import Dict, List, Optional, Tuple, Union
from PIL import Image
from io import BytesIO
import base64

from nano_banana_client import NanoBananaClient
from character_extractor import CharacterExtractor
from story_memory import StoryMemory
from auto_character_intelligence import AutoCharacterIntelligence
from config import Config


class CharacterConsistencyManager:
    """Manages character consistency across comic panels using Nano Banana + LlamaIndex"""

    def __init__(
        self,
        character_library_dir: str = "./character_library",
        enable_nano_banana: bool = None,
        enable_character_library: bool = None
    ):
        """
        Initialize Character Consistency Manager

        Args:
            character_library_dir: Directory to store character library
            enable_nano_banana: Enable Nano Banana image generation (defaults to Config)
            enable_character_library: Enable character library persistence (defaults to Config)
        """
        self.character_library_dir = character_library_dir
        self.enable_nano_banana = enable_nano_banana if enable_nano_banana is not None else Config.ENABLE_NANO_BANANA
        self.enable_character_library = enable_character_library if enable_character_library is not None else Config.ENABLE_CHARACTER_LIBRARY

        # Create character library directory
        os.makedirs(character_library_dir, exist_ok=True)
        os.makedirs(os.path.join(character_library_dir, "images"), exist_ok=True)
        os.makedirs(os.path.join(character_library_dir, "profiles"), exist_ok=True)

        # Initialize components
        self.nano_banana = None
        self.character_extractor = None
        self.story_memory = None
        self.auto_intelligence = None

        if Config.GEMINI_API_KEY:
            try:
                if self.enable_nano_banana:
                    self.nano_banana = NanoBananaClient(model=Config.NANO_BANANA_MODEL)
                    print("✓ Nano Banana client initialized")
            except Exception as e:
                print(f"Warning: Could not initialize Nano Banana: {e}")

            try:
                self.character_extractor = CharacterExtractor()
                print("✓ Character Extractor initialized")
            except Exception as e:
                print(f"Warning: Could not initialize Character Extractor: {e}")

            try:
                self.auto_intelligence = AutoCharacterIntelligence()
                print("✓ Auto Character Intelligence initialized")
            except Exception as e:
                print(f"Warning: Could not initialize Auto Intelligence: {e}")

        if self.enable_character_library and Config.ENABLE_LLAMAINDEX:
            try:
                self.story_memory = StoryMemory()
                print("✓ Story Memory (LlamaIndex) initialized")
            except Exception as e:
                print(f"Warning: Could not initialize Story Memory: {e}")

        # Character cache (in-memory for current session)
        self.character_cache: Dict[str, Dict] = {}

        print(f"Character Consistency Manager initialized")
        print(f"  Nano Banana: {'✓ Enabled' if self.enable_nano_banana and self.nano_banana else '✗ Disabled'}")
        print(f"  Character Library: {'✓ Enabled' if self.enable_character_library and self.story_memory else '✗ Disabled'}")

    def create_character(
        self,
        character_description: str,
        character_name: str,
        style: str = "manga"
    ) -> Dict:
        """
        Create a new character with reference images

        Args:
            character_description: Detailed character description
            character_name: Name of the character
            style: Art style

        Returns:
            Character profile dictionary with reference images
        """
        if not self.nano_banana:
            raise ValueError("Nano Banana is not enabled. Set ENABLE_NANO_BANANA=true")

        print(f"\n=== Creating Character: {character_name} ===")

        # Generate character reference sheet
        print("Generating character reference images...")
        try:
            reference_images = self.nano_banana.generate_character_sheet(
                character_description=character_description,
                style=style
            )

            if not reference_images:
                raise ValueError("No reference images generated")

            print(f"Generated {len(reference_images)} reference images")

            # Extract character profile from first image
            print("Extracting character profile...")
            character_profile = self.character_extractor.extract_character_profile(
                image=reference_images[0],
                character_name=character_name
            )

            # Save reference images
            image_paths = []
            for i, img in enumerate(reference_images):
                img_filename = f"{character_name.lower().replace(' ', '_')}_ref_{i}.png"
                img_path = os.path.join(self.character_library_dir, "images", img_filename)
                img.save(img_path)
                image_paths.append(img_path)
                print(f"Saved reference image: {img_filename}")

            # Update profile with image paths
            character_profile["reference_images"] = image_paths
            character_profile["style"] = style
            character_profile["original_description"] = character_description

            # Save character profile
            profile_filename = f"{character_name.lower().replace(' ', '_')}_profile.json"
            profile_path = os.path.join(self.character_library_dir, "profiles", profile_filename)
            with open(profile_path, 'w') as f:
                json.dump(character_profile, f, indent=2)
            print(f"Saved character profile: {profile_filename}")

            # Cache character
            self.character_cache[character_name] = character_profile

            # Store in LlamaIndex if enabled
            if self.story_memory and self.enable_character_library:
                consistency_prompt = self.character_extractor.generate_consistency_prompt(character_profile)
                self.story_memory.store_story(
                    story_text=f"Character: {character_name}\n{consistency_prompt}",
                    metadata={
                        "type": "character_profile",
                        "character_name": character_name,
                        "style": style
                    }
                )
                print(f"Stored {character_name} in character library")

            print(f"✓ Character '{character_name}' created successfully")
            return character_profile

        except Exception as e:
            print(f"Error creating character: {e}")
            raise

    def get_character(self, character_name: str) -> Optional[Dict]:
        """
        Retrieve character from library

        Args:
            character_name: Name of the character

        Returns:
            Character profile dictionary or None if not found
        """
        # Check cache first
        if character_name in self.character_cache:
            return self.character_cache[character_name]

        # Check disk
        profile_filename = f"{character_name.lower().replace(' ', '_')}_profile.json"
        profile_path = os.path.join(self.character_library_dir, "profiles", profile_filename)

        if os.path.exists(profile_path):
            with open(profile_path, 'r') as f:
                character_profile = json.load(f)
                self.character_cache[character_name] = character_profile
                print(f"Loaded character '{character_name}' from library")
                return character_profile

        # Check LlamaIndex
        if self.story_memory and self.enable_character_library:
            char_info = self.story_memory.get_character_info(character_name)
            if char_info:
                print(f"Found character '{character_name}' in story memory")
                return char_info

        print(f"Character '{character_name}' not found in library")
        return None

    def generate_panel_with_character(
        self,
        scene_description: str,
        character_name: str,
        panel_number: int = 1,
        create_if_missing: bool = False,
        character_description: Optional[str] = None
    ) -> Image.Image:
        """
        Generate a comic panel with character consistency

        Args:
            scene_description: Description of the scene
            character_name: Name of the character
            panel_number: Panel number in sequence
            create_if_missing: If True, create character if not in library
            character_description: Character description (required if create_if_missing=True)

        Returns:
            Generated PIL Image
        """
        if not self.nano_banana:
            raise ValueError("Nano Banana is not enabled")

        # Get or create character
        character_profile = self.get_character(character_name)

        if not character_profile:
            if create_if_missing and character_description:
                print(f"Character '{character_name}' not found. Creating new character...")
                character_profile = self.create_character(
                    character_description=character_description,
                    character_name=character_name
                )
            else:
                raise ValueError(f"Character '{character_name}' not found in library. Set create_if_missing=True to create.")

        # Load reference images
        reference_images = []
        for img_path in character_profile.get("reference_images", []):
            if os.path.exists(img_path):
                reference_images.append(Image.open(img_path))

        if not reference_images:
            print(f"Warning: No reference images found for {character_name}")

        # Generate consistency prompt
        consistency_prompt = self.character_extractor.generate_consistency_prompt(character_profile)

        # Generate panel
        print(f"\nGenerating Panel {panel_number} with character '{character_name}'...")
        generated_image = self.nano_banana.generate_with_character_consistency(
            prompt=scene_description,
            character_references=reference_images,
            character_description=consistency_prompt,
            scene_number=panel_number
        )

        # Verify consistency
        if panel_number > 1 and reference_images and self.character_extractor:
            print("Verifying character consistency...")
            consistency_result = self.character_extractor.compare_character_consistency(
                image1=reference_images[0],
                image2=generated_image,
                character_name=character_name
            )
            print(f"Consistency score: {consistency_result.get('consistency_score', 0.0):.2%}")

            if consistency_result.get('consistency_score', 0.0) < 0.7:
                print(f"⚠ Low consistency score! Differences: {consistency_result.get('differences', [])}")

        return generated_image

    def generate_story_panels(
        self,
        story_panels: List[Dict],
        main_character_name: str,
        main_character_description: Optional[str] = None
    ) -> List[Image.Image]:
        """
        Generate all panels for a story with character consistency

        Args:
            story_panels: List of panel dictionaries with 'description' field
            main_character_name: Name of the main character
            main_character_description: Character description (for first-time creation)

        Returns:
            List of generated PIL Images
        """
        generated_panels = []

        # Get or create character
        character_profile = self.get_character(main_character_name)

        if not character_profile and main_character_description:
            print(f"Creating new character '{main_character_name}'...")
            character_profile = self.create_character(
                character_description=main_character_description,
                character_name=main_character_name
            )

        # Generate each panel
        for i, panel in enumerate(story_panels, 1):
            scene_description = panel.get("description", "")

            try:
                panel_image = self.generate_panel_with_character(
                    scene_description=scene_description,
                    character_name=main_character_name,
                    panel_number=i,
                    create_if_missing=False
                )
                generated_panels.append(panel_image)
                print(f"✓ Panel {i}/{len(story_panels)} generated")

            except Exception as e:
                print(f"✗ Error generating panel {i}: {e}")
                # Create placeholder
                placeholder = Image.new('RGB', (1024, 1024), color='gray')
                generated_panels.append(placeholder)

        return generated_panels

    def generate_panel_with_multiple_characters(
        self,
        scene_description: str,
        character_names: List[str],
        panel_number: int = 1,
        create_missing: bool = False,
        character_descriptions: Optional[Dict[str, str]] = None
    ) -> Image.Image:
        """
        Generate a panel with MULTIPLE characters maintaining consistency

        Args:
            scene_description: Description of the scene
            character_names: List of character names to include
            panel_number: Panel number in sequence
            create_missing: Create characters if not in library
            character_descriptions: Dict mapping character names to descriptions (for creation)

        Returns:
            Generated PIL Image with all characters

        Note:
            Nano Banana Pro can maintain up to 5 people and 6 objects simultaneously.
            For best results, limit to 2-3 main characters per panel.
        """
        if not self.nano_banana:
            raise ValueError("Nano Banana is not enabled")

        if len(character_names) > 5:
            print(f"⚠️ Warning: {len(character_names)} characters requested. Nano Banana Pro supports up to 5 people.")
            print(f"   Using first 5 characters only.")
            character_names = character_names[:5]

        # Get or create all characters
        all_character_profiles = []
        all_reference_images = []
        combined_consistency_prompts = []

        for char_name in character_names:
            char_profile = self.get_character(char_name)

            # Create if missing
            if not char_profile:
                if create_missing and character_descriptions and char_name in character_descriptions:
                    print(f"Creating new character '{char_name}'...")
                    char_profile = self.create_character(
                        character_description=character_descriptions[char_name],
                        character_name=char_name
                    )
                else:
                    raise ValueError(f"Character '{char_name}' not found. Set create_missing=True and provide description.")

            all_character_profiles.append(char_profile)

            # Load reference images (limit per character to fit within 14 total)
            refs_per_character = min(14 // len(character_names), 4)
            char_refs = []
            for img_path in char_profile.get("reference_images", [])[:refs_per_character]:
                if os.path.exists(img_path):
                    char_refs.append(Image.open(img_path))

            all_reference_images.extend(char_refs)

            # Build consistency prompt
            consistency_prompt = self.character_extractor.generate_consistency_prompt(char_profile)
            combined_consistency_prompts.append(f"CHARACTER {len(all_character_profiles)} ({char_name}):\n{consistency_prompt}")

        # Build combined prompt
        characters_description = "\n\n".join(combined_consistency_prompts)

        full_prompt = f"""IMPORTANT: Maintain EXACT character appearances from reference images.

{characters_description}

SCENE (Panel {panel_number}):
{scene_description}

Generate this scene with ALL {len(character_names)} characters maintaining their EXACT appearance from the references.
Ensure each character is visually distinct and identifiable.
"""

        print(f"\nGenerating Panel {panel_number} with {len(character_names)} characters:")
        for name in character_names:
            print(f"  - {name}")

        # Generate with all reference images
        try:
            generated_image = self.nano_banana.generate_image(
                prompt=full_prompt,
                reference_images=all_reference_images,
                aspect_ratio="1:1",
                num_images=1
            )[0]

            print(f"✓ Panel {panel_number} generated with {len(character_names)} characters")
            return generated_image

        except Exception as e:
            print(f"✗ Error generating multi-character panel: {e}")
            raise

    def list_characters(self) -> List[str]:
        """List all characters in the library"""
        characters = []

        # Check profiles directory
        profiles_dir = os.path.join(self.character_library_dir, "profiles")
        if os.path.exists(profiles_dir):
            for filename in os.listdir(profiles_dir):
                if filename.endswith("_profile.json"):
                    character_name = filename.replace("_profile.json", "").replace("_", " ").title()
                    characters.append(character_name)

        return sorted(characters)

    def get_character_info_summary(self, character_name: str) -> str:
        """Get a human-readable summary of character information"""
        character_profile = self.get_character(character_name)

        if not character_profile:
            return f"Character '{character_name}' not found."

        summary = f"=== Character: {character_name} ===\n"
        summary += f"Physical Description: {character_profile.get('physical_description', 'N/A')}\n"
        summary += f"Age Range: {character_profile.get('age_range', 'N/A')}\n"
        summary += f"Hair: {character_profile.get('hair', 'N/A')}\n"
        summary += f"Clothing: {character_profile.get('clothing', 'N/A')}\n"
        summary += f"Body Type: {character_profile.get('body_type', 'N/A')}\n"
        summary += f"Art Style: {character_profile.get('style', 'N/A')}\n"
        summary += f"Reference Images: {len(character_profile.get('reference_images', []))}\n"

        return summary

    def generate_comic_from_prompt(
        self,
        story_prompt: str,
        num_panels: int = 3,
        art_style: str = "manga",
        include_surprise_character: bool = True,
        output_dir: str = "./auto_generated_comics"
    ) -> Dict:
        """
        FULLY AUTOMATIC: Generate complete comic from just a text prompt

        This method automatically:
        1. Extracts ALL characters from the prompt
        2. Adds a surprise character (optional)
        3. Generates detailed descriptions for each character
        4. Creates character reference sheets
        5. Breaks story into panel-by-panel scenes
        6. Generates all panels with consistent characters
        7. Saves everything to disk

        Args:
            story_prompt: User's story idea (e.g., "A yogi teaching meditation...")
            num_panels: Number of comic panels to generate
            art_style: Art style (manga, anime, comic, etc.)
            include_surprise_character: Add a surprise character for intrigue
            output_dir: Directory to save generated comic

        Returns:
            Dictionary with:
            - characters: Dict of created characters
            - panels: List of generated panel images
            - panel_info: List of panel descriptions
            - surprise_character: Name of surprise character (if added)

        Example:
            ```python
            manager = CharacterConsistencyManager()
            result = manager.generate_comic_from_prompt(
                "A yogi teaching a young aspirant meditation and levitation in a jungle"
            )
            # Automatically creates: Yogi, Aspirant, + Surprise character
            # Generates 3 panels with all characters consistent
            ```
        """
        if not self.auto_intelligence or not self.nano_banana:
            raise ValueError("Auto Intelligence and Nano Banana must be enabled")

        print("\n" + "="*70)
        print("  🤖 AUTOMATIC COMIC GENERATION")
        print("="*70)
        print(f"Story: {story_prompt}")
        print(f"Panels: {num_panels}")
        print(f"Style: {art_style}")
        print(f"Surprise Character: {'Yes' if include_surprise_character else 'No'}")
        print("="*70)

        # Create output directory
        os.makedirs(output_dir, exist_ok=True)

        # Step 1: Extract characters automatically
        print("\n📖 Step 1: Extracting characters from prompt...")
        characters_dict = self.auto_intelligence.extract_characters_from_prompt(
            story_prompt=story_prompt,
            art_style=art_style,
            include_surprise_character=include_surprise_character
        )

        surprise_char_name = None
        if include_surprise_character:
            # Identify surprise character (last one added)
            all_names = list(characters_dict.keys())
            surprise_char_name = all_names[-1] if len(all_names) > 2 else None

        # Step 2: Create all characters with reference sheets
        print(f"\n🎨 Step 2: Creating {len(characters_dict)} characters...")
        created_characters = {}

        for char_name, char_description in characters_dict.items():
            # Check if character already exists
            existing = self.get_character(char_name)
            if existing:
                print(f"  ✓ {char_name} already in library (reusing)")
                created_characters[char_name] = existing
            else:
                char_type = "🎁 SURPRISE" if char_name == surprise_char_name else "📖"
                print(f"  {char_type} Creating {char_name}...")
                try:
                    profile = self.create_character(
                        character_description=char_description,
                        character_name=char_name,
                        style=art_style
                    )
                    created_characters[char_name] = profile
                    print(f"  ✓ {char_name} created")
                except Exception as e:
                    print(f"  ✗ Error creating {char_name}: {e}")

        # Step 3: Generate panel breakdown
        print(f"\n📋 Step 3: Breaking story into {num_panels} panels...")
        panel_breakdown = self.auto_intelligence.generate_panel_breakdown(
            story_prompt=story_prompt,
            num_panels=num_panels
        )

        # Step 4: Generate each panel
        print(f"\n🖼️  Step 4: Generating {num_panels} panels...")
        generated_panels = []
        panel_info = []

        for i, panel_desc in enumerate(panel_breakdown[:num_panels], 1):
            scene_desc = panel_desc.get("scene_description", story_prompt)
            panel_characters = panel_desc.get("characters", list(created_characters.keys()))

            # Filter to only use created characters
            valid_characters = [c for c in panel_characters if c in created_characters]

            # If no characters specified, use all created characters
            if not valid_characters:
                valid_characters = list(created_characters.keys())[:3]  # Limit to 3 for best results

            print(f"\n  Panel {i}/{num_panels}:")
            print(f"    Scene: {scene_desc[:80]}...")
            print(f"    Characters: {', '.join(valid_characters)}")

            try:
                # Generate panel with all characters
                if len(valid_characters) == 1:
                    panel_image = self.generate_panel_with_character(
                        scene_description=scene_desc,
                        character_name=valid_characters[0],
                        panel_number=i
                    )
                else:
                    panel_image = self.generate_panel_with_multiple_characters(
                        scene_description=scene_desc,
                        character_names=valid_characters,
                        panel_number=i
                    )

                # Save panel
                panel_filename = f"panel_{i:02d}.png"
                panel_path = os.path.join(output_dir, panel_filename)
                panel_image.save(panel_path)

                generated_panels.append(panel_image)
                panel_info.append({
                    "panel_number": i,
                    "characters": valid_characters,
                    "scene": scene_desc,
                    "filename": panel_filename,
                    "path": panel_path
                })

                print(f"    ✓ Saved: {panel_filename}")

            except Exception as e:
                print(f"    ✗ Error generating panel {i}: {e}")

        # Step 5: Summary
        print("\n" + "="*70)
        print("  ✅ COMIC GENERATION COMPLETE")
        print("="*70)
        print(f"\nCharacters Created: {len(created_characters)}")
        for char_name in created_characters.keys():
            char_type = "🎁 SURPRISE" if char_name == surprise_char_name else "📖"
            print(f"  {char_type} {char_name}")

        print(f"\nPanels Generated: {len(generated_panels)}/{num_panels}")
        print(f"Output Directory: {output_dir}")

        if surprise_char_name:
            print(f"\n🎁 Surprise Character: {surprise_char_name}")
            print(f"   Check the panels to see how they appear in the story!")

        # Return results
        return {
            "characters": created_characters,
            "panels": generated_panels,
            "panel_info": panel_info,
            "surprise_character": surprise_char_name,
            "output_dir": output_dir,
            "story_prompt": story_prompt
        }


# Example usage
if __name__ == "__main__":
    print("=== Character Consistency Manager Test ===\n")

    # Initialize manager
    manager = CharacterConsistencyManager()

    # Example: Create a character
    if manager.nano_banana:
        print("Test 1: Creating a character")
        try:
            yogi_profile = manager.create_character(
                character_description="Elderly Indian yogi with long white beard, saffron robes, peaceful expression, meditation beads",
                character_name="Master Yogi",
                style="manga"
            )
            print("✓ Character created successfully\n")
        except Exception as e:
            print(f"✗ Test 1 failed: {e}\n")

        # Example: List characters
        print("Test 2: Listing characters")
        characters = manager.list_characters()
        print(f"Characters in library: {characters}\n")

        # Example: Generate panel with character
        if characters:
            print("Test 3: Generating panel with character")
            try:
                panel = manager.generate_panel_with_character(
                    scene_description="The yogi meditating peacefully in a vibrant jungle",
                    character_name=characters[0],
                    panel_number=1
                )
                print("✓ Panel generated successfully\n")
            except Exception as e:
                print(f"✗ Test 3 failed: {e}\n")
    else:
        print("Nano Banana not enabled. Set ENABLE_NANO_BANANA=true in .env to test.")

    print("=== Tests Complete ===")