Infographics

scripts/generate_infographic_ai.py

#!/usr/bin/env python3
"""
AI-powered infographic generation using Nano Banana 2.

This script uses a smart iterative refinement approach:
1. (Optional) Research phase - gather facts and data using Perplexity Sonar
2. Generate initial infographic with Nano Banana 2
3. AI quality review using Gemini 3.7 Flash for infographic critique
4. Only regenerate if quality is below threshold for document type
5. Repeat until quality meets standards (max iterations)

Requirements:
    - OPENROUTER_API_KEY environment variable
    - requests library

Usage:
    python generate_infographic_ai.py "5 benefits of exercise" -o benefits.png --type list
    python generate_infographic_ai.py "Global AI market size" -o ai_market.png --type statistical --research
    python generate_infographic_ai.py "Company history 2010-2025" -o timeline.png --type timeline --style corporate
"""

import argparse
import base64
import json
import os
import re
import sys
import time
from pathlib import Path
from datetime import date
from typing import Optional, Dict, Any, List, NamedTuple

try:
    import requests
except ImportError:
    print("Error: requests library not found. Install with: uv pip install requests")
    sys.exit(1)


class ReviewResult(NamedTuple):
    """The outcome of one quality review.

    A named tuple rather than a bare tuple because the review has several
    failure modes -- the model returns no choices, the request raises, the
    answer arrives in an unexpected shape -- and each of them used to be a
    chance to return the wrong number of values to the caller.

    `score` is None whenever no usable number was parsed, and `reviewed` is
    False whenever the review never produced an answer at all. Neither case
    invents a score: an infographic whose quality was not measured is reported
    as unmeasured, not as passing.
    """

    critique: str
    score: Optional[float]
    needs_improvement: bool
    reviewed: bool
    error: Optional[str] = None


# Markdown decoration the reviewer routinely wraps its answer in. Gemini writes
# "**SCORE:** 8.5" far more often than the bare "SCORE: 8.5" the rubric asks
# for, and a pattern that does not allow for it silently scores every image
# at whatever the default happens to be.
_MARKUP = re.compile(r"[*_`#]")

_SCORE_PATTERN = re.compile(r"\bSCORE\s*:?\s*(\d+(?:\.\d+)?)", re.IGNORECASE)
_LOOSE_SCORE_PATTERN = re.compile(
    r"(?:score|rating|quality)\s*[:\s]\s*(\d+(?:\.\d+)?)\s*(?:/\s*10)?", re.IGNORECASE
)
# The underscore is optional because normalization strips it out of
# NEEDS_IMPROVEMENT along with the surrounding markdown.
_VERDICT_PATTERN = re.compile(
    r"^\s*VERDICT\s*:?\s*(ACCEPTABLE|NEEDS[_ ]?IMPROVEMENT)", re.IGNORECASE | re.MULTILINE
)


def _normalize_review_text(text: str) -> str:
    """Strip markdown decoration so the rubric's fields can be found."""
    return _MARKUP.sub("", text or "")


def _parse_score(text: str) -> Optional[float]:
    """Return the reviewer's 0-10 score, or None when there isn't a usable one.

    None rather than a default: a fabricated score is indistinguishable from a
    real one once it reaches the review log, and a default near the middle of
    the range quietly passes lenient document types while failing strict ones.
    """
    normalized = _normalize_review_text(text)
    match = _SCORE_PATTERN.search(normalized) or _LOOSE_SCORE_PATTERN.search(normalized)
    if not match:
        return None
    score = float(match.group(1))
    # A number outside the rubric's range means something other than the score
    # was matched -- a figure count, a percentage, an iteration number.
    return score if 0.0 <= score <= 10.0 else None


def _parse_verdict(text: str) -> Optional[str]:
    """Return ACCEPTABLE / NEEDS_IMPROVEMENT from the verdict line, or None.

    Anchored to the start of a line. The rubric prompt spells out both tokens
    when it explains the response format, so a reviewer that echoes those
    instructions back would fail a perfectly good image under a plain
    substring search.
    """
    match = _VERDICT_PATTERN.search(_normalize_review_text(text))
    if not match:
        return None
    verdict = match.group(1).upper().replace(" ", "").replace("_", "")
    return "NEEDS_IMPROVEMENT" if verdict == "NEEDSIMPROVEMENT" else "ACCEPTABLE"


def _resolve_api_key(explicit: Optional[str] = None) -> Optional[str]:
    """Resolve the OpenRouter key from --api-key, the environment, then any .env file.

    The .env scan walks up from the working directory and finally checks the
    script's own directory, so running from anywhere inside a project picks up
    the key at its root. Only the standard library is used: python-dotenv is a
    common omission, and a missing optional dependency should not read as a
    missing credential.
    """
    if explicit:
        return explicit

    from_env = os.environ.get("OPENROUTER_API_KEY", "").strip()
    if from_env:
        return from_env

    cwd = Path.cwd()
    for directory in [cwd, *cwd.parents, Path(__file__).resolve().parent]:
        env_file = directory / ".env"
        if not env_file.is_file():
            continue
        try:
            content = env_file.read_text(encoding="utf-8", errors="replace")
        except OSError:
            continue
        for raw in content.splitlines():
            line = raw.strip()
            if line.startswith("#") or "=" not in line:
                continue
            name, _, value = line.partition("=")
            if name.strip() == "OPENROUTER_API_KEY":
                value = value.strip().strip('"').strip("'")
                if value:
                    return value

    return None


