1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3"""
4Icon Generation Script using Gemini 3.1 Pro Preview API
5Generates SVG icons via text generation (SVG is XML text format)
6 
7Model: gemini-3.1-pro-preview - best thinking, token efficiency, factual consistency
8 
9Usage:
10    python generate.py --prompt "settings gear icon" --style outlined
11    python generate.py --prompt "shopping cart" --style filled --color "#6366F1"
12    python generate.py --name "dashboard" --category navigation --style duotone
13    python generate.py --prompt "cloud upload" --batch 4 --output-dir ./icons
14    python generate.py --prompt "user profile" --sizes "16,24,32,48"
15"""
16 
17import argparse
18import json
19import os
20import re
21import sys
22import time
23from pathlib import Path
24from datetime import datetime
25 
26 
27def load_env():
28    """Load .env files in priority order"""
29    env_paths = [
30        Path(__file__).parent.parent.parent / ".env",
31        Path.home() / ".claude" / "skills" / ".env",
32        Path.home() / ".claude" / ".env"
33    ]
34    for env_path in env_paths:
35        if env_path.exists():
36            with open(env_path) as f:
37                for line in f:
38                    line = line.strip()
39                    if line and not line.startswith('#') and '=' in line:
40                        key, value = line.split('=', 1)
41                        if key not in os.environ:
42                            os.environ[key] = value.strip('"\'')
43 
44load_env()
45 
46try:
47    from google import genai
48    from google.genai import types
49except ImportError:
50    print("Error: google-genai package not installed.")
51    print("Install with: pip install google-genai")
52    sys.exit(1)
53 
54 
55# ============ CONFIGURATION ============
56GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
57MODEL = "gemini-3.1-pro-preview"
58 
59# Icon styles with SVG-specific instructions
60ICON_STYLES = {
61    "outlined": "outlined stroke icons, 2px stroke width, no fill, clean open paths",
62    "filled": "solid filled icons, no stroke, flat color fills, bold shapes",
63    "duotone": "duotone style with primary color at full opacity and secondary color at 30% opacity, layered shapes",
64    "thin": "thin line icons, 1px or 1.5px stroke width, delicate minimalist lines",
65    "bold": "bold thick line icons, 3px stroke width, heavy weight, impactful",
66    "rounded": "rounded icons with round line caps and joins, soft corners, friendly feel",
67    "sharp": "sharp angular icons, square line caps and mitered joins, precise edges",
68    "flat": "flat design icons, solid fills, no gradients or shadows, geometric simplicity",
69    "gradient": "linear or radial gradient fills, modern vibrant color transitions",
70    "glassmorphism": "glassmorphism style with semi-transparent fills, blur backdrop effect simulation, frosted glass",
71    "pixel": "pixel art style icons on a grid, retro 8-bit aesthetic, crisp edges",
72    "hand-drawn": "hand-drawn sketch style, slightly irregular strokes, organic feel, imperfect lines",
73    "isometric": "isometric 3D projection, 30-degree angles, dimensional depth",
74    "glyph": "simple glyph style, single solid shape, minimal detail, pictogram",
75    "animated-ready": "animated-ready SVG with named groups and IDs for CSS/JS animation targets",
76}
77 
78ICON_CATEGORIES = {
79    "navigation": "arrows, menus, hamburger, chevrons, home, back, forward, breadcrumb",
80    "action": "edit, delete, save, download, upload, share, copy, paste, print, search",
81    "communication": "email, chat, phone, video call, notification, bell, message bubble",
82    "media": "play, pause, stop, skip, volume, microphone, camera, image, gallery",
83    "file": "document, folder, archive, attachment, cloud, database, storage",
84    "user": "person, group, avatar, profile, settings, lock, key, shield",
85    "commerce": "cart, bag, wallet, credit card, receipt, tag, gift, store",
86    "data": "chart, graph, analytics, dashboard, table, filter, sort, calendar",
87    "development": "code, terminal, bug, git, API, server, database, deploy",
88    "social": "heart, star, thumbs up, bookmark, flag, trophy, badge, crown",
89    "weather": "sun, moon, cloud, rain, snow, wind, thunder, temperature",
90    "map": "pin, location, compass, globe, route, directions, map marker",
91}
92 
93# SVG generation prompt template
94SVG_PROMPT_TEMPLATE = """Generate a clean, production-ready SVG icon.
95 
96Requirements:
97- Output ONLY valid SVG code, nothing else
98- ViewBox: "0 0 {viewbox} {viewbox}"
99- Use currentColor for strokes/fills (inherits CSS color)
100- No embedded fonts or text elements unless specifically requested
101- No raster images or external references
102- Optimized paths with minimal nodes
103- Accessible: include <title> element with icon description
104{style_instructions}
105{color_instructions}
106{size_instructions}
107 
108Icon to generate: {prompt}
109 
110Output the SVG code only, wrapped in ```svg``` code block."""
111 
112SVG_BATCH_PROMPT_TEMPLATE = """Generate {count} distinct SVG icon variations for: {prompt}
113 
114Requirements for EACH icon:
115- Output ONLY valid SVG code
116- ViewBox: "0 0 {viewbox} {viewbox}"
117- Use currentColor for strokes/fills (inherits CSS color)
118- No embedded fonts, raster images, or external references
119- Optimized paths with minimal nodes
120- Include <title> element with icon description
121{style_instructions}
122{color_instructions}
123 
124Generate {count} different visual interpretations. Output each SVG in a separate ```svg``` code block.
125Label each variation (e.g., "Variation 1: [brief description]")."""
126 
127 
128def extract_svgs(text):
129    """Extract SVG code blocks from model response"""
130    svgs = []
131 
132    # Try ```svg code blocks first
133    pattern = r'```svg\s*\n(.*?)```'
134    matches = re.findall(pattern, text, re.DOTALL)
135    if matches:
136        svgs.extend(matches)
137 
138    # Fallback: try ```xml code blocks
139    if not svgs:
140        pattern = r'```xml\s*\n(.*?)```'
141        matches = re.findall(pattern, text, re.DOTALL)
142        svgs.extend(matches)
143 
144    # Fallback: try bare <svg> tags
145    if not svgs:
146        pattern = r'(<svg[^>]*>.*?</svg>)'
147        matches = re.findall(pattern, text, re.DOTALL)
148        svgs.extend(matches)
149 
150    # Clean up extracted SVGs
151    cleaned = []
152    for svg in svgs:
153        svg = svg.strip()
154        if not svg.startswith('<svg'):
155            # Try to find <svg> within the extracted text
156            match = re.search(r'(<svg[^>]*>.*?</svg>)', svg, re.DOTALL)
157            if match:
158                svg = match.group(1)
159            else:
160                continue
161        cleaned.append(svg)
162 
163    return cleaned
164 
165 
166def apply_color(svg_code, color):
167    """Replace currentColor with specific color if provided"""
168    if color:
169        # Replace currentColor with the specified color
170        svg_code = svg_code.replace('currentColor', color)
171        # If no currentColor was present, add fill/stroke color
172        if color not in svg_code:
173            svg_code = svg_code.replace('<svg', f'<svg color="{color}"', 1)
174    return svg_code
175 
176 
177def apply_viewbox_size(svg_code, size):
178    """Adjust SVG viewBox to target size"""
179    if size:
180        # Update width/height attributes if present
181        svg_code = re.sub(r'width="[^"]*"', f'width="{size}"', svg_code)
182        svg_code = re.sub(r'height="[^"]*"', f'height="{size}"', svg_code)
183        # Add width/height if not present
184        if 'width=' not in svg_code:
185            svg_code = svg_code.replace('<svg', f'<svg width="{size}" height="{size}"', 1)
186    return svg_code
187 
188 
189def generate_icon(prompt, style=None, category=None, name=None,
190                  color=None, size=24, output_path=None, viewbox=24):
191    """Generate a single SVG icon using Gemini 3.1 Pro Preview"""
192 
193    if not GEMINI_API_KEY:
194        print("Error: GEMINI_API_KEY not set")
195        print("Set it with: export GEMINI_API_KEY='your-key'")
196        return None
197 
198    client = genai.Client(api_key=GEMINI_API_KEY)
199 
200    # Build style instructions
201    style_instructions = ""
202    if style and style in ICON_STYLES:
203        style_instructions = f"- Style: {ICON_STYLES[style]}"
204 
205    # Build color instructions
206    color_instructions = "- Use currentColor for all strokes and fills"
207    if color:
208        color_instructions = f"- Use color: {color} for primary elements, currentColor for secondary"
209 
210    # Build size instructions
211    size_instructions = f"- Design for {size}px display size, optimize detail level accordingly"
212 
213    # Build final prompt
214    icon_prompt = prompt
215    if category and category in ICON_CATEGORIES:
216        icon_prompt = f"{prompt} (category: {ICON_CATEGORIES[category]})"
217    if name:
218        icon_prompt = f"'{name}' icon: {icon_prompt}"
219 
220    full_prompt = SVG_PROMPT_TEMPLATE.format(
221        prompt=icon_prompt,
222        viewbox=viewbox,
223        style_instructions=style_instructions,
224        color_instructions=color_instructions,
225        size_instructions=size_instructions
226    )
227 
228    print(f"Generating icon with {MODEL}...")
229    print(f"Prompt: {prompt}")
230    if style:
231        print(f"Style: {style}")
232    print()
233 
234    try:
235        response = client.models.generate_content(
236            model=MODEL,
237            contents=full_prompt,
238            config=types.GenerateContentConfig(
239                temperature=0.7,
240                max_output_tokens=4096,
241            )
242        )
243 
244        # Extract SVG from response
245        response_text = response.text if hasattr(response, 'text') else ""
246        if not response_text:
247            for part in response.candidates[0].content.parts:
248                if hasattr(part, 'text') and part.text:
249                    response_text += part.text
250 
251        svgs = extract_svgs(response_text)
252 
253        if not svgs:
254            print("No valid SVG generated. Model response:")
255            print(response_text[:500])
256            return None
257 
258        svg_code = svgs[0]
259 
260        # Apply color if specified
261        svg_code = apply_color(svg_code, color)
262 
263        # Apply size
264        svg_code = apply_viewbox_size(svg_code, size)
265 
266        # Determine output path
267        if output_path is None:
268            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
269            slug = name or prompt.split()[0] if prompt else "icon"
270            slug = re.sub(r'[^a-zA-Z0-9_-]', '_', slug.lower())
271            style_suffix = f"_{style}" if style else ""
272            output_path = f"{slug}{style_suffix}_{timestamp}.svg"
273 
274        # Save SVG
275        with open(output_path, "w", encoding="utf-8") as f:
276            f.write(svg_code)
277 
278        print(f"Icon saved to: {output_path}")
279        return output_path
280 
281    except Exception as e:
282        print(f"Error generating icon: {e}")
283        return None
284 
285 
286def generate_batch(prompt, count, output_dir, style=None, color=None,
287                   viewbox=24, name=None):
288    """Generate multiple icon variations"""
289 
290    if not GEMINI_API_KEY:
291        print("Error: GEMINI_API_KEY not set")
292        return []
293 
294    client = genai.Client(api_key=GEMINI_API_KEY)
295    os.makedirs(output_dir, exist_ok=True)
296 
297    # Build instructions
298    style_instructions = ""
299    if style and style in ICON_STYLES:
300        style_instructions = f"- Style: {ICON_STYLES[style]}"
301 
302    color_instructions = "- Use currentColor for all strokes and fills"
303    if color:
304        color_instructions = f"- Use color: {color} for primary elements"
305 
306    full_prompt = SVG_BATCH_PROMPT_TEMPLATE.format(
307        prompt=prompt,
308        count=count,
309        viewbox=viewbox,
310        style_instructions=style_instructions,
311        color_instructions=color_instructions
312    )
313 
314    print(f"\n{'='*60}")
315    print(f"  BATCH ICON GENERATION")
316    print(f"  Model: {MODEL}")
317    print(f"  Prompt: {prompt}")
318    print(f"  Variants: {count}")
319    print(f"  Output: {output_dir}")
320    print(f"{'='*60}\n")
321 
322    try:
323        response = client.models.generate_content(
324            model=MODEL,
325            contents=full_prompt,
326            config=types.GenerateContentConfig(
327                temperature=0.9,
328                max_output_tokens=16384,
329            )
330        )
331 
332        response_text = response.text if hasattr(response, 'text') else ""
333        if not response_text:
334            for part in response.candidates[0].content.parts:
335                if hasattr(part, 'text') and part.text:
336                    response_text += part.text
337 
338        svgs = extract_svgs(response_text)
339 
340        if not svgs:
341            print("No valid SVGs generated.")
342            print(response_text[:500])
343            return []
344 
345        results = []
346        slug = name or re.sub(r'[^a-zA-Z0-9_-]', '_', prompt.split()[0].lower())
347        style_suffix = f"_{style}" if style else ""
348 
349        for i, svg_code in enumerate(svgs[:count]):
350            svg_code = apply_color(svg_code, color)
351            filename = f"{slug}{style_suffix}_{i+1:02d}.svg"
352            filepath = os.path.join(output_dir, filename)
353 
354            with open(filepath, "w", encoding="utf-8") as f:
355                f.write(svg_code)
356 
357            results.append(filepath)
358            print(f"  [{i+1}/{len(svgs[:count])}] Saved: {filename}")
359 
360        print(f"\n{'='*60}")
361        print(f"  BATCH COMPLETE: {len(results)}/{count} icons generated")
362        print(f"{'='*60}\n")
363 
364        return results
365 
366    except Exception as e:
367        print(f"Error generating icons: {e}")
368        return []
369 
370 
371def generate_sizes(prompt, sizes, style=None, color=None, output_dir=None, name=None):
372    """Generate same icon at multiple sizes"""
373    if output_dir is None:
374        output_dir = "."
375    os.makedirs(output_dir, exist_ok=True)
376 
377    results = []
378    slug = name or re.sub(r'[^a-zA-Z0-9_-]', '_', prompt.split()[0].lower())
379    style_suffix = f"_{style}" if style else ""
380 
381    for size in sizes:
382        print(f"Generating {size}px variant...")
383        filename = f"{slug}{style_suffix}_{size}px.svg"
384        filepath = os.path.join(output_dir, filename)
385 
386        result = generate_icon(
387            prompt=prompt,
388            style=style,
389            color=color,
390            size=size,
391            output_path=filepath,
392            viewbox=size
393        )
394 
395        if result:
396            results.append(result)
397 
398        time.sleep(1)
399 
400    return results
401 
402 
403def main():
404    parser = argparse.ArgumentParser(
405        description="Generate SVG icons using Gemini 3.1 Pro Preview"
406    )
407    parser.add_argument("--prompt", "-p", type=str, help="Icon description")
408    parser.add_argument("--name", "-n", type=str, help="Icon name (for filename)")
409    parser.add_argument("--style", "-s", choices=list(ICON_STYLES.keys()),
410                        help="Icon style")
411    parser.add_argument("--category", "-c", choices=list(ICON_CATEGORIES.keys()),
412                        help="Icon category for context")
413    parser.add_argument("--color", type=str,
414                        help="Primary color (hex, e.g. #6366F1). Default: currentColor")
415    parser.add_argument("--size", type=int, default=24,
416                        help="Icon size in px (default: 24)")
417    parser.add_argument("--viewbox", type=int, default=24,
418                        help="SVG viewBox size (default: 24)")
419    parser.add_argument("--output", "-o", type=str, help="Output file path")
420    parser.add_argument("--output-dir", type=str, help="Output directory for batch")
421    parser.add_argument("--batch", type=int,
422                        help="Number of icon variants to generate")
423    parser.add_argument("--sizes", type=str,
424                        help="Comma-separated sizes (e.g. '16,24,32,48')")
425    parser.add_argument("--list-styles", action="store_true",
426                        help="List available icon styles")
427    parser.add_argument("--list-categories", action="store_true",
428                        help="List available icon categories")
429 
430    args = parser.parse_args()
431 
432    if args.list_styles:
433        print("Available icon styles:")
434        for style, desc in ICON_STYLES.items():
435            print(f"  {style}: {desc[:70]}...")
436        return
437 
438    if args.list_categories:
439        print("Available icon categories:")
440        for cat, desc in ICON_CATEGORIES.items():
441            print(f"  {cat}: {desc}")
442        return
443 
444    if not args.prompt and not args.name:
445        parser.error("Either --prompt or --name is required")
446 
447    prompt = args.prompt or args.name
448 
449    # Multi-size mode
450    if args.sizes:
451        sizes = [int(s.strip()) for s in args.sizes.split(",")]
452        generate_sizes(
453            prompt=prompt,
454            sizes=sizes,
455            style=args.style,
456            color=args.color,
457            output_dir=args.output_dir or "./icons",
458            name=args.name
459        )
460    # Batch mode
461    elif args.batch:
462        output_dir = args.output_dir or "./icons"
463        generate_batch(
464            prompt=prompt,
465            count=args.batch,
466            output_dir=output_dir,
467            style=args.style,
468            color=args.color,
469            viewbox=args.viewbox,
470            name=args.name
471        )
472    # Single icon
473    else:
474        generate_icon(
475            prompt=prompt,
476            style=args.style,
477            category=args.category,
478            name=args.name,
479            color=args.color,
480            size=args.size,
481            output_path=args.output,
482            viewbox=args.viewbox
483        )
484 
485 
486if __name__ == "__main__":
487    main()
488