services/creative_modifier.py

"""
Creative Modifier Service
Analyzes uploaded creatives and generates modified versions based on user-provided angles/concepts.
"""

import os
import sys
import logging
import time
import uuid
from datetime import datetime
from typing import Dict, Any, Optional, Tuple, List

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from pydantic import BaseModel
from config import settings
from services.llm import llm_service
from services.image import image_service

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger("creative_modifier")

# Optional R2 storage import
try:
    from services.r2_storage import get_r2_storage
    r2_storage_available = True
except ImportError:
    r2_storage_available = False


class CreativeAnalysis(BaseModel):
    """Structured analysis of an uploaded creative."""
    visual_style: str
    color_palette: List[str]
    mood: str
    composition: str
    subject_matter: str
    text_content: Optional[str] = None
    current_angle: Optional[str] = None
    current_concept: Optional[str] = None
    target_audience: Optional[str] = None
    strengths: List[str]
    areas_for_improvement: List[str]


class CreativeModifierService:
    """Service for analyzing and modifying creative images."""
    
    def __init__(self):
        """Initialize the creative modifier service."""
        self.output_dir = settings.output_dir
        os.makedirs(self.output_dir, exist_ok=True)
    
    def _should_save_locally(self) -> bool:
        """Determine if images should be saved locally based on environment settings."""
        if settings.environment.lower() == "production":
            return settings.save_images_locally
        return True
    
    def _save_image_locally(self, image_bytes: bytes, filename: str) -> Optional[str]:
        """Conditionally save image locally based on environment settings."""
        if not self._should_save_locally():
            return None
        
        try:
            filepath = os.path.join(self.output_dir, filename)
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            with open(filepath, "wb") as f:
                f.write(image_bytes)
            return filepath
        except Exception as e:
            logger.warning(f"Failed to save image locally: {e}")
            return None
    
    async def analyze_creative(
        self,
        image_bytes: bytes,
    ) -> Dict[str, Any]:
        """
        Analyze an uploaded creative image using GPT-4 Vision.
        
        Args:
            image_bytes: Image file bytes to analyze
            
        Returns:
            Analysis results dictionary with structured data
        """
        start_time = time.time()
        logger.info("=" * 60)
        logger.info("Starting creative analysis")
        logger.info(f"Image size: {len(image_bytes)} bytes")
        
        system_prompt = """You are an expert creative director and marketing analyst specializing in ad creatives.
Your task is to thoroughly analyze advertising images to understand their visual style, messaging strategy, and effectiveness.
You must provide detailed, actionable insights that can be used to modify or improve the creative."""
        
        analysis_prompt = """Please analyze this advertising creative in detail. Provide your analysis as a JSON object with the following structure:

{
    "visual_style": "Description of the visual style (e.g., 'photorealistic', 'illustrated', 'minimalist', 'vibrant', 'muted')",
    "color_palette": ["List of dominant colors used"],
    "mood": "The emotional mood conveyed (e.g., 'urgent', 'calm', 'exciting', 'trustworthy')",
    "composition": "Description of the layout and composition (e.g., 'centered subject', 'rule of thirds', 'text-heavy')",
    "subject_matter": "What is depicted in the image",
    "text_content": "Any text visible in the image (or null if none)",
    "current_angle": "The psychological angle being used (e.g., 'fear of missing out', 'social proof', 'authority')",
    "current_concept": "The creative concept/format (e.g., 'testimonial', 'before/after', 'lifestyle shot')",
    "target_audience": "Who this creative seems to target",
    "strengths": ["List of what works well in this creative"],
    "areas_for_improvement": ["List of potential improvements"]
}

Be specific and detailed in your analysis. If you cannot determine something with confidence, make your best assessment based on visual cues."""
        
        try:
            logger.info("Calling vision API for creative analysis...")
            analysis_text = await llm_service.analyze_image_with_vision(
                image_bytes=image_bytes,
                analysis_prompt=analysis_prompt,
                system_prompt=system_prompt,
            )
            
            # Parse the JSON response
            import json
            analysis_text = analysis_text.strip()
            if analysis_text.startswith("```json"):
                analysis_text = analysis_text[7:]
            if analysis_text.startswith("```"):
                analysis_text = analysis_text[3:]
            if analysis_text.endswith("```"):
                analysis_text = analysis_text[:-3]
            
            analysis_data = json.loads(analysis_text.strip())
            
            elapsed_time = time.time() - start_time
            logger.info(f"✓ Creative analysis completed successfully in {elapsed_time:.2f}s")
            
            # Generate suggested angles and concepts based on analysis
            suggested_angles = self._generate_suggested_angles(analysis_data)
            suggested_concepts = self._generate_suggested_concepts(analysis_data)
            
            return {
                "status": "success",
                "analysis": analysis_data,
                "suggested_angles": suggested_angles,
                "suggested_concepts": suggested_concepts,
            }
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse analysis JSON: {e}")
            logger.error(f"Raw response: {analysis_text[:500]}...")
            return {
                "status": "error",
                "error": f"Failed to parse analysis response: {str(e)}",
            }
        except Exception as e:
            elapsed_time = time.time() - start_time
            logger.error(f"✗ Creative analysis failed after {elapsed_time:.2f}s: {str(e)}")
            return {
                "status": "error",
                "error": str(e),
            }
    
    def _generate_suggested_angles(self, analysis: Dict[str, Any]) -> List[str]:
        """Generate suggested angles based on the analysis using all available angles."""
        from data.angles import get_all_angles, get_random_angles
        
        current_angle = analysis.get("current_angle", "").lower()
        
        # Get all angles from the data module
        all_angles = get_all_angles()
        
        # Format angles as "Name (Trigger)" for display
        formatted_angles = []
        for angle in all_angles:
            formatted = f"{angle['name']} ({angle['trigger']})"
            # Filter out angles that match the current angle
            if current_angle and current_angle in angle['name'].lower():
                continue
            formatted_angles.append(formatted)
        
        # Get diverse random angles for suggestions (from different categories)
        random_angles = get_random_angles(count=12, diverse=True)
        suggestions = []
        for angle in random_angles:
            formatted = f"{angle['name']} ({angle['trigger']})"
            if current_angle and current_angle in angle['name'].lower():
                continue
            suggestions.append(formatted)
        
        return suggestions[:8]  # Return top 8 suggestions
    
    def _generate_suggested_concepts(self, analysis: Dict[str, Any]) -> List[str]:
        """Generate suggested concepts based on the analysis using all available concepts."""
        from data.concepts import get_all_concepts, get_random_concepts
        
        current_concept = analysis.get("current_concept", "").lower()
        
        # Get all concepts from the data module
        all_concepts = get_all_concepts()
        
        # Format concepts as "Name - Structure" for display
        formatted_concepts = []
        for concept in all_concepts:
            formatted = f"{concept['name']} ({concept['structure']})"
            # Filter out concepts that match the current concept
            if current_concept and current_concept in concept['name'].lower():
                continue
            formatted_concepts.append(formatted)
        
        # Get diverse random concepts for suggestions (from different categories)
        random_concepts = get_random_concepts(count=12, diverse=True)
        suggestions = []
        for concept in random_concepts:
            formatted = f"{concept['name']} ({concept['structure']})"
            if current_concept and current_concept in concept['name'].lower():
                continue
            suggestions.append(formatted)
        
        return suggestions[:8]  # Return top 8 suggestions
    
    async def generate_modification_prompt(
        self,
        analysis: Dict[str, Any],
        user_angle: Optional[str] = None,
        user_concept: Optional[str] = None,
        mode: str = "modify",
        user_prompt: Optional[str] = None,
    ) -> str:
        """
        Generate a prompt for modifying the creative based on analysis and user input.
        
        Args:
            analysis: The creative analysis data
            user_angle: User-provided angle to apply
            user_concept: User-provided concept to apply
            mode: "modify" for image-to-image, "inspired" for new generation
            user_prompt: Custom user instructions for modification
            
        Returns:
            Generated prompt string
        """
        logger.info(f"Generating modification prompt (mode: {mode})")
        logger.info(f"User angle: {user_angle}")
        logger.info(f"User concept: {user_concept}")
        logger.info(f"User prompt: {user_prompt}")
        
        # If user provided a custom prompt, use it directly
        if user_prompt and user_prompt.strip():
            logger.info("Using custom user prompt directly")
            return user_prompt.strip()
        
        system_prompt = """You are an expert advertising creative director with 20+ years experience.
    Your task is to create seamless, organic modifications that enhance existing creatives without appearing forced.
    You understand that effective ads feel authentic and natural, not like concepts were "pasted on."
    Focus on subtlety, consistency, and maintaining the original's visual language while applying new psychological angles.
    CRITICAL: ALL generated images MUST be photorealistic. Never mention "AI-generated" or similar terms."""
        
        # Try to enrich angle/concept with library data if not already detailed
        angle_info = user_angle
        concept_info = user_concept
        
        try:
            from data.angles import get_all_angles
            from data.concepts import get_all_concepts
            
            # Lookup angle details if it's just a name
            if user_angle and "(" not in user_angle:
                all_angles = get_all_angles()
                for a in all_angles:
                    if a["name"].lower() == user_angle.lower():
                        angle_info = f"{a['name']} (Psychological trigger: {a['trigger']})"
                        break
            
            # Lookup concept details if it's just a name
            if user_concept and "(" not in user_concept:
                all_concepts = get_all_concepts()
                for c in all_concepts:
                    if c["name"].lower() == user_concept.lower():
                        concept_info = f"{c['name']} (Visual structure: {c['structure']})"
                        break
        except Exception as e:
            logger.warning(f"Failed to enrich angle/concept data: {e}")

        # Extract original creative's essence for natural integration
        original_mood = analysis.get('mood', 'Unknown')
        original_style = analysis.get('visual_style', 'Unknown')
        color_palette = analysis.get('color_palette', [])
        subject_matter = analysis.get('subject_matter', 'Unknown')
        composition = analysis.get('composition', 'Unknown')
        
        # Guard against common hallucinations (e.g. "fan")
        if "fan of" in subject_matter.lower() and "bills" in subject_matter.lower():
            subject_matter = subject_matter.replace("fan of", "fanned-out stack of").replace("bills", "money/dollar bills")

        if mode == "modify":
            # For image-to-image: visually clear modifications
            prompt_request = f"""ORIGINAL IMAGE CONTEXT:
    - Subject: {subject_matter}
    - Mood: {original_mood}
    - Style: {original_style}
    - Composition: {composition}
    - Colors: {', '.join(color_palette[:3]) if color_palette else 'Natural tones'}

    TRANSFORMATION TASK:
    - Apply Angle: {angle_info or 'natural enhancement'}
    - Apply Concept: {concept_info or 'subtle adjustment'}

    Generate a clear, descriptive transformation prompt (20-40 words) that:
    1. Specifically describes the VISUAL change needed to reflect the new angle and concept.
    2. Makes the transformation feel like a professional edit of the original - NOT a different image.
    3. Keeps the core subjects ({subject_matter}) but adapts their presentation, lighting, or surrounding elements.
    4. Maintain the {original_style} style and {original_mood} tone.
    5. Be explicit about what to change (e.g., "Change the lighting to be more dramatic", "Rearrange elements for [concept]")

    CRITICAL: The instruction must be strong enough for an image-to-image AI to actually make a visible change.
    
    Return ONLY the prompt text, no explanations."""
        else:
            # For inspired generation: create new but keep the same authentic vibe
            prompt_request = f"""ORIGINAL CAMPAIGN INSPIRATION:
    - Core Vibe: {original_mood} mood, {original_style} aesthetic
    - Visual Language: {composition}, {subject_matter}

    NEW AD CREATIVE REQUIREMENTS:
    - Target Angle: {angle_info or 'natural persuasion'}
    - Target Concept: {concept_info or 'authentic presentation'}

    Generate a premium, authentic advertising photography prompt (60-90 words) that:
    1. Creates a NEW photo that looks like it belongs in the same campaign as the original.
    2. Centers the visual around {concept_info or 'the concept'}.
    3. Conveys the {angle_info or 'the angle'} through unposed, authentic human moments and natural lighting.
    4. Maintains high photorealistic quality with realistic textures and real-world lighting.
    5. Do NOT describe a generic stock photo; describe a high-end, cinematic brand image.

    Return ONLY the prompt text, no explanations."""
        
        try:
            prompt = await llm_service.generate(
                prompt=prompt_request,
                system_prompt=system_prompt,
                temperature=0.7,
            )
            
            # Clean up and ensure natural language
            prompt = prompt.strip().strip('"').strip("'")
            
            # Enhance for seamlessness if needed
            seamless_keywords = ["seamless", "organic", "natural", "authentic", "integrated"]
            has_seamless = any(keyword in prompt.lower() for keyword in seamless_keywords)
            
            photorealistic_keywords = [
                "photorealistic", "realistic photograph", "cinematic photography",
                "authentic photo", "natural lighting", "real-world"
            ]
            has_photorealistic = any(keyword.lower() in prompt.lower() for keyword in photorealistic_keywords)
            
            # Add subtle guidance if missing key elements
            if not has_seamless and mode == "modify":
                prompt = f"Seamlessly integrated, {prompt}"
            
            if not has_photorealistic:
                # Add natural photography terms
                prompt = f"Cinematic photography, authentic moment, natural lighting. {prompt}"
                logger.info("Added natural photography emphasis to prompt")
            
            # Remove any phrases that sound artificial or forced
            forced_phrases = [
                "add", "insert", "place", "include the concept of",
                "visibly show", "clearly demonstrate", "obviously"
            ]
            
            for phrase in forced_phrases:
                if phrase in prompt.lower():
                    # Replace with more natural alternatives
                    if phrase == "add":
                        prompt = prompt.lower().replace("add", "naturally incorporate")
                    elif phrase == "insert":
                        prompt = prompt.lower().replace("insert", "subtly integrate")
                    elif phrase == "place":
                        prompt = prompt.lower().replace("place", "position naturally")
                    elif "visibly show" in prompt.lower():
                        prompt = prompt.lower().replace("visibly show", "suggest through")
            
            logger.info(f"Generated natural prompt: {prompt[:100]}...")
            
            return prompt
        except Exception as e:
            logger.error(f"Failed to generate modification prompt: {e}")
            # Fallback prompts with emphasis on natural integration
            if mode == "modify":
                return f"Cinematic photography, seamless integration. Subtly enhance {subject_matter} to naturally incorporate {user_angle or 'emotional resonance'} through {user_concept or 'visual storytelling'}. Authentic moment, natural lighting, feels like original campaign."
            else:
                return f"Cinematic advertising photograph, authentic human moment. {subject_matter} presented with {original_mood} emotional tone, naturally integrating {user_angle or 'persuasive angle'} through {user_concept or 'visual concept'}. Real-world lighting, natural skin textures, campaign-consistent styling."

    async def modify_creative(
        self,
        image_url: str,
        analysis: Dict[str, Any],
        user_angle: Optional[str] = None,
        user_concept: Optional[str] = None,
        mode: str = "modify",
        image_model: Optional[str] = None,
        user_prompt: Optional[str] = None,
        width: int = 1024,
        height: int = 1024,
    ) -> Dict[str, Any]:
        """
        Modify or generate a new creative based on the original and user input.
        
        Args:
            image_url: URL of the original image
            analysis: Creative analysis data
            user_angle: User-provided angle
            user_concept: User-provided concept
            mode: "modify" for image-to-image, "inspired" for new generation
            image_model: Model to use for generation
            user_prompt: Custom user instructions for modification
            width: Output width
            height: Output height
            
        Returns:
            Result dictionary with generated image info
        """
        workflow_start = time.time()
        logger.info("=" * 80)
        logger.info("CREATIVE MODIFICATION WORKFLOW STARTED")
        logger.info(f"Mode: {mode}")
        logger.info(f"Image URL: {image_url}")
        logger.info(f"User angle: {user_angle}")
        logger.info(f"User concept: {user_concept}")
        logger.info(f"User prompt: {user_prompt}")
        logger.info(f"Image model: {image_model}")
        logger.info("=" * 80)
        
        result = {
            "status": "pending",
            "prompt": None,
            "image": None,
            "error": None,
        }
        
        # Step 1: Generate modification prompt
        logger.info("STEP 1: Generating modification prompt...")
        prompt = await self.generate_modification_prompt(
            analysis=analysis,
            user_angle=user_angle,
            user_concept=user_concept,
            mode=mode,
            user_prompt=user_prompt,
        )
        result["prompt"] = prompt
        logger.info(f"✓ Generated prompt: {prompt}")
        
        # Step 2: Generate/modify image
        logger.info("STEP 2: Generating modified image...")
        try:
            if mode == "modify":
                # For subtle modifications with image-to-image
                # Note: nano-banana models don't support guidance_scale or strength parameters
                # The model will naturally preserve the original image based on the prompt
                generation_params = {
                    "prompt": prompt,
                    "model_key": image_model or "nano-banana",
                    "width": width,
                    "height": height,
                    "image_url": image_url,
                }
            else:
                # For inspired generation (text-to-image)
                generation_params = {
                    "prompt": prompt,
                    "model_key": image_model or "nano-banana",
                    "width": width,
                    "height": height,
                }
            
            image_bytes, model_used, generated_url = await image_service.generate(**generation_params)
            
            if not image_bytes:
                raise Exception("Image generation returned no data")
            
            logger.info(f"✓ Image generated successfully ({len(image_bytes)} bytes)")
            
        except Exception as e:
            logger.error(f"✗ Image generation failed: {e}")
            result["status"] = "error"
            result["error"] = f"Image generation failed: {str(e)}"
            return result
        
        # Step 3: Save and upload image
        logger.info("STEP 3: Saving generated image...")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        unique_id = uuid.uuid4().hex[:8]
        filename = f"modified_{timestamp}_{unique_id}.png"
        
        # Upload to R2 if available
        r2_url = None
        if r2_storage_available:
            try:
                logger.info("Uploading to R2 storage...")
                r2_storage = get_r2_storage()
                if r2_storage:
                    r2_url = r2_storage.upload_image(
                        image_bytes=image_bytes,
                        filename=filename,
                        niche="modified",
                    )
                    logger.info(f"✓ Uploaded to R2: {r2_url}")
            except Exception as e:
                logger.warning(f"R2 upload failed: {e}")
        
        # Save locally if configured
        filepath = self._save_image_locally(image_bytes, filename)
        if filepath:
            logger.info(f"✓ Saved locally: {filepath}")
        
        # Determine final URL
        final_url = r2_url or generated_url
        
        result["status"] = "success"
        result["image"] = {
            "filename": filename,
            "filepath": filepath,
            "image_url": final_url,
            "r2_url": r2_url,
            "model_used": model_used,
            "mode": mode,
            "applied_angle": user_angle,
            "applied_concept": user_concept,
        }
        
        total_time = time.time() - workflow_start
        logger.info("=" * 80)
        logger.info("✓ CREATIVE MODIFICATION COMPLETED SUCCESSFULLY")
        logger.info(f"Total time: {total_time:.2f}s")
        logger.info(f"Output URL: {final_url}")
        logger.info("=" * 80)
        
        return result


# Global instance
creative_modifier_service = CreativeModifierService()