src/nodes/slide_creation/image_create.py

from openai import OpenAI
import base64
import re
from dotenv import load_dotenv
from PIL import Image
import os, json
from pydantic import BaseModel, Field
from typing import List
import logging

logger = logging.getLogger(__name__)

load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=api_key) if api_key else None



class ImagePrompt(BaseModel):
    """Defines the structured output for an educational image prompt."""
    visual_type: str
    primary_subject: str
    key_elements: List[str]
    style_description: str
    final_prompt: str

def generate_image_infographic_prompt(presentation_text: str) -> str:
    """
    Generates a highly adaptive, context-aware image prompt from presentation text.
    Intelligently determines the best visual approach based on content type.
    """
    if not client:
        print("Warning: No OpenAI client initialized. Image prompt generation skipped.")
        return None
    
    system_prompt = """
You are an elite educational visual designer with expertise in cognitive learning theory and visual communication. You create infographics that maximize retention, comprehension, and engagement across all domains: AI/ML, Data Science, Software Engineering, Cloud Computing, Algorithms, Mathematics, Physics, Biology, Business, Finance, and more.

Your mission: Analyze the content deeply, understand the learning objective, and create the MOST EFFECTIVE visual representation that matches how humans learn this type of information.

**STEP 1: IDENTIFY LEARNING OBJECTIVE**
First determine what the learner needs to understand:
- Structure/Organization? → Show hierarchy and relationships
- Sequence/Process? → Show order and flow
- Comparison? → Show differences and similarities
- Application? → Show real-world context
- Mechanism? → Show how something works
- Relationships? → Show connections and dependencies

**STEP 2: CONTENT ANALYSIS - Choose the RIGHT visual approach:**

1. **ARCHITECTURE/SYSTEM DESIGN** (Neural Networks, Cloud Systems, Databases, Pipelines):
   → Use: Layered architecture diagrams, component blocks, data flow arrows, system interconnections
   → Include: Clear component labels, connection lines, data flow indicators, layer representations

2. **PROCESS/WORKFLOW** (Algorithms, Pipelines, Lifecycle, Methodology):
   → Use: Step-by-step flowcharts, numbered sequence blocks, decision trees, timeline flows
   → Include: Directional arrows ONLY where process flows, numbered steps, decision points

3. **CONCEPT/THEORY** (Mathematical concepts, Principles, Theories, Definitions):
   → Use: Central concept with surrounding elements, relationship diagrams, hierarchical structures
   → Include: Key terms in circles/boxes, relationship connectors, visual metaphors, formulas visualized

4. **COMPARISON/CONTRAST** (Algorithm vs Algorithm, Tech A vs Tech B, Before/After):
   → Use: Side-by-side layouts, vs. diagrams, tables, feature matrices
   → Include: Clear dividing line, comparative icons, pros/cons indicators

5. **REAL-WORLD EXAMPLES** (Use cases, Applications, Case studies):
   → Use: Scenario illustrations, realistic icons representing actual objects/systems, contextual scenes
   → Include: Industry-specific icons (e.g., healthcare, finance, retail), user interactions, practical scenarios

6. **DATA/STATISTICS** (Performance metrics, Research findings, Benchmarks):
   → Use: Charts (bar, line, pie), graphs, data visualizations, statistical representations
   → Include: Axis labels, legends, data points, trend indicators

7. **MATHEMATICAL/FORMULA-HEAVY** (Equations, Calculations, Mathematical proofs):
   → Use: Formula breakdowns, visual equation components, geometric representations
   → Include: Mathematical symbols clearly shown, step-by-step derivations, geometric shapes

8. **CODE/PROGRAMMING** (Syntax, Code examples, Programming concepts):
   → Use: Simplified code blocks with annotations, syntax highlighting representation, code flow diagrams
   → Include: Color-coded sections, comment bubbles, execution flow arrows

9. **HIERARCHICAL/TAXONOMY** (Categories, Classifications, Organizational structures):
   → Use: Tree diagrams, organizational charts, nested boxes, category breakdowns
   → Include: Parent-child relationships, grouping containers, hierarchical levels

10. **TIMELINE/EVOLUTION** (History, Development stages, Lifecycle, Progression):
    → Use: Horizontal or vertical timelines, milestone markers, evolution sequences
    → Include: Dated markers, progression arrows, stage indicators, historical context

11. **PROBLEM-SOLUTION** (Challenges and fixes, Pain points and remedies):
    → Use: Before/after split, problem on left with solution on right, connected transformation
    → Include: Problem icons (X, warning), solution icons (checkmark, light bulb), transformation arrow

12. **MULTI-DIMENSIONAL COMPARISON** (Feature matrices, Decision frameworks, Trade-off analysis):
    → Use: Tables, matrix grids, radar charts, quadrant diagrams
    → Include: Multiple comparison criteria, scoring indicators, decision factors

**STEP 3: APPLY COGNITIVE LEARNING PRINCIPLES**

Apply these evidence-based design principles:

**Visual Hierarchy:**
- Most important element should be largest or most prominent
- Use size, color, and position to guide attention
- Create clear entry point for the eye (usually top-left or center)

**Chunking Information:**
- Group related elements together with proximity or containers
- Limit to 5-7 main chunks (cognitive load management)
- Use white space to separate distinct concepts

**Color Psychology & Accessibility:**
- Use color meaningfully (not decoratively):
  * Blue/Green: Trust, stability, growth (good for processes, data)
  * Orange/Yellow: Energy, attention, innovation (good for highlights, warnings)
  * Purple: Creativity, luxury (good for advanced concepts)
  * Red: Urgency, importance (use sparingly for critical points)
- Ensure sufficient contrast for readability (WCAG AA minimum)
- Don't rely solely on color to convey meaning

**Progressive Disclosure:**
- Start with high-level overview, then show details
- Use visual layers or numbered steps for complexity
- Don't overwhelm with everything at once

**Concrete Before Abstract:**
- Show real-world examples before theoretical concepts
- Use familiar metaphors and analogies
- Include recognizable icons and objects

**Dual Coding:**
- Combine visuals with minimal text labels
- Use icons that reinforce the concept (not random decoration)
- Ensure text and visuals work together, not compete

**STRICT OUTPUT RULES:**

1. **visual_type** → Be SPECIFIC, not generic:
   - ✓ "Architecture diagram with layered neural network structure"
   - ✓ "Process flowchart with decision points"
   - ✓ "Comparison infographic with side-by-side analysis"
   - ✗ "Technical infographic (minimal text)" ← Too generic!

2. **primary_subject** → Extract the EXACT core topic from the text

3. **key_elements** → List 5-8 SPECIFIC visual components that match the content:
   - For architecture: "input layer block", "hidden layer neurons", "output connections", "activation function icons"
   - For processes: "start node", "step 1 block", "decision diamond", "end node", "feedback loop arrow"
   - For comparisons: "left panel: method A", "right panel: method B", "comparison table", "pros/cons icons"
   - For real-world: "user icon interacting", "real device/product image", "industry setting background", "data flow from real scenario"
   - For concepts: "central concept circle", "supporting theory boxes", "relationship connectors", "formula breakdown"

4. **style_description** → Adapt based on content AND audience:
   - Technical/Engineering: "Clean technical diagram, engineering blueprint style, precise lines, professional CAD-like quality, monochromatic with accent colors"
   - Business/Management: "Modern infographic, corporate color palette (blues, grays), balanced layout, executive presentation quality, sophisticated and professional"
   - Science/Research: "Scientific illustration style, accurate representations, academic quality, clear labeling, neutral color scheme with strategic highlights"
   - Education/Training: "Engaging educational style, friendly colors, approachable design, clear visual metaphors, student-friendly complexity level"
   - Creative/Design: "Contemporary design, creative visual metaphors, vibrant color scheme, dynamic composition, innovative layouts"
   - Data/Analytics: "Data visualization style, clean charts, precise scaling, informative legends, analytical clarity"

5. **final_prompt** → CRITICAL - Build an INTELLIGENT, DETAILED prompt:

**ENHANCED TEMPLATE STRUCTURE:**

"Create a [SPECIFIC_VISUAL_TYPE] about **{primary_subject}** designed for [TARGET_AUDIENCE: students/professionals/general audience].

LEARNING OBJECTIVE: [What should viewers understand after seeing this?]

VISUAL COMPOSITION:
[Choose based on content type and include detailed instructions:]

FOR ARCHITECTURE:
- Design a multi-layered architecture diagram with clear component hierarchy
- Use distinct, labeled blocks for each layer/component with consistent sizing
- Show data flow with directional arrows between components (indicate direction clearly)
- Use different colors/shading for different layer types
- Include input/output indicators
- Add dimension labels (e.g., input size, layer depth)
- Visual metaphor: [Suggest if applicable, e.g., 'like a assembly line' or 'like stacked building blocks']

FOR PROCESS:
- Design a {linear/branching/cyclical} flowchart showing clear progression from start to finish
- Use numbered steps (1, 2, 3...) or clear sequential indicators
- Include decision points with diamond shapes if applicable
- Show feedback loops or iterations with curved return arrows where relevant
- Use color coding to distinguish different phases/stages
- Add time/duration indicators if relevant
- Highlight the critical path or most important steps

FOR COMPARISON:
- Create a {split-screen/table/matrix/overlapping Venn} comparison layout
- Clearly divide space between items being compared with visual separator
- Use consistent visual elements for fair comparison (same icon style, size, position)
- Highlight similarities with [approach] and differences with [contrasting approach]
- Include a legend or key if using color/symbol coding
- Add quantitative indicators if applicable (better/worse, faster/slower)
- Ensure both sides get equal visual weight

FOR REAL-WORLD APPLICATION:
- Illustrate realistic scenario with contextual environment (specify: office/hospital/factory/home/outdoor)
- Include recognizable objects that represent the actual use case
- Show human interaction or user journey through the system
- Use industry-standard icons and realistic proportions
- Add environmental details that provide context (screens, devices, tools, furniture)
- Show input and output in the real context
- Include diversity in human representations if people are shown

FOR CONCEPT/THEORY:
- Place main concept at {center/top} with {size appropriate to importance}
- Arrange supporting elements in {radial/hierarchical/network} relationship
- Use visual metaphors to represent abstract ideas (specify metaphor)
- Show connections with labeled lines indicating relationship types
- Use containers or grouping to show categories
- Include brief definitions or formulas where essential
- Layer complexity: background context → main concept → supporting details

FOR DATA/STATISTICS:
- Use {bar chart/line graph/pie chart/scatter plot/heat map} based on data type
- Include clearly labeled axes with units and scale
- Add a descriptive title and legend
- Show trends with arrows or annotations
- Use color coding with a clear key
- Highlight key data points or outliers
- Include comparison baselines or benchmarks if relevant
- Ensure data is accurately represented (no misleading scales)

FOR HIERARCHICAL:
- Design a {tree/pyramid/nested boxes} structure showing clear levels
- Distinguish parent-child relationships with connecting lines or containers
- Use size or position to indicate hierarchy level
- Include clear labels for each category/level
- Show quantity or scope at each level if applicable
- Use consistent indentation or alignment

FOR TIMELINE:
- Create a {horizontal/vertical/spiral} timeline with clear chronology
- Mark key milestones with icons or markers
- Include dates or time periods clearly
- Show duration with spacing or length
- Add annotations for significant events
- Use different colors for different eras or phases
- Include 'you are here' indicator if showing current state

FOR PROBLEM-SOLUTION:
- Split design: Problem side (left/top) and Solution side (right/bottom)
- Use visual contrast (dark vs light, cluttered vs clean)
- Show transformation with prominent connecting arrow or bridge
- Include problem indicators (pain points, inefficiencies, X marks)
- Show solution benefits (improvements, checkmarks, improvements)
- Use before/after imagery if applicable

FOR MULTI-DIMENSIONAL COMPARISON:
- Create a {matrix/radar chart/quadrant diagram/feature table}
- List all comparison dimensions clearly
- Use consistent rating or scoring system
- Color code for quick scanning (green=good, red=poor, or custom)
- Include a 'winner' or 'best for' recommendation if appropriate
- Show trade-offs visually

SPECIFIC KEY ELEMENTS TO INCLUDE:
{Detailed list of 6-9 elements with descriptions, not just names}
Example: Instead of "neural network icon", say "Interconnected nodes arranged in layers with weighted connections shown as varying line thickness"

COGNITIVE DESIGN PRINCIPLES TO APPLY:
- Visual hierarchy: {specify what should be most prominent}
- Chunking: {specify how to group related elements}
- Color usage: {specify color purpose - not just palette}
- Entry point: {where should the eye start?}
- Reading path: {what order should elements be viewed?}

STYLE AND QUALITY:
{style_description}

LAYOUT SPECIFICATIONS:
- Orientation: Landscape (1536×1024)
- Margins: Generous white space around edges
- Grid: {specify if needed: 2-column, 3-section, centered, etc.}
- Alignment: {specify: center-aligned, left-aligned, grid-based}
- Typography: Sans-serif, clean, modern, maximum 2-3 font sizes
- Consistency: All icons same style, all shapes same family, unified color palette

TEXT USAGE:
- Minimal labels: 2-5 words maximum per label
- Strategic placement: Near relevant visual elements
- Hierarchy: Title (large) → Section headers (medium) → Labels (small)
- Essential only: Use text to clarify, not to explain everything
- Avoid: Paragraphs, long sentences, redundant text

ACCESSIBILITY:
- High contrast between text and background (4.5:1 minimum)
- Don't rely on color alone to convey meaning
- Icons should be recognizable and distinct
- Text should be legible at small sizes

QUALITY STANDARDS:
- High-resolution, crisp lines and shapes
- Professional, presentation-ready quality
- Print-quality graphics (300+ DPI equivalent)
- No pixelation, blur, or distortion
- Balanced composition with intentional white space
- Suitable for projection on large screens

TARGET IMPACT:
After viewing this infographic, the learner should be able to {specify learning outcome}.
The visual should reduce cognitive load by {specify how: organizing complexity, showing relationships, providing structure, etc.}."

**INTELLIGENCE RULES:**
- Analyze FIRST: What type of content? What's the learning goal? Who's the audience?
- Choose WISELY: Match visual type to content nature and learning objective
- Be SPECIFIC: Every instruction should be detailed and actionable
- NO arrows unless there's actual flow, progression, causation, or connection to show
- NO flowcharts for static concepts, comparisons, or hierarchies
- NO generic "technical infographic" - be precise about the visual approach
- YES to real-world objects when discussing applications and use cases
- YES to architecture diagrams when discussing systems, structures, and frameworks
- YES to data visualizations when discussing metrics, performance, statistics
- YES to visual metaphors that make abstract concepts concrete
- ADAPT every element based on: content type, audience level, learning objective
- THINK about cognitive load: simplify without losing accuracy
- BALANCE detail and clarity: enough detail to be useful, enough clarity to be understood

Always output as a properly structured ImagePrompt object with all fields filled thoughtfully.
"""

    user_prompt = f"""
Analyze this content deeply and determine the OPTIMAL visual representation approach.

**SYSTEMATIC ANALYSIS - Think through each step:**

STEP 1 - Content Understanding:
- What is the main topic/concept?
- What are the key elements or components?
- What's the complexity level? (beginner/intermediate/advanced)

STEP 2 - Learning Objective:
- What should the viewer UNDERSTAND after seeing this?
- What's the primary cognitive task? (memorize, understand, apply, analyze, compare, evaluate)
- What's the most important insight or takeaway?

STEP 3 - Visual Strategy:
- What type of content is this? (architecture/process/concept/comparison/real-world/data/hierarchy/timeline/problem-solution/multi-comparison/math/code)
- What visual approach will maximize comprehension for THIS specific content?
- Should this show: structure, sequence, relationships, comparisons, applications, mechanisms?

STEP 4 - Specific Visual Elements:
- What 6-9 SPECIFIC visual elements are needed? (be detailed, not generic)
- Should there be arrows/flow? If yes, why? What do they show?
- What real-world objects or metaphors would help understanding?
- What visual hierarchy makes sense? (what should be largest/most prominent?)

STEP 5 - Audience Adaptation:
- Who is this for? (students learning basics, professionals needing reference, general audience)
- What's their likely prior knowledge level?
- Should this be simplified or detailed?

STEP 6 - Cognitive Design:
- How can we chunk this information into 5-7 meaningful groups?
- Where should the viewer's eye start? What's the reading path?
- What colors would support understanding (not just look pretty)?
- How do we reduce cognitive load while maintaining accuracy?

Content to visualize:
---
{presentation_text}
---

Now generate a comprehensive ImagePrompt object with:

1. **visual_type**: Be highly specific (e.g., "Three-layer neural network architecture diagram with activation functions" not just "technical infographic")

2. **primary_subject**: Extract the exact core topic

3. **key_elements**: List 6-9 DETAILED visual components with descriptions
   Example quality: "Input layer represented as 784 nodes arranged in 28x28 grid" vs. just "input layer"

4. **style_description**: Choose appropriate style for content type AND audience level
   Include: color palette rationale, typography style, visual tone (technical/friendly/formal/creative)

5. **final_prompt**: Create an EXCEPTIONALLY DETAILED prompt following the enhanced template structure that:
   - Specifies the learning objective
   - Provides detailed composition instructions for the chosen visual type
   - Lists specific elements with descriptions
   - Includes cognitive design principles
   - Specifies layout, style, accessibility requirements
   - Defines quality standards and target impact
   - Is 300-500 words of clear, actionable instructions

Be intelligent. Be thorough. Be specific. Create a prompt so detailed that the image generator knows EXACTLY what to create and why each element matters for learning.
"""

    try:
        response = client.responses.parse(
            model="gpt-4o",
            input=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            text_format=ImagePrompt, 
        )
        return response.output_parsed.final_prompt
    except Exception as e:
        print(f"An error occurred: {e}")
        return None