# Infographic type configurations with detailed prompting
INFOGRAPHIC_TYPES = {
    "statistical": {
        "name": "Statistical/Data-Driven",
        "guidelines": """
STATISTICAL INFOGRAPHIC REQUIREMENTS:
- Large, bold numbers that are immediately readable
- Clear data visualization (bar charts, pie charts, or donut charts)
- Data callouts with context (e.g., "15% increase")
- Trend indicators (arrows, growth symbols)
- Legend if multiple data series
- Source attribution area at bottom
- Clean grid or alignment for data elements
"""
    },
    "timeline": {
        "name": "Timeline/Chronological",
        "guidelines": """
TIMELINE INFOGRAPHIC REQUIREMENTS:
- Clear chronological flow (horizontal or vertical)
- Prominent date/year markers
- Connecting line or path between events
- Event nodes (circles, icons, or markers)
- Brief event descriptions
- Consistent spacing between events
- Visual progression indicating time direction
- Start and end points clearly marked
"""
    },
    "process": {
        "name": "Process/How-To",
        "guidelines": """
PROCESS INFOGRAPHIC REQUIREMENTS:
- Numbered steps (1, 2, 3, etc.) clearly visible
- Directional arrows between steps
- Action-oriented icons for each step
- Brief action text for each step
- Clear start and end indicators
- Logical flow direction (top-to-bottom or left-to-right)
- Visual connection between sequential steps
"""
    },
    "comparison": {
        "name": "Comparison",
        "guidelines": """
COMPARISON INFOGRAPHIC REQUIREMENTS:
- Symmetrical side-by-side layout
- Clear headers for each option being compared
- Matching rows/categories for fair comparison
- Visual indicators (checkmarks, X marks, ratings)
- Balanced visual weight on both sides
- Color differentiation between options
- Summary or verdict section if applicable
"""
    },
    "list": {
        "name": "List/Informational",
        "guidelines": """
LIST INFOGRAPHIC REQUIREMENTS:
- Large, clear numbers or bullet points
- Icons representing each list item
- Brief, scannable text descriptions
- Consistent visual treatment for all items
- Clear hierarchy (title, items, details)
- Adequate spacing between items
- Visual variety to prevent monotony
"""
    },
    "geographic": {
        "name": "Geographic/Map-Based",
        "guidelines": """
GEOGRAPHIC INFOGRAPHIC REQUIREMENTS:
- Accurate map representation
- Color-coded regions based on data
- Clear legend explaining color scale
- Data callouts for key regions
- Region labels where space permits
- Scale or context indicator
- Clean cartographic style
"""
    },
    "hierarchical": {
        "name": "Hierarchical/Pyramid",
        "guidelines": """
HIERARCHICAL INFOGRAPHIC REQUIREMENTS:
- Clear pyramid or tree structure
- Distinct levels with visual separation
- Size progression (larger at base, smaller at top)
- Labels for each tier
- Color gradient or distinct colors per level
- Clear top-to-bottom or bottom-to-top hierarchy
- Balanced, centered composition
"""
    },
    "anatomical": {
        "name": "Anatomical/Visual Metaphor",
        "guidelines": """
ANATOMICAL INFOGRAPHIC REQUIREMENTS:
- Central metaphor image (body, tree, machine, etc.)
- Labeled components with callout lines
- Clear connection between labels and parts
- Explanatory text for each labeled part
- Consistent callout style throughout
- Educational, diagram-like appearance
"""
    },
    "resume": {
        "name": "Resume/Professional",
        "guidelines": """
RESUME INFOGRAPHIC REQUIREMENTS:
- Professional photo/avatar placeholder area
- Skill visualization (bars, charts, ratings)
- Experience timeline or list
- Contact information section with icons
- Achievement or certification badges
- Clean, professional layout
- Personal branding colors
"""
    },
    "social": {
        "name": "Social Media",
        "guidelines": """
SOCIAL MEDIA INFOGRAPHIC REQUIREMENTS:
- Bold, attention-grabbing headline
- Large, impactful central statistic or visual
- Minimal text, maximum visual impact
- Platform-appropriate format (square for Instagram)
- Vibrant, eye-catching colors
- Call-to-action element
- Brand/logo placement area
"""
    },
}

# Industry style configurations
STYLE_PRESETS = {
    "corporate": {
        "name": "Corporate/Business",
        "colors": "navy blue (#1E3A5F), steel blue (#4A90A4), gold (#F5A623) accents",
        "description": "Clean, professional, minimal design with structured layout",
    },
    "healthcare": {
        "name": "Healthcare/Medical",
        "colors": "medical blue (#0077B6), cyan (#00B4D8), light cyan (#90E0EF)",
        "description": "Trust-inducing, clinical, clean design",
    },
    "technology": {
        "name": "Technology/Data",
        "colors": "tech blue (#2563EB), slate gray (#475569), violet (#7C3AED) accents",
        "description": "Modern, innovative, futuristic design",
    },
    "nature": {
        "name": "Nature/Environmental",
        "colors": "forest green (#2D6A4F), mint (#95D5B2), earth brown (#8B4513)",
        "description": "Organic, natural, earth-toned design",
    },
    "education": {
        "name": "Education/Academic",
        "colors": "academic blue (#3D5A80), light blue (#98C1D9), coral (#EE6C4D) accents",
        "description": "Friendly, approachable, educational design",
    },
    "marketing": {
        "name": "Marketing/Creative",
        "colors": "coral (#FF6B6B), teal (#4ECDC4), yellow (#FFE66D)",
        "description": "Bold, vibrant, eye-catching design",
    },
    "finance": {
        "name": "Finance/Investment",
        "colors": "navy (#14213D), gold (#FCA311), green (#2ECC71) for positive",
        "description": "Conservative, trustworthy, professional design",
    },
    "nonprofit": {
        "name": "Nonprofit/Cause",
        "colors": "warm orange (#E07A5F), sage green (#81B29A), sand (#F2CC8F)",
        "description": "Warm, human-centered, impactful design",
    },
}

# Colorblind-safe palette options
PALETTE_PRESETS = {
    "wong": {
        "name": "Wong's Palette",
        "colors": "orange (#E69F00), sky blue (#56B4E9), bluish green (#009E73), yellow (#F0E442), blue (#0072B2), vermillion (#D55E00), reddish purple (#CC79A7), black (#000000)",
    },
    "ibm": {
        "name": "IBM Colorblind-Safe",
        "colors": "ultramarine (#648FFF), indigo (#785EF0), magenta (#DC267F), orange (#FE6100), gold (#FFB000)",
    },
    "tol": {
        "name": "Tol's Muted Qualitative",
        "colors": "rose (#CC6677), indigo (#332288), sand (#DDCC77), green (#117733), cyan (#88CCEE), wine (#882255), teal (#44AA99), olive (#999933), purple (#AA4499); pale gray (#DDDDDD) only for missing data",
    },
}


class InfographicGenerator:
    """Generate infographics using AI with smart iterative refinement.
    
    Uses Gemini 3.7 Flash for quality review to determine if regeneration is needed.
    Multiple passes only occur if the generated infographic doesn't meet the
    quality threshold for the target document type.
    """
    
    # Quality thresholds by document type (score out of 10)
    QUALITY_THRESHOLDS = {
        "marketing": 8.5,     # Marketing materials - must be compelling
        "report": 8.0,        # Business reports - professional quality
        "presentation": 7.5,  # Slides/talks - clear and engaging
        "social": 7.0,        # Social media - eye-catching
        "internal": 7.0,      # Internal use - good quality
        "draft": 6.5,         # Draft/working - acceptable
        "default": 7.5,       # Default threshold
    }
    
    # Base infographic design guidelines
    INFOGRAPHIC_GUIDELINES = """
Create a high-quality professional infographic with these requirements:

VISUAL QUALITY:
- Clean white or light solid color background (no busy textures)
- High contrast for readability
- Professional, polished appearance
- Sharp, clear graphics and text
- Adequate spacing to prevent crowding
- Balanced composition

TYPOGRAPHY:
- Bold, attention-grabbing headline
- Clear, readable sans-serif fonts
- Minimum 12pt font size for body text
- Larger text for key statistics and headlines
- Consistent font hierarchy throughout
- No overlapping text

LAYOUT:
- Clear visual hierarchy (most important info largest/first)
- Logical reading flow (top-to-bottom or left-to-right)
- 60% visual elements, 40% text (approximately)
- Adequate white space between sections
- Balanced left-right composition
- Clear section divisions

DATA VISUALIZATION:
- Large, bold numbers for key statistics
- Clear, accurate charts and graphs
- Labeled axes and data points
- Legend where needed
- Icons that clearly represent concepts

ACCESSIBILITY:
- Colorblind-friendly color choices
- High contrast between text and background
- Shapes and labels, not just colors, to convey meaning
- Works in grayscale

IMPORTANT - NO META CONTENT:
- Do NOT include the prompt, instructions, or metadata in the image
- Do NOT include layout descriptions like "left panel", "right panel"
- Do NOT include font or color specifications
- Only include the actual infographic content
"""

    def __init__(self, api_key: Optional[str] = None, verbose: bool = False):
        """Initialize the generator."""
        self.api_key = _resolve_api_key(api_key)

        if not self.api_key:
            raise ValueError(
                "OPENROUTER_API_KEY not found. Please either:\n"
                "  1. Set the OPENROUTER_API_KEY environment variable\n"
                "  2. Add OPENROUTER_API_KEY to your .env file\n"
                "  3. Pass api_key parameter to the constructor\n"
                "Get your API key from: https://openrouter.ai/keys"
            )
        
        self.verbose = verbose
        self._last_error = None
        self.base_url = "https://openrouter.ai/api/v1"
        # Nano Banana 2 for image generation. The slug must be an image-output
        # model; a text-only chat model is rejected with "No endpoints found that
        # support the requested output modalities".
        # https://openrouter.ai/google/gemini-3.1-flash-image
        self.image_model = "google/gemini-3.1-flash-image"
        # Gemini 3.7 Flash for quality review - reads the image, answers in text
        self.review_model = "google/gemini-3.7-flash"
        self.research_model = "perplexity/sonar-pro"
        
    def _log(self, message: str):
        """Log message if verbose mode is enabled."""
        if self.verbose:
            print(f"[{time.strftime('%H:%M:%S')}] {message}")
    
    # ========== RESEARCH METHODS ==========
    
    def research_topic(self, topic: str, infographic_type: Optional[str] = None) -> Dict[str, Any]:
        """
        Research a topic using Perplexity Sonar Pro to gather facts and data.
        
        Args:
            topic: The topic to research
            infographic_type: Type of infographic to tailor the research
            
        Returns:
            Dictionary with research results including facts, statistics, and sources
        """
        self._log(f"Researching topic: {topic}")
        
        # Build research query based on infographic type
        type_context = ""
        if infographic_type:
            if infographic_type == "statistical":
                type_context = "Focus on statistics, numbers, percentages, and quantitative data."
            elif infographic_type == "timeline":
                type_context = "Focus on key dates, milestones, and chronological events."
            elif infographic_type == "process":
                type_context = "Focus on steps, procedures, and sequential information."
            elif infographic_type == "comparison":
                type_context = "Focus on comparing different options, pros/cons, and differences."
            elif infographic_type == "list":
                type_context = "Focus on key points, tips, facts, and organized information."
            elif infographic_type == "geographic":
                type_context = "Focus on regional data, location-based statistics, and geographic distribution."
            elif infographic_type == "hierarchical":
                type_context = "Focus on levels, rankings, and hierarchical relationships."
        
        research_prompt = f"""You are a research assistant gathering information for an infographic.

TOPIC: {topic}

{type_context}

Please provide:
1. KEY FACTS: 5-8 key facts or statistics about this topic (with specific numbers where possible)
2. CONTEXT: Brief background context (2-3 sentences)
3. SOURCES: Provide primary-source URLs, publication dates, and which claims they support
4. DATA POINTS: Any specific data points that would make good visualizations

Format your response as structured data that can be easily incorporated into an infographic.
Be specific with numbers, percentages, and dates.
As of {date.today().isoformat()}, prioritize relevant recent evidence; retain original dates for historical events.
Do not invent missing numbers. Include units, populations, denominators, and source URLs."""

        messages = [
            {
                "role": "system",
                "content": "You are an expert research assistant. Provide accurate, well-sourced information formatted for infographic creation. Include numbers, dates, and statistics only when supported by cited evidence."
            },
            {"role": "user", "content": research_prompt}
        ]
        
        return self._research_request(messages, max_tokens=2000, search_context="high")

    @staticmethod
    def _text_content(message: Dict[str, Any]) -> str:
        """Read returned answer text without treating internal reasoning as an answer."""
        content = message.get("content")
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            return "\n".join(
                block["text"] for block in content
                if isinstance(block, dict) and block.get("type") == "text"
                and isinstance(block.get("text"), str)
            )
        return ""

    def _research_request(self, messages: List[Dict[str, Any]], max_tokens: int,
                          search_context: str = "medium") -> Dict[str, Any]:
        """Preserve normalized citations from OpenRouter's Sonar response."""
        try:
            result = self._post_request("chat/completions", {
                "model": self.research_model,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": 0.1,
                "web_search_options": {"search_context_size": search_context},
            })
            choices = result.get("choices", [])
            if not choices:
                raise RuntimeError("No response from research model")
            message = choices[0].get("message", {})
            content = self._text_content(message)
            if not content.strip():
                raise RuntimeError("Research response contained no answer text")
            annotations = message.get("annotations") or []
            sources = [
                item["url_citation"] for item in annotations
                if isinstance(item, dict) and item.get("type") == "url_citation"
                and isinstance(item.get("url_citation"), dict)
            ]
            return {
                "success": True, "content": content, "sources": sources,
                "annotations": annotations,
                # Preserve provider fields if present; normalized citations above
                # are the documented OpenRouter response contract.
                "search_results": result.get("search_results", []),
                "citations": result.get("citations", []),
                "model": self.research_model,
                "retrieved_on": date.today().isoformat(),
            }
        except Exception as exc:
            detail = str(exc).replace(self.api_key, "[redacted]")
            self._log(f"Research failed: {detail}")
            return {"success": False, "error": detail, "model": self.research_model}

    def web_search(self, query: str) -> Dict[str, Any]:
        """
        Perform a quick web search for current information.
        
        Args:
            query: Search query
            
        Returns:
            Dictionary with search results
        """
        self._log(f"Web search: {query}")
        
        search_prompt = f"""Search for current information about: {query}

Provide:
1. The most relevant and recent facts
2. Any statistics or numbers
3. Key dates if applicable
4. Brief source attribution

Be concise and factual. Focus on information useful for an infographic."""

        messages = [
            {
                "role": "system",
                "content": "You are a web search assistant. Provide accurate, current information with sources."
            },
            {"role": "user", "content": search_prompt}
        ]
        
        return self._research_request(messages, max_tokens=1000)

    def _enhance_prompt_with_research(self, user_prompt: str, research_data: Dict[str, Any]) -> str:
        """
        Enhance the user prompt with researched information.
        
        Args:
            user_prompt: Original user prompt
            research_data: Research results dictionary
            
        Returns:
            Enhanced prompt with research data
        """
        if not research_data.get("success") or not research_data.get("content"):
            return user_prompt
        
        enhanced = f"""{user_prompt}

RESEARCH CANDIDATES (check cited evidence before publication; do not invent additional claims):
{research_data['content']}

SOURCE RECORDS:
{json.dumps(research_data.get('sources', []), ensure_ascii=False)}

Use the above researched facts, statistics, and data points to create an accurate, informative infographic.
Incorporate specific numbers, percentages, and dates from the research."""
        
        return enhanced
    
    # ========== END RESEARCH METHODS ==========
    
    def _post_request(self, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Send one non-streaming request; never automatically replay a paid call."""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "X-OpenRouter-Title": "Infographic Generator",
        }
        self._log(f"POST {endpoint} using {payload['model']}")
        try:
            response = requests.post(
                f"{self.base_url}/{endpoint}", headers=headers,
                json=payload, timeout=120,
            )
            try:
                response_json = response.json()
            except ValueError:
                raise RuntimeError(
                    f"API returned non-JSON data (HTTP {response.status_code})"
                ) from None
            if not isinstance(response_json, dict):
                raise RuntimeError("API response must be a JSON object")
            if response.status_code != 200 or "error" in response_json:
                error = response_json.get("error", {})
                detail = error.get("message", "request rejected") if isinstance(error, dict) else str(error)
                # Providers occasionally echo input; do not dump whole bodies or credentials.
                detail = str(detail).replace(self.api_key, "[redacted]")[:500]
                raise RuntimeError(f"API request failed (HTTP {response.status_code}): {detail}")
            return response_json
        except requests.exceptions.Timeout:
            raise RuntimeError("API request timed out after 120 seconds") from None
        except requests.exceptions.RequestException as exc:
            detail = str(exc).replace(self.api_key, "[redacted]")
            raise RuntimeError(f"API request failed: {detail}") from None

    def _make_request(self, model: str, messages: List[Dict[str, Any]],
                      modalities: Optional[List[str]] = None) -> Dict[str, Any]:
        """Use chat completions for the text-and-image vision review."""
        payload = {"model": model, "messages": messages}
        if modalities:
            payload["modalities"] = modalities
        return self._post_request("chat/completions", payload)

    def _extract_image_from_response(self, response: Dict[str, Any]) -> bytes:
        """Decode the dedicated Image API response and enforce this helper's PNG contract.

        The provider may omit media_type; the PNG signature still has to match.
        This is a format check, not a full image-decoder or scientific-content check.
        """
        data = response.get("data")
        if not isinstance(data, list) or not data or not isinstance(data[0], dict):
            raise RuntimeError("Image API returned no data[0] image")
        item = data[0]
        if item.get("media_type") not in (None, "image/png"):
            raise RuntimeError(f"Expected PNG output; received {item['media_type']}")
        encoded = item.get("b64_json")
        if not isinstance(encoded, str) or not encoded:
            raise RuntimeError("Image API returned no b64_json image")
        try:
            image_data = base64.b64decode(encoded, validate=True)
        except ValueError:
            raise RuntimeError("Image API returned invalid base64 image data") from None
        if not image_data.startswith(b"\x89PNG\r\n\x1a\n"):
            raise RuntimeError("Image API output does not have a PNG signature")
        return image_data

    def _image_to_base64(self, image_path: str) -> str:
        """Convert image file to base64 data URL."""
        ext = Path(image_path).suffix.lower()
        if ext not in {".png", ".jpg", ".jpeg", ".gif", ".webp"}:
            raise ValueError("Reference images must be PNG, JPEG, GIF, or WebP")
        with open(image_path, "rb") as f:
            image_data = f.read()
        mime_type = {
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
            ".gif": "image/gif",
            ".webp": "image/webp"
        }.get(ext, "image/png")
        
        base64_data = base64.b64encode(image_data).decode("utf-8")
        return f"data:{mime_type};base64,{base64_data}"
    
    def _build_generation_prompt(self, user_prompt: str, 
                                  infographic_type: Optional[str] = None,
                                  style: Optional[str] = None,
                                  palette: Optional[str] = None,
                                  background: str = "white",
                                  context_image_count: int = 0) -> str:
        """Build the full generation prompt with all enhancements."""
        parts = [self.INFOGRAPHIC_GUIDELINES]
        
        # Add type-specific guidelines
        if infographic_type and infographic_type in INFOGRAPHIC_TYPES:
            type_config = INFOGRAPHIC_TYPES[infographic_type]
            parts.append(f"\nINFOGRAPHIC TYPE: {type_config['name']}")
            parts.append(type_config['guidelines'])
        
        # Add style preset
        if style and style in STYLE_PRESETS:
            style_config = STYLE_PRESETS[style]
            parts.append(f"\nSTYLE: {style_config['name']}")
            parts.append(f"Colors: {style_config['colors']}")
            parts.append(f"Design: {style_config['description']}")
        
        # Add colorblind-safe palette
        if palette and palette in PALETTE_PRESETS:
            palette_config = PALETTE_PRESETS[palette]
            parts.append(f"\nCOLORBLIND-SAFE PALETTE: {palette_config['name']}")
            parts.append(f"Use these colors: {palette_config['colors']}")
        
        # Add user request
        parts.append(f"\nUSER REQUEST: {user_prompt}")
        
        if context_image_count:
            plural = "image" if context_image_count == 1 else "images"
            parts.append(
                f"\nREFERENCE CONTEXT: {context_image_count} reference {plural} attached. "
                "Use them as visual/content context for the infographic while following the user request."
            )
        
        # Add background
        parts.append(f"\nBackground: {background} background")
        
        # Final instruction
        parts.append("\nGenerate a professional, publication-quality infographic that meets all the guidelines above.")
        
        return "\n".join(parts)
    
    def generate_image(self, prompt: str, context_images: Optional[List[str]] = None) -> Optional[bytes]:
        """Generate a PNG with the dedicated Image API and optional references."""
        self._last_error = None
        context_images = context_images or []
        if len(context_images) > 14:
            raise ValueError("Gemini 3.1 Flash Image accepts at most 14 reference images")
        payload = {"model": self.image_model, "prompt": prompt, "n": 1}
        if context_images:
            payload["input_references"] = [
                {"type": "image_url", "image_url": {"url": self._image_to_base64(path)}}
                for path in context_images
            ]
        try:
            # Current Gemini endpoints do not advertise output_format. Accept
            # their default only after checking the returned PNG MIME/signature.
            response = self._post_request("images", payload)
            image_data = self._extract_image_from_response(response)
            self._log(f"[OK] Generated PNG ({len(image_data)} bytes)")
            return image_data
        except RuntimeError as exc:
            self._last_error = str(exc)
            self._log(f"[FAIL] Generation failed: {self._last_error}")
            return None

    def review_image(self, image_path: str, original_prompt: str,
                    infographic_type: Optional[str],
                    iteration: int, doc_type: str = "default",
                    max_iterations: int = 3) -> ReviewResult:
        """
        Review generated infographic using Gemini 3.7 Flash for quality analysis.
        
        Evaluates the infographic on multiple criteria specific to good
        infographic design and determines if regeneration is needed.
        """
        image_data_url = self._image_to_base64(image_path)
        
        threshold = self.QUALITY_THRESHOLDS.get(doc_type.lower(), 
                                                 self.QUALITY_THRESHOLDS["default"])
        
        type_name = "general"
        if infographic_type and infographic_type in INFOGRAPHIC_TYPES:
            type_name = INFOGRAPHIC_TYPES[infographic_type]["name"]
        
        review_prompt = f"""You are an expert infographic designer reviewing a generated infographic for quality.

ORIGINAL REQUEST: {original_prompt}

INFOGRAPHIC TYPE: {type_name}
QUALITY THRESHOLD: {threshold}/10
ITERATION: {iteration}/{max_iterations}

Carefully evaluate this infographic on these criteria:

1. **Visual Hierarchy & Layout** (0-2 points)
   - Clear visual hierarchy (most important elements prominent)
   - Logical reading flow
   - Balanced composition
   - Adequate white space

2. **Typography & Readability** (0-2 points)
   - Text is readable and appropriately sized
   - Headlines are bold and attention-grabbing
   - No overlapping or cramped text
   - Consistent font usage

3. **Data Visualization** (0-2 points)
   - Numbers and statistics are prominent
   - Charts/icons are clear and accurate
   - Data is easy to understand at a glance
   - Labels are present where needed

4. **Color & Accessibility** (0-2 points)
   - Colors are harmonious and professional
   - Sufficient contrast for readability
   - Works for colorblind viewers
   - Colors support the content hierarchy

5. **Overall Impact & Professionalism** (0-2 points)
   - Looks professional and polished
   - Engaging and visually appealing
   - Free of visual bugs or artifacts
   - Achieves its communication goal

RESPOND IN THIS EXACT FORMAT:
SCORE: [total score 0-10]

STRENGTHS:
- [strength 1]
- [strength 2]

ISSUES:
- [issue 1 if any]
- [issue 2 if any]

SPECIFIC_IMPROVEMENTS:
- [specific improvement 1]
- [specific improvement 2]

VERDICT: [ACCEPTABLE or NEEDS_IMPROVEMENT]

If score >= {threshold}, the infographic is ACCEPTABLE.
If score < {threshold}, mark as NEEDS_IMPROVEMENT with specific suggestions."""

        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": review_prompt
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image_data_url
                        }
                    }
                ]
            }
        ]
        
        try:
            response = self._make_request(
                model=self.review_model,
                messages=messages
            )
            
            choices = response.get("choices", [])
            if not choices:
                # A normal enough outcome -- a rate limit, a content filter, a
                # provider hiccup. The image itself is fine; only its review is
                # missing, and saying so beats inventing a score.
                reason = "the review model returned no choices"
                self._log(f"[WARN] Review unavailable: {reason}")
                return ReviewResult(
                    f"Review unavailable: {reason}.",
                    None, False, reviewed=False, error=reason,
                )

            message = choices[0].get("message", {})
            content = self._text_content(message)

            score = _parse_score(content)
            verdict = _parse_verdict(content)

            if score is None and verdict is None:
                reason = "no score or verdict found in the review"
                self._log(f"[WARN] Review unusable: {reason}")
                return ReviewResult(
                    content if content else f"Review unusable: {reason}.",
                    None, False, reviewed=True, error=reason,
                )

            needs_improvement = verdict == "NEEDS_IMPROVEMENT" or (
                score is not None and score < threshold
            )

            shown = f"{score}/10" if score is not None else "not stated"
            self._log(f"[OK] Review complete (Score: {shown}, Threshold: {threshold}/10)")
            self._log(f"  Verdict: {'Needs improvement' if needs_improvement else 'Acceptable'}")

            return ReviewResult(
                content if content else "Review returned no text",
                score,
                needs_improvement,
                reviewed=True,
                error=None if score is not None else "no score found in the review",
            )
        except Exception as e:
            # A failed review must not fail the run -- the image is already
            # generated and saved -- but it must not read as a pass either.
            self._log(f"[WARN] Review failed: {str(e)}")
            return ReviewResult(
                f"Review failed: {str(e)}",
                None, False, reviewed=False, error=str(e),
            )
    
    def improve_prompt(self, original_prompt: str, critique: str, 
                      infographic_type: Optional[str],
                      style: Optional[str],
                      palette: Optional[str],
                      background: str,
                      iteration: int,
                      context_image_count: int = 0) -> str:
        """Improve the generation prompt based on critique."""
        
        parts = [self.INFOGRAPHIC_GUIDELINES]
        
        # Add type-specific guidelines
        if infographic_type and infographic_type in INFOGRAPHIC_TYPES:
            type_config = INFOGRAPHIC_TYPES[infographic_type]
            parts.append(f"\nINFOGRAPHIC TYPE: {type_config['name']}")
            parts.append(type_config['guidelines'])
        
        # Add style preset
        if style and style in STYLE_PRESETS:
            style_config = STYLE_PRESETS[style]
            parts.append(f"\nSTYLE: {style_config['name']}")
            parts.append(f"Colors: {style_config['colors']}")
            parts.append(f"Design: {style_config['description']}")
        
        # Add palette
        if palette and palette in PALETTE_PRESETS:
            palette_config = PALETTE_PRESETS[palette]
            parts.append(f"\nCOLORBLIND-SAFE PALETTE: {palette_config['name']}")
            parts.append(f"Use these colors: {palette_config['colors']}")
        
        # Add original request
        parts.append(f"\nUSER REQUEST: {original_prompt}")
        
        if context_image_count:
            plural = "image" if context_image_count == 1 else "images"
            parts.append(
                f"\nREFERENCE CONTEXT: {context_image_count} reference {plural} attached. "
                "Continue using them as visual/content context while addressing the review feedback."
            )
        
        parts.append(f"\nBackground: {background} background")
        
        # Add improvement instructions
        parts.append(f"""
ITERATION {iteration}: Based on previous review, address these specific improvements:
{critique}

Generate an improved version that:
1. Fixes ALL the issues mentioned in the review
2. Maintains the requested infographic type and style
3. Ensures professional, publication-ready quality
4. Has no visual bugs, overlapping elements, or readability issues
""")
        
        return "\n".join(parts)
    
    def generate_iterative(self, user_prompt: str, output_path: str,
                          infographic_type: Optional[str] = None,
                          style: Optional[str] = None,
                          palette: Optional[str] = None,
                          background: str = "white",
                          iterations: int = 3,
                          doc_type: str = "default",
                          research: bool = False,
                          context_images: Optional[List[str]] = None) -> Dict[str, Any]:
        """
        Generate infographic with smart iterative refinement.
        
        Only regenerates if the quality score is below the threshold.
        
        Args:
            user_prompt: Description of the infographic content
            output_path: Path to save final image
            infographic_type: Type preset (statistical, timeline, etc.)
            style: Industry style preset
            palette: Colorblind-safe palette
            background: Background color
            iterations: Maximum refinement iterations
            doc_type: Document type for quality threshold
            research: If True, research the topic first for better data
            context_images: Optional image paths to attach as generation context
        """
        if isinstance(iterations, bool) or not isinstance(iterations, int) or iterations < 1:
            raise ValueError("iterations must be a positive integer")
        if doc_type not in self.QUALITY_THRESHOLDS:
            raise ValueError(f"Unknown document type: {doc_type}")
        for name, value, choices in (("type", infographic_type, INFOGRAPHIC_TYPES),
                                     ("style", style, STYLE_PRESETS),
                                     ("palette", palette, PALETTE_PRESETS)):
            if value is not None and value not in choices:
                raise ValueError(f"Unknown {name}: {value}")
        output_path = Path(output_path)
        if output_path.suffix.lower() != ".png":
            raise ValueError("Output path must end in .png")
        if len(context_images or []) > 14:
            raise ValueError("Gemini 3.1 Flash Image accepts at most 14 reference images")
        output_dir = output_path.parent
        output_dir.mkdir(parents=True, exist_ok=True)
        
        context_images = [str(Path(image_path)) for image_path in (context_images or [])]
        for image_path in context_images:
            if not Path(image_path).is_file():
                raise FileNotFoundError(f"Context image not found: {image_path}")
            if Path(image_path).suffix.lower() not in {".png", ".jpg", ".jpeg", ".gif", ".webp"}:
                raise ValueError("Reference images must be PNG, JPEG, GIF, or WebP")
        
        base_name = output_path.stem
        extension = output_path.suffix or ".png"
        
        threshold = self.QUALITY_THRESHOLDS.get(doc_type.lower(), 
                                                 self.QUALITY_THRESHOLDS["default"])
        
        type_name = infographic_type if infographic_type else "general"
        style_name = style if style else "default"
        
        results = {
            "user_prompt": user_prompt,
            "infographic_type": infographic_type,
            "style": style,
            "palette": palette,
            "doc_type": doc_type,
            "quality_threshold": threshold,
            "research_enabled": research,
            "image_model": self.image_model,
            "review_model": self.review_model,
            "research_model": self.research_model if research else None,
            "quality_met": False,
            "termination_reason": None,
            "research_data": None,
            "context_images": context_images,
            "iterations": [],
            "final_image": None,
            # None, not 0.0 -- a run whose review never completed has no score,
            # and a number here would be read as one the reviewer gave.
            "final_score": None,
            "final_reviewed": False,
            "success": False,
            "early_stop": False,
            "early_stop_reason": None
        }
        
        print(f"\n{'='*60}")
        print(f"Generating Infographic with Nano Banana 2")
        print(f"{'='*60}")
        print(f"Content: {user_prompt}")
        print(f"Type: {type_name}")
        print(f"Style: {style_name}")
        print(f"Research: {'Enabled' if research else 'Disabled'}")
        print(f"Context Images: {len(context_images)}")
        print(f"Quality Threshold: {threshold}/10")
        print(f"Max Iterations: {iterations}")
        print(f"Output: {output_path}")
        print(f"{'='*60}\n")
        
        # ===== RESEARCH PHASE =====
        enhanced_prompt = user_prompt
        if research:
            print(f"\n[Research Phase]")
            print("-" * 40)
            print(f"Researching topic for accurate data...")
            
            research_result = self.research_topic(user_prompt, infographic_type)
            results["research_data"] = research_result
            
            if research_result.get("success"):
                print(f"[OK] Research complete - gathered facts and statistics")
                
                # Enhance the prompt with researched data
                enhanced_prompt = self._enhance_prompt_with_research(user_prompt, research_result)
                
                # Save research data to file
                research_path = output_dir / f"{base_name}_research.json"
                with open(research_path, "w", encoding="utf-8") as f:
                    json.dump(research_result, f, indent=2)
                print(f"[OK] Research saved: {research_path}")
            else:
                print(f"[WARN] Research failed: {research_result.get('error', 'Unknown error')}")
                print(f"  Proceeding with original prompt...")
        
        # Build initial prompt (using enhanced prompt if research was done)
        current_prompt = self._build_generation_prompt(
            enhanced_prompt, infographic_type, style, palette, background, len(context_images)
        )
        
        for i in range(1, iterations + 1):
            print(f"\n[Iteration {i}/{iterations}]")
            print("-" * 40)
            
            # Generate image
            print(f"Generating infographic with Nano Banana 2...")
            image_data = self.generate_image(current_prompt, context_images)
            
            if not image_data:
                error_msg = getattr(self, '_last_error', 'Generation failed')
                print(f"[FAIL] Generation failed: {error_msg}")
                results["iterations"].append({
                    "iteration": i,
                    "success": False,
                    "error": error_msg
                })
                results["termination_reason"] = "generation_failed"
                break
            
            # Save iteration image
            iter_path = output_dir / f"{base_name}_v{i}{extension}"
            with open(iter_path, "wb") as f:
                f.write(image_data)
            print(f"[OK] Saved: {iter_path}")
            
            # Review the image with the vision review model
            print(f"Reviewing with {self.review_model}...")
            review = self.review_image(
                str(iter_path), enhanced_prompt, infographic_type, i, doc_type, iterations
            )
            if review.score is not None:
                print(f"[OK] Score: {review.score}/10 (threshold: {threshold}/10)")
            else:
                print(f"[WARN] Review unavailable - image kept, quality not verified")
                print(f"  Reason: {review.error}")

            # Save iteration results
            iteration_result = {
                "iteration": i,
                "image_path": str(iter_path),
                "prompt_length": len(current_prompt),
                "critique": review.critique,
                "score": review.score,
                "reviewed": review.reviewed,
                "review_error": review.error,
                "needs_improvement": review.needs_improvement,
                "success": True
            }
            results["iterations"].append(iteration_result)
            # Retain each saved draft if a later generation fails.
            results["final_image"] = str(iter_path)
            results["final_score"] = review.score
            results["final_reviewed"] = review.reviewed and review.score is not None
            results["success"] = True
            results["quality_met"] = (
                results["final_reviewed"] and review.score >= threshold
                and not review.needs_improvement
            )

            # Check if quality is acceptable
            if not review.needs_improvement:
                if results["quality_met"]:
                    print(f"\n[OK] Quality meets threshold ({review.score} >= {threshold})")
                    print(f"  No further iterations needed!")
                    reason = f"Quality score {review.score} meets threshold {threshold}"
                else:
                    # Regenerating cannot fix a reviewer that did not answer.
                    print(f"\n[WARN] Stopping without a verified score - review the image yourself")
                    reason = f"Review did not produce a score: {review.error}"
                results["final_image"] = str(iter_path)
                results["final_score"] = review.score
                results["final_reviewed"] = review.reviewed and review.score is not None
                results["success"] = True
                results["early_stop"] = True
                results["early_stop_reason"] = reason
                results["termination_reason"] = "quality_met" if results["quality_met"] else "review_unavailable"
                break

            # If this is the last iteration, keep the draft without claiming a pass.
            if i == iterations:
                results["termination_reason"] = "max_iterations"
                print(f"\n[WARN] Maximum iterations reached")
                results["final_image"] = str(iter_path)
                results["final_score"] = review.score
                results["final_reviewed"] = review.reviewed and review.score is not None
                results["success"] = True
                break

            # Quality below threshold - improve prompt
            print("\n[WARN] Reviewer requested improvements")
            print(f"Improving prompt based on feedback...")
            current_prompt = self.improve_prompt(
                enhanced_prompt, review.critique, infographic_type, style, palette, background, i + 1,
                len(context_images)
            )
        
        # Copy final version to output path
        if results["success"] and results["final_image"]:
            final_iter_path = Path(results["final_image"])
            if final_iter_path != output_path:
                import shutil
                shutil.copy(final_iter_path, output_path)
                print(f"\n[OK] Final image: {output_path}")
        
        # Save review log
        log_path = output_dir / f"{base_name}_review_log.json"
        with open(log_path, "w", encoding="utf-8") as f:
            json.dump(results, f, indent=2)
        print(f"[OK] Review log: {log_path}")
        
        print(f"\n{'='*60}")
        print(f"Generation Complete!")
        if results["final_score"] is not None:
            print(f"Final Score: {results['final_score']}/10")
        else:
            print(f"Final Score: unavailable (the review did not produce one)")
        if results["early_stop"]:
            iterations_used = len([r for r in results['iterations'] if r.get('success')])
            print(f"Iterations Used: {iterations_used}/{iterations} (early stop)")
        print(f"{'='*60}\n")
        
        return results


def main():
    """Command-line interface."""
    parser = argparse.ArgumentParser(
        description="Generate infographics using Nano Banana 2 with smart iterative refinement",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate a list infographic
  python generate_infographic_ai.py "5 benefits of meditation" -o benefits.png --type list
  
  # Generate a timeline with corporate style
  python generate_infographic_ai.py "Company history 2010-2025" -o timeline.png --type timeline --style corporate
  
  # Generate with colorblind-safe palette
  python generate_infographic_ai.py "Heart disease stats" -o stats.png --type statistical --palette wong
  
  # Generate with RESEARCH for accurate data (uses Perplexity Sonar)
  python generate_infographic_ai.py "Global AI market 2025" -o ai_market.png --type statistical --research
  
  # Generate using reference images as visual/content context
  python generate_infographic_ai.py "Create a branded hiring infographic" -o hiring.png --context-image logo.png --context-image brand_chart.png
  
  # Verbose output
  python generate_infographic_ai.py "Process diagram" -o process.png --type process -v

Infographic Types:
  statistical   - Data-driven with charts and numbers
  timeline      - Chronological events
  process       - Step-by-step instructions
  comparison    - Side-by-side comparison
  list          - Tips, facts, key points
  geographic    - Map-based visualization
  hierarchical  - Pyramid or tree structure
  anatomical    - Visual metaphor
  resume        - Professional/CV style
  social        - Social media optimized

Style Presets:
  corporate, healthcare, technology, nature, education, marketing, finance, nonprofit

Colorblind-Safe Palettes:
  wong, ibm, tol

Document Types (quality thresholds):
  marketing     8.5/10  - Marketing materials
  report        8.0/10  - Business reports
  presentation  7.5/10  - Slides/talks
  social        7.0/10  - Social media
  internal      7.0/10  - Internal use
  draft         6.5/10  - Working drafts
  default       7.5/10  - General purpose

Environment:
  OPENROUTER_API_KEY    OpenRouter API key (required)
        """
    )
    
    parser.add_argument("prompt", help="Description of the infographic content")
    parser.add_argument("-o", "--output", required=True, 
                       help="Output image path (e.g., infographic.png)")
    parser.add_argument("--type", "-t", choices=list(INFOGRAPHIC_TYPES.keys()),
                       help="Infographic type preset")
    parser.add_argument("--style", "-s", choices=list(STYLE_PRESETS.keys()),
                       help="Industry style preset")
    parser.add_argument("--palette", "-p", choices=list(PALETTE_PRESETS.keys()),
                       help="Colorblind-safe palette")
    parser.add_argument("--background", "-b", default="white",
                       help="Background color (default: white)")
    parser.add_argument("--iterations", type=int, default=3,
                       help="Maximum refinement iterations (default: 3)")
    parser.add_argument("--doc-type", default="default",
                       choices=["marketing", "report", "presentation", "social", 
                               "internal", "draft", "default"],
                       help="Document type for quality threshold (default: default)")
    parser.add_argument("--api-key", help="OpenRouter API key (or set OPENROUTER_API_KEY)")
    parser.add_argument("-v", "--verbose", action="store_true",
                       help="Verbose output")
    parser.add_argument("--research", "-r", action="store_true",
                       help="Research the topic first using Perplexity Sonar for accurate data")
    parser.add_argument("--context-image", action="append", default=[],
                       help="Reference image path to include as Nano Banana 2 context; repeat for multiple images")
    
    args = parser.parse_args()
    
    if args.iterations < 1:
        parser.error("--iterations must be at least 1")
    if Path(args.output).suffix.lower() != ".png":
        parser.error("--output must end in .png")

    # Check for API key - resolves --api-key, the environment, then any .env file
    api_key = _resolve_api_key(args.api_key)
    if not api_key:
        print("Error: OPENROUTER_API_KEY not found")
        print("\nSet it with:")
        print("  export OPENROUTER_API_KEY='your_api_key'")
        print("\nOr add OPENROUTER_API_KEY=your_api_key to a .env file")
        print("Or provide via --api-key flag")
        sys.exit(1)
    
    try:
        generator = InfographicGenerator(api_key=api_key, verbose=args.verbose)
        results = generator.generate_iterative(
            user_prompt=args.prompt,
            output_path=args.output,
            infographic_type=args.type,
            style=args.style,
            palette=args.palette,
            background=args.background,
            iterations=args.iterations,
            doc_type=args.doc_type,
            research=args.research,
            context_images=args.context_image
        )
        
        if results["success"]:
            print(f"\n[OK] Success! Infographic saved to: {args.output}")
            if results.get("quality_met"):
                iterations_used = len([r for r in results['iterations'] if r.get('success')])
                print(f"  (Completed in {iterations_used} iteration(s) - quality threshold met)")
            else:
                print(f"  [WARN] Draft retained; quality not verified as passing ({results['termination_reason']})")
            sys.exit(0)
        else:
            print(f"\n[FAIL] Generation failed. Check review log for details.")
            sys.exit(1)
    except Exception as e:
        print(f"\n[FAIL] Error: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()