"""
Generate documentation images for the project.
"""
from PIL import Image, ImageDraw, ImageFont
import os

def create_directory(path):
    """Create directory if it doesn't exist."""
    if not os.path.exists(path):
        os.makedirs(path)

def create_interface_screenshot():
    """Create a mock interface screenshot."""
    # Create a new image with a white background
    img = Image.new('RGB', (1200, 800), 'white')
    draw = ImageDraw.Draw(img)
    
    # Add title
    draw.rectangle([(50, 50), (1150, 100)], fill='#f0f0f0')
    draw.text((100, 60), "Agent2Robot - AI-Powered Robot Design", fill='black')
    
    # Add main sections
    sections = [
        ("Design Parameters", 150, 200, 500),
        ("Simulation View", 550, 200, 500),
        ("Performance Metrics", 150, 500, 500),
        ("Design History", 550, 500, 500)
    ]
    
    for title, x, y, width in sections:
        draw.rectangle([(x, y), (x + width, y + 200)], fill='#e8e8e8')
        draw.text((x + 20, y + 20), title, fill='black')
    
    # Save the image
    img.save('docs/images/interface.png')

def create_simulation_gif():
    """Create a mock simulation GIF."""
    # Create a series of frames
    frames = []
    for i in range(10):
        img = Image.new('RGB', (400, 300), 'white')
        draw = ImageDraw.Draw(img)
        
        # Draw a simple robot animation
        x = 50 + i * 30
        draw.rectangle([(x, 150), (x + 100, 200)], fill='blue')
        draw.ellipse([(x + 20, 200), (x + 40, 220)], fill='black')
        draw.ellipse([(x + 60, 200), (x + 80, 220)], fill='black')
        
        frames.append(img)
    
    # Save as GIF
    frames[0].save(
        'docs/images/simulation.gif',
        save_all=True,
        append_images=frames[1:],
        duration=200,
        loop=0
    )

def create_workflow_diagram():
    """Create a workflow diagram."""
    img = Image.new('RGB', (800, 600), 'white')
    draw = ImageDraw.Draw(img)
    
    # Draw workflow boxes
    boxes = [
        ("Design Input", 100, 100),
        ("AI Processing", 300, 100),
        ("Physics Simulation", 500, 100),
        ("Performance Analysis", 300, 300),
        ("Design Optimization", 300, 500)
    ]
    
    for text, x, y in boxes:
        draw.rectangle([(x, y), (x + 150, y + 50)], fill='#e8e8e8')
        draw.text((x + 10, y + 15), text, fill='black')
    
    # Draw arrows
    arrows = [
        (250, 125, 300, 125),
        (450, 125, 500, 125),
        (575, 150, 375, 300),
        (375, 350, 375, 500)
    ]
    
    for x1, y1, x2, y2 in arrows:
        draw.line([(x1, y1), (x2, y2)], fill='black', width=2)
    
    img.save('docs/images/design_flow.png')

def create_architecture_diagram():
    """Create an architecture diagram."""
    img = Image.new('RGB', (1000, 800), 'white')
    draw = ImageDraw.Draw(img)
    
    # Draw components
    components = [
        ("Gradio Interface", 100, 100),
        ("LLM Integration", 400, 100),
        ("Physics Engine", 700, 100),
        ("Design Generator", 400, 300),
        ("Simulation Manager", 400, 500),
        ("Performance Analyzer", 400, 700)
    ]
    
    for text, x, y in components:
        draw.rectangle([(x, y), (x + 200, y + 60)], fill='#e8e8e8')
        draw.text((x + 10, y + 20), text, fill='black')
    
    # Draw connections
    connections = [
        (300, 130, 400, 130),
        (600, 130, 700, 130),
        (500, 160, 500, 300),
        (500, 360, 500, 500),
        (500, 560, 500, 700)
    ]
    
    for x1, y1, x2, y2 in connections:
        draw.line([(x1, y1), (x2, y2)], fill='black', width=2)
    
    img.save('docs/images/architecture.png')

def main():
    """Generate all documentation images."""
    create_directory('docs/images')
    create_interface_screenshot()
    create_simulation_gif()
    create_workflow_diagram()
    create_architecture_diagram()
    print("Documentation images generated successfully!")

if __name__ == "__main__":
    main()