1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3"""
4CIP Design Generator - Generate corporate identity mockups using Gemini Nano Banana
5 
6Uses Gemini's native image generation (Nano Banana Flash/Pro) for high-quality mockups.
7Supports text-and-image-to-image generation for using actual brand logos.
8 
9- gemini-2.5-flash-image: Fast generation, cost-effective (default)
10- gemini-3-pro-image-preview: Pro quality, 4K text rendering
11 
12Image Editing (text-and-image-to-image):
13  When --logo is provided, the script uses Gemini's image editing capability
14  to incorporate the actual logo into CIP mockups instead of generating one.
15"""
16 
17import argparse
18import json
19import os
20import sys
21from pathlib import Path
22from datetime import datetime
23 
24# Add parent directory for imports
25sys.path.insert(0, str(Path(__file__).parent))
26from core import search, get_cip_brief
27 
28# Model options
29MODELS = {
30    "flash": "gemini-2.5-flash-image",      # Nano Banana Flash - fast, default
31    "pro": "gemini-3-pro-image-preview"      # Nano Banana Pro - quality, 4K text
32}
33DEFAULT_MODEL = "flash"
34 
35 
36def load_logo_image(logo_path):
37    """Load logo image using PIL for Gemini image editing"""
38    try:
39        from PIL import Image
40    except ImportError:
41        print("Error: pillow package not installed.")
42        print("Install with: pip install pillow")
43        return None
44 
45    logo_path = Path(logo_path)
46    if not logo_path.exists():
47        print(f"Error: Logo file not found: {logo_path}")
48        return None
49 
50    try:
51        img = Image.open(logo_path)
52        # Convert to RGB if necessary (Gemini works best with RGB)
53        if img.mode in ('RGBA', 'P'):
54            # Create white background for transparent images
55            background = Image.new('RGB', img.size, (255, 255, 255))
56            if img.mode == 'RGBA':
57                background.paste(img, mask=img.split()[3])  # Use alpha channel as mask
58            else:
59                background.paste(img)
60            img = background
61        elif img.mode != 'RGB':
62            img = img.convert('RGB')
63        return img
64    except Exception as e:
65        print(f"Error loading logo: {e}")
66        return None
67 
68# Load environment variables
69def load_env():
70    """Load environment variables from .env files"""
71    env_paths = [
72        Path(__file__).parent.parent.parent / ".env",
73        Path.home() / ".claude" / "skills" / ".env",
74        Path.home() / ".claude" / ".env"
75    ]
76    for env_path in env_paths:
77        if env_path.exists():
78            with open(env_path) as f:
79                for line in f:
80                    line = line.strip()
81                    if line and not line.startswith("#") and "=" in line:
82                        key, value = line.split("=", 1)
83                        if key not in os.environ:
84                            os.environ[key] = value.strip('"\'')
85 
86load_env()
87 
88 
89def build_cip_prompt(deliverable, brand_name, style=None, industry=None, mockup=None, use_logo_image=False):
90    """Build an optimized prompt for CIP mockup generation
91 
92    Args:
93        deliverable: Type of deliverable (business card, letterhead, etc.)
94        brand_name: Name of the brand
95        style: Design style preference
96        industry: Industry for style recommendations
97        mockup: Mockup context override
98        use_logo_image: If True, prompt is optimized for image editing with logo
99    """
100 
101    # Get deliverable details
102    deliverable_info = search(deliverable, "deliverable", 1)
103    deliverable_data = deliverable_info.get("results", [{}])[0] if deliverable_info.get("results") else {}
104 
105    # Get style details
106    style_info = search(style or "corporate minimal", "style", 1) if style else {}
107    style_data = style_info.get("results", [{}])[0] if style_info.get("results") else {}
108 
109    # Get industry details
110    industry_info = search(industry or "technology", "industry", 1) if industry else {}
111    industry_data = industry_info.get("results", [{}])[0] if industry_info.get("results") else {}
112 
113    # Get mockup context
114    mockup_context = deliverable_data.get("Mockup Context", "clean professional")
115    if mockup:
116        mockup_info = search(mockup, "mockup", 1)
117        if mockup_info.get("results"):
118            mockup_data = mockup_info["results"][0]
119            mockup_context = mockup_data.get("Scene Description", mockup_context)
120 
121    # Build prompt components
122    deliverable_name = deliverable_data.get("Deliverable", deliverable)
123    description = deliverable_data.get("Description", "")
124    dimensions = deliverable_data.get("Dimensions", "")
125    logo_placement = deliverable_data.get("Logo Placement", "center")
126 
127    style_name = style_data.get("Style Name", style or "corporate")
128    primary_colors = style_data.get("Primary Colors", industry_data.get("Primary Colors", "#0F172A #FFFFFF"))
129    typography = style_data.get("Typography", industry_data.get("Typography", "clean sans-serif"))
130    materials = style_data.get("Materials", "premium quality")
131    finishes = style_data.get("Finishes", "professional")
132 
133    mood = style_data.get("Mood", industry_data.get("Mood", "professional"))
134 
135    # Construct the prompt - different for image editing vs pure generation
136    if use_logo_image:
137        # Image editing prompt: instructs to USE the provided logo image
138        prompt_parts = [
139            f"Create a professional corporate identity mockup photograph of a {deliverable_name}",
140            f"Use the EXACT logo from the provided image - do NOT modify or recreate the logo",
141            f"The logo MUST appear exactly as shown in the input image",
142            f"Place the logo on the {deliverable_name} at: {logo_placement}",
143            f"Brand name: '{brand_name}'",
144            f"{description}" if description else "",
145            f"Design style: {style_name}",
146            f"Color scheme matching the logo colors",
147            f"Materials: {materials} with {finishes} finish",
148            f"Setting: {mockup_context}",
149            f"Mood: {mood}",
150            "Photorealistic product photography",
151            "Soft natural lighting, professional studio quality",
152            "8K resolution, sharp details"
153        ]
154    else:
155        # Pure text-to-image prompt
156        prompt_parts = [
157            f"Professional corporate identity mockup photograph",
158            f"showing {deliverable_name} for brand '{brand_name}'",
159            f"{description}" if description else "",
160            f"{style_name} design style",
161            f"using colors {primary_colors}",
162            f"{typography} typography",
163            f"logo placement: {logo_placement}",
164            f"{materials} materials with {finishes} finish",
165            f"{mockup_context} setting",
166            f"{mood} mood",
167            "photorealistic product photography",
168            "soft natural lighting",
169            "high quality professional shot",
170            "8k resolution detailed"
171        ]
172 
173    prompt = ", ".join([p for p in prompt_parts if p])
174 
175    return {
176        "prompt": prompt,
177        "deliverable": deliverable_name,
178        "style": style_name,
179        "brand": brand_name,
180        "colors": primary_colors,
181        "mockup_context": mockup_context,
182        "logo_placement": logo_placement
183    }
184 
185 
186def generate_with_nano_banana(prompt_data, output_dir=None, model_key="flash", aspect_ratio="1:1", logo_image=None):
187    """Generate image using Gemini Nano Banana (native image generation)
188 
189    Supports two modes:
190    1. Text-to-image: Pure prompt-based generation (logo_image=None)
191    2. Image editing: Text-and-image-to-image using provided logo (logo_image=PIL.Image)
192 
193    Models:
194    - flash: gemini-2.5-flash-image (fast, cost-effective) - DEFAULT
195    - pro: gemini-3-pro-image-preview (quality, 4K text rendering)
196 
197    Args:
198        prompt_data: Dict with prompt, deliverable, brand, etc.
199        output_dir: Output directory for generated images
200        model_key: 'flash' or 'pro'
201        aspect_ratio: Output aspect ratio (1:1, 16:9, etc.)
202        logo_image: PIL.Image object of the brand logo for image editing mode
203    """
204    try:
205        from google import genai
206        from google.genai import types
207    except ImportError:
208        print("Error: google-genai package not installed.")
209        print("Install with: pip install google-genai")
210        return None
211 
212    api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
213    if not api_key:
214        print("Error: GEMINI_API_KEY or GOOGLE_API_KEY not set")
215        return None
216 
217    client = genai.Client(api_key=api_key)
218 
219    prompt = prompt_data["prompt"]
220    model_name = MODELS.get(model_key, MODELS[DEFAULT_MODEL])
221 
222    # Determine mode
223    mode = "image-editing" if logo_image else "text-to-image"
224 
225    print(f"\n🎨 Generating CIP mockup...")
226    print(f"   Mode: {mode}")
227    print(f"   Deliverable: {prompt_data['deliverable']}")
228    print(f"   Brand: {prompt_data['brand']}")
229    print(f"   Style: {prompt_data['style']}")
230    print(f"   Model: {model_name}")
231    print(f"   Context: {prompt_data['mockup_context']}")
232    if logo_image:
233        print(f"   Logo: Using provided image ({logo_image.size[0]}x{logo_image.size[1]})")
234 
235    try:
236        # Build contents: either just prompt or [prompt, image] for image editing
237        if logo_image:
238            # Image editing mode: pass both prompt and logo image
239            contents = [prompt, logo_image]
240        else:
241            # Text-to-image mode: just the prompt
242            contents = prompt
243 
244        # Use generate_content with response_modalities=['IMAGE'] for Nano Banana
245        response = client.models.generate_content(
246            model=model_name,
247            contents=contents,
248            config=types.GenerateContentConfig(
249                response_modalities=['IMAGE'],  # Uppercase required
250                image_config=types.ImageConfig(
251                    aspect_ratio=aspect_ratio
252                )
253            )
254        )
255 
256        # Extract image from response
257        if response.candidates and response.candidates[0].content.parts:
258            for part in response.candidates[0].content.parts:
259                if hasattr(part, 'inline_data') and part.inline_data:
260                    # Save image
261                    output_dir = output_dir or Path.cwd()
262                    output_dir = Path(output_dir)
263                    output_dir.mkdir(parents=True, exist_ok=True)
264 
265                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
266                    brand_slug = prompt_data["brand"].lower().replace(" ", "-")
267                    deliverable_slug = prompt_data["deliverable"].lower().replace(" ", "-")
268                    filename = f"{brand_slug}-{deliverable_slug}-{timestamp}.png"
269                    filepath = output_dir / filename
270 
271                    image_data = part.inline_data.data
272                    with open(filepath, "wb") as f:
273                        f.write(image_data)
274 
275                    print(f"\n✅ Generated: {filepath}")
276                    return str(filepath)
277 
278        print("No image generated in response")
279        return None
280 
281    except Exception as e:
282        print(f"Error generating image: {e}")
283        return None
284 
285 
286def generate_cip_set(brand_name, industry, style=None, deliverables=None, output_dir=None, model_key="flash", logo_path=None, aspect_ratio="1:1"):
287    """Generate a complete CIP set for a brand
288 
289    Args:
290        brand_name: Brand name to generate for
291        industry: Industry type for style recommendations
292        style: Optional specific style override
293        deliverables: List of deliverables to generate (default: core set)
294        output_dir: Output directory for images
295        model_key: 'flash' (fast) or 'pro' (quality)
296        logo_path: Path to brand logo image for image editing mode
297        aspect_ratio: Output aspect ratio
298    """
299 
300    # Load logo image if provided
301    logo_image = None
302    if logo_path:
303        logo_image = load_logo_image(logo_path)
304        if not logo_image:
305            print("Warning: Could not load logo, falling back to text-to-image mode")
306 
307    # Get CIP brief for the brand
308    brief = get_cip_brief(brand_name, industry, style)
309 
310    # Default deliverables if not specified
311    if not deliverables:
312        deliverables = ["business card", "letterhead", "office signage", "vehicle", "polo shirt"]
313 
314    results = []
315    for deliverable in deliverables:
316        prompt_data = build_cip_prompt(
317            deliverable=deliverable,
318            brand_name=brand_name,
319            style=brief.get("style", {}).get("Style Name"),
320            industry=industry,
321            use_logo_image=(logo_image is not None)
322        )
323 
324        filepath = generate_with_nano_banana(
325            prompt_data,
326            output_dir,
327            model_key=model_key,
328            aspect_ratio=aspect_ratio,
329            logo_image=logo_image
330        )
331        if filepath:
332            results.append({
333                "deliverable": deliverable,
334                "filepath": filepath,
335                "prompt": prompt_data["prompt"]
336            })
337 
338    return results
339 
340 
341def check_logo_required(brand_name, skip_prompt=False):
342    """Check if logo is required and suggest logo-design skill if not provided
343 
344    Returns:
345        str: 'continue' to proceed without logo, 'generate' to use logo-design skill, 'exit' to abort
346    """
347    if skip_prompt:
348        return 'continue'
349 
350    print(f"\n⚠️  No logo image provided for '{brand_name}'")
351    print("   Without a logo, AI will generate its own interpretation of the brand logo.")
352    print("")
353    print("   Options:")
354    print("   1. Continue without logo (AI-generated logo interpretation)")
355    print("   2. Generate a logo first using 'logo-design' skill")
356    print("   3. Exit and provide a logo path with --logo")
357    print("")
358 
359    try:
360        choice = input("   Enter choice [1/2/3] (default: 1): ").strip()
361        if choice == '2':
362            return 'generate'
363        elif choice == '3':
364            return 'exit'
365        return 'continue'
366    except (EOFError, KeyboardInterrupt):
367        return 'continue'
368 
369 
370def main():
371    parser = argparse.ArgumentParser(
372        description="Generate CIP mockups using Gemini Nano Banana",
373        formatter_class=argparse.RawDescriptionHelpFormatter,
374        epilog="""
375Examples:
376  # Generate with brand logo (RECOMMENDED)
377  python generate.py --brand "TopGroup" --logo /path/to/logo.png --deliverable "business card"
378 
379  # Generate CIP set with logo
380  python generate.py --brand "TopGroup" --logo /path/to/logo.png --industry "consulting" --set
381 
382  # Generate without logo (AI interprets brand)
383  python generate.py --brand "TechFlow" --deliverable "business card" --no-logo-prompt
384 
385  # Generate with Pro model (higher quality, 4K text)
386  python generate.py --brand "TechFlow" --logo logo.png --deliverable "business card" --model pro
387 
388  # Specify output directory and aspect ratio
389  python generate.py --brand "MyBrand" --logo logo.png --deliverable "vehicle" --output ./mockups --ratio 16:9
390 
391Models:
392  flash (default): gemini-2.5-flash-image - Fast, cost-effective
393  pro: gemini-3-pro-image-preview - Quality, 4K text rendering
394 
395Image Editing Mode:
396  When --logo is provided, uses Gemini's text-and-image-to-image capability
397  to incorporate your ACTUAL logo into the CIP mockups.
398        """
399    )
400 
401    parser.add_argument("--brand", "-b", required=True, help="Brand name")
402    parser.add_argument("--logo", "-l", help="Path to brand logo image (enables image editing mode)")
403    parser.add_argument("--deliverable", "-d", help="Single deliverable to generate")
404    parser.add_argument("--deliverables", help="Comma-separated list of deliverables")
405    parser.add_argument("--industry", "-i", default="technology", help="Industry type")
406    parser.add_argument("--style", "-s", help="Design style")
407    parser.add_argument("--mockup", "-m", help="Mockup context")
408    parser.add_argument("--set", action="store_true", help="Generate full CIP set")
409    parser.add_argument("--output", "-o", help="Output directory")
410    parser.add_argument("--model", default="flash", choices=["flash", "pro"], help="Model: flash (fast) or pro (quality)")
411    parser.add_argument("--ratio", default="1:1", help="Aspect ratio (1:1, 16:9, 4:3, etc.)")
412    parser.add_argument("--prompt-only", action="store_true", help="Only show prompt, don't generate")
413    parser.add_argument("--json", "-j", action="store_true", help="Output as JSON")
414    parser.add_argument("--no-logo-prompt", action="store_true", help="Skip logo prompt, proceed without logo")
415 
416    args = parser.parse_args()
417 
418    # Check if logo is provided, prompt user if not
419    logo_image = None
420    if args.logo:
421        logo_image = load_logo_image(args.logo)
422        if not logo_image:
423            print("Error: Could not load logo image")
424            sys.exit(1)
425    elif not args.prompt_only:
426        # No logo provided - ask user what to do
427        action = check_logo_required(args.brand, skip_prompt=args.no_logo_prompt)
428        if action == 'generate':
429            print("\n💡 To generate a logo, use the logo-design skill:")
430            print(f"   python ~/.claude/skills/design/scripts/logo/generate.py --brand \"{args.brand}\" --industry \"{args.industry}\"")
431            print("\n   Then re-run this command with --logo <generated_logo.png>")
432            sys.exit(0)
433        elif action == 'exit':
434            print("\n   Provide logo with: --logo /path/to/your/logo.png")
435            sys.exit(0)
436        # else: continue without logo
437 
438    use_logo = logo_image is not None
439 
440    if args.set or args.deliverables:
441        # Generate multiple deliverables
442        deliverables = args.deliverables.split(",") if args.deliverables else None
443 
444        if args.prompt_only:
445            results = []
446            deliverables = deliverables or ["business card", "letterhead", "office signage", "vehicle", "polo shirt"]
447            for d in deliverables:
448                prompt_data = build_cip_prompt(d, args.brand, args.style, args.industry, args.mockup, use_logo_image=use_logo)
449                results.append(prompt_data)
450            if args.json:
451                print(json.dumps(results, indent=2))
452            else:
453                for r in results:
454                    print(f"\n{r['deliverable']}:\n{r['prompt']}\n")
455        else:
456            results = generate_cip_set(
457                args.brand, args.industry, args.style, deliverables, args.output,
458                model_key=args.model, logo_path=args.logo, aspect_ratio=args.ratio
459            )
460            if args.json:
461                print(json.dumps(results, indent=2))
462            else:
463                print(f"\n✅ Generated {len(results)} CIP mockups")
464    else:
465        # Generate single deliverable
466        deliverable = args.deliverable or "business card"
467        prompt_data = build_cip_prompt(deliverable, args.brand, args.style, args.industry, args.mockup, use_logo_image=use_logo)
468 
469        if args.prompt_only:
470            if args.json:
471                print(json.dumps(prompt_data, indent=2))
472            else:
473                print(f"\nPrompt:\n{prompt_data['prompt']}")
474        else:
475            filepath = generate_with_nano_banana(
476                prompt_data, args.output, model_key=args.model,
477                aspect_ratio=args.ratio, logo_image=logo_image
478            )
479            if args.json:
480                print(json.dumps({"filepath": filepath, **prompt_data}, indent=2))
481 
482 
483if __name__ == "__main__":
484    main()
485