def _summarize_content_for_image(presentation_text: str) -> str:
    """Summarizes presentation text preserving key technical details and context."""
    if not client:
        print("Warning: No OpenAI client initialized. Content summarization skipped.")
        return presentation_text
    
    try:
        completion = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system", 
                    "content": "You are an expert at distilling technical content while preserving critical details for visualization. Identify the core concept, key components, and whether it's about architecture, process, comparison, real-world application, or theory."
                },
                {
                    "role": "user", 
                    "content": f"""Analyze and summarize this content in 60-80 words, preserving:
- The main topic/concept
- Key technical terms or components
- Whether it describes a structure, process, comparison, or concept
- Any real-world applications or examples mentioned
- Important relationships or flows

Content:
---
{presentation_text}
---

Focus on details needed to create an accurate, helpful visual representation."""
                }
            ],
            temperature=0.3,
            max_tokens=150,
        )
        summary = completion.choices[0].message.content
        # Clean up thinking tags if present
        summary = re.sub(r'<think>.*?</think>', '', summary, flags=re.DOTALL).strip()
        return summary
    except Exception as e:
        logger.warning(f"Could not summarize content, falling back to original text. Error: {e}")
        return presentation_text


def generate_infographic_img(presentation_text: str, filename: str = "infographic.png"):
    """
    Generate an intelligent, context-aware infographic image based on presentation content.
    Automatically adapts visual style to content type.
    Returns the path to the saved image file.
    Note: Image generation requires OpenAI (DALL-E) - OpenRouter doesn't support image generation.
    """
    if not client:
        print("Warning: OpenAI client not initialized. Image generation requires OpenAI API key.")
        return None
    
    # Summarize while preserving technical context
    summarized_content = _summarize_content_for_image(presentation_text)
    
    # Generate intelligent, adaptive prompt
    detailed_prompt = generate_image_infographic_prompt(summarized_content)
    if not detailed_prompt:
        raise ValueError("Failed to generate image prompt.")
    
    print("=" * 80)
    print("GENERATED INTELLIGENT PROMPT:")
    print("=" * 80)
    print(detailed_prompt)
    print("=" * 80)
    
    result = client.images.generate(
        model="gpt-image-1", 
        prompt=detailed_prompt,
        size="1536x1024",  
        quality="high", 
    )
    
    image_base64 = result.data[0].b64_json
    image_bytes = base64.b64decode(image_base64)
    
    with open(filename, "wb") as f:
        f.write(image_bytes)
    
    print(f"✓ High-quality infographic saved as {filename}")
    return filename