ai-game-2/tools/ora_editor/routes/images.py

"""Image serving routes for ORA Editor."""

import io
import zipfile
from flask import Blueprint, request, jsonify, send_file, Response
from pathlib import Path
from PIL import Image, ImageDraw

from ora_editor.ora_ops import parse_stack_xml
from ora_editor.config import PROJECT_ROOT

images_bp = Blueprint('images', __name__)


@images_bp.route('/api/image/<layer_name>')
def api_image(layer_name):
    """Serve a specific layer image."""
    ora_path = request.args.get('ora_path')
    
    if not ora_path:
        return jsonify({'error': 'Missing ora_path'}), 400
    
    try:
        with zipfile.ZipFile(ora_path, 'r') as zf:
            img_data = zf.read('mergedimage.png')
            return send_file(io.BytesIO(img_data), mimetype='image/png')
    except Exception as e:
        return Response(f"Error loading image: {e}", status=500)


@images_bp.route('/api/image/base')
def api_image_base():
    """Serve the base/merged image."""
    ora_path = request.args.get('ora_path')
    
    if not ora_path:
        return jsonify({'error': 'Missing ora_path'}), 400
    
    try:
        with zipfile.ZipFile(ora_path, 'r') as zf:
            img_data = zf.read('mergedimage.png')
            return send_file(io.BytesIO(img_data), mimetype='image/png')
    except Exception as e:
        return Response(f"Error loading image: {e}", status=500)


@images_bp.route('/api/image/layer/<layer_name>')
def api_image_layer(layer_name):
    """Serve a specific layer as image."""
    ora_path = request.args.get('ora_path')
    
    if not ora_path:
        return jsonify({'error': 'Missing ora_path'}), 400
    
    try:
        from xml.etree import ElementTree as ET
        
        with zipfile.ZipFile(ora_path, 'r') as zf:
            stack = ET.fromstring(zf.read('stack.xml'))
            
            for child in stack.find('stack'):
                if child.tag == 'layer' and child.get('name') == layer_name:
                    src = child.get('src')
                    img_data = zf.read(src)
                    return send_file(io.BytesIO(img_data), mimetype='image/png')
                
                elif child.tag == 'stack':
                    for layer in child:
                        if layer.tag == 'layer' and layer.get('name') == layer_name:
                            src = layer.get('src')
                            img_data = zf.read(src)
                            return send_file(io.BytesIO(img_data), mimetype='image/png')
        
        return jsonify({'error': 'Layer not found'}), 404
    except Exception as e:
        import traceback
        traceback.print_exc()
        return Response(f"Error loading layer: {e}", status=500)


@images_bp.route('/api/image/polygon')
def api_image_polygon():
    """Serve polygon overlay if stored."""
    from ora_editor.services import polygon_storage
    
    ora_path = request.args.get('ora_path')
    
    if not ora_path:
        return Response("Missing ora_path", status=400)
    
    poly_data = polygon_storage.get(ora_path)
    if not poly_data:
        return Response("No polygon data", status=404)
    
    try:
        with zipfile.ZipFile(ora_path, 'r') as zf:
            base_img = Image.open(zf.open('mergedimage.png')).convert('RGBA')
        
        points = poly_data.get('points', [])
        color = poly_data.get('color', '#FF0000')
        width = poly_data.get('width', 2)
        
        if len(points) < 3:
            return Response("Not enough points", status=404)
        
        w, h = base_img.size
        pixel_points = [(int(p['x'] * w), int(p['y'] * h)) for p in points]
        
        draw = ImageDraw.Draw(base_img)
        hex_color = color if len(color) == 7 else color + 'FF'
        draw.polygon(pixel_points, outline=hex_color, width=width)
        
        img_io = io.BytesIO()
        base_img.save(img_io, format='PNG')
        img_io.seek(0)
        
        return send_file(img_io, mimetype='image/png')
    except Exception as e:
        return Response(f"Error drawing polygon: {e}", status=500)


@images_bp.route('/api/image/masked')
def api_image_masked():
    """Serve the base image with mask applied as alpha channel."""
    ora_path = request.args.get('ora_path')
    mask_path = request.args.get('mask_path')
    
    if not ora_path or not mask_path:
        return jsonify({'error': 'Missing ora_path or mask_path'}), 400
    
    try:
        with zipfile.ZipFile(ora_path, 'r') as zf:
            base_img = Image.open(zf.open('mergedimage.png')).convert('RGBA')
        
        mask_img = Image.open(mask_path).convert('L')
        
        if mask_img.size != base_img.size:
            mask_img = mask_img.resize(base_img.size, Image.Resampling.BILINEAR)
        
        r, g, b, _ = base_img.split()
        result = Image.merge('RGBA', (r, g, b, mask_img))
        
        img_io = io.BytesIO()
        result.save(img_io, format='PNG')
        img_io.seek(0)
        
        return send_file(img_io, mimetype='image/png')
    except Exception as e:
        return Response(f"Error creating masked image: {e}", status=500)


@images_bp.route('/api/file/mask')
def api_file_mask():
    """Serve a mask file from temp directory."""
    path = request.args.get('path')
    
    if not path:
        return jsonify({'error': 'Missing path'}), 400
    
    full_path = Path(path)
    
    if not full_path.exists():
        return jsonify({'error': 'File not found'}), 404
    
    return send_file(full_path, mimetype='image/png')