ai-game-2/tools/ora_editor/ora_ops.py

#!/usr/bin/env python3
"""ORA operations wrapper for the web editor."""

import io
import os
import re
import shutil
import sys
import tempfile
import zipfile
from pathlib import Path
from typing import Any
from xml.etree import ElementTree as ET

from PIL import Image, ImageChops


def parse_stack_xml(ora_path: str) -> ET.Element:
    """Parse stack.xml from an ORA file."""
    with zipfile.ZipFile(ora_path, 'r') as zf:
        return ET.fromstring(zf.read('stack.xml'))


def get_image_size(root: ET.Element) -> tuple[int, int]:
    """Get image dimensions from stack.xml root element."""
    return int(root.get('w', 0)), int(root.get('h', 0))


def get_groups_and_layers(root: ET.Element) -> list[dict[str, Any]]:
    """Extract groups and their layers from stack.xml."""
    groups = []
    root_stack = root.find('stack')
    if root_stack is None:
        return groups

    for child in root_stack:
        if child.tag == 'stack':
            group_name = child.get('name', 'unnamed')
            layers = []
            for layer in child:
                if layer.tag == 'layer':
                    layers.append({
                        'name': layer.get('name', 'unnamed'),
                        'src': layer.get('src', ''),
                        'opacity': layer.get('opacity', '1.0'),
                        'visibility': layer.get('visibility', 'visible')
                    })
            groups.append({
                'name': group_name,
                'is_group': True,
                'layers': layers
            })
        elif child.tag == 'layer':
            groups.append({
                'name': child.get('name', 'unnamed'),
                'is_group': False,
                'layers': [{
                    'name': child.get('name', 'unnamed'),
                    'src': child.get('src', ''),
                    'opacity': child.get('opacity', '1.0'),
                    'visibility': child.get('visibility', 'visible')
                }]
            })

    return groups


def get_next_layer_index(group: dict, entity_name: str) -> int:
    """Get the next auto-increment index for layers in an entity group."""
    max_index = -1
    pattern = re.compile(rf'^{re.escape(entity_name)}_(\d+)$')

    for layer in group['layers']:
        match = pattern.match(layer['name'])
        if match:
            max_index = max(max_index, int(match.group(1)))

    return max_index + 1


def find_entity_group(groups: list, entity_name: str) -> dict | None:
    """Find an existing entity group by name."""
    for group in groups:
        if group.get('is_group') and group['name'] == entity_name:
            return group
    return None


def find_layer(groups: list, layer_name: str) -> tuple[dict, dict] | None:
    """Find a layer by name across all groups. Returns (group, layer) tuple or None."""
    for group in groups:
        for layer in group['layers']:
            if layer['name'] == layer_name:
                return group, layer
    return None


def extract_layer_image(ora_path: str, layer_src: str) -> Image.Image:
    """Extract a layer image from the ORA file."""
    with zipfile.ZipFile(ora_path, 'r') as zf:
        image_data = zf.read(layer_src)
        return Image.open(io.BytesIO(image_data)).convert('RGBA')


def apply_mask(source_image: Image.Image, mask_image: Image.Image) -> Image.Image:
    """Apply a mask as the alpha channel of an image."""
    if mask_image.mode != 'L':
        mask_image = mask_image.convert('L')

    if mask_image.size != source_image.size:
        mask_image = mask_image.resize(source_image.size, Image.Resampling.BILINEAR)

    result = source_image.copy()
    r, g, b, a = result.split()
    new_alpha = ImageChops.multiply(a, mask_image)
    result.putalpha(new_alpha)

    return result


def build_stack_xml(width: int, height: int, groups: list[dict]) -> bytes:
    """Build stack.xml content from group structure."""
    image = ET.Element('image', {
        'version': '0.0.3',
        'w': str(width),
        'h': str(height)
    })

    root_stack = ET.SubElement(image, 'stack')

    for group in groups:
        if group.get('is_group'):
            group_el = ET.SubElement(root_stack, 'stack', {
                'name': group['name']
            })

            for layer in group['layers']:
                layer_attrs = {
                    'name': layer['name'],
                    'src': layer['src'],
                    'opacity': layer.get('opacity', '1.0')
                }
                if layer.get('visibility') and layer['visibility'] != 'visible':
                    layer_attrs['visibility'] = layer['visibility']
                ET.SubElement(group_el, 'layer', layer_attrs)
        else:
            layer = group['layers'][0]
            layer_attrs = {
                'name': layer['name'],
                'src': layer['src'],
                'opacity': layer.get('opacity', '1.0')
            }
            if layer.get('visibility') and layer['visibility'] != 'visible':
                layer_attrs['visibility'] = layer['visibility']
            ET.SubElement(root_stack, 'layer', layer_attrs)

    return ET.tostring(image, encoding='utf-8', xml_declaration=True)


def create_ora_from_structure(
    groups: list[dict],
    width: int,
    height: int,
    layer_images: dict[str, Image.Image],
    output_path: str
) -> None:
    """Create an ORA file from the structure and layer images."""
    tmp = tempfile.mkdtemp()

    try:
        data_dir = os.path.join(tmp, 'data')
        thumb_dir = os.path.join(tmp, 'Thumbnails')
        os.makedirs(data_dir)
        os.makedirs(thumb_dir)

        for src_path, img in layer_images.items():
            full_path = os.path.join(tmp, src_path)
            os.makedirs(os.path.dirname(full_path), exist_ok=True)
            img.save(full_path)

        stack_xml = build_stack_xml(width, height, groups)
        with open(os.path.join(tmp, 'stack.xml'), 'wb') as f:
            f.write(stack_xml)

        with open(os.path.join(tmp, 'mimetype'), 'w') as f:
            f.write('image/openraster')

        merged = None
        for group in groups:
            if not group.get('is_group') and group['layers']:
                layer_src = group['layers'][0]['src']
                merged = layer_images.get(layer_src)
                break

        if merged is None:
            merged = Image.new('RGBA', (width, height), (0, 0, 0, 0))

        merged.save(os.path.join(tmp, 'mergedimage.png'))

        thumb = merged.copy()
        thumb.thumbnail((256, 256))
        thumb.save(os.path.join(thumb_dir, 'thumbnail.png'))

        with zipfile.ZipFile(output_path, 'w') as zf:
            zf.write(os.path.join(tmp, 'mimetype'), 'mimetype',
                    compress_type=zipfile.ZIP_STORED)

            for root, _, files in os.walk(tmp):
                for file in files:
                    if file == 'mimetype':
                        continue
                    full = os.path.join(root, file)
                    rel = os.path.relpath(full, tmp)
                    zf.write(full, rel)

    finally:
        shutil.rmtree(tmp)


def load_ora(ora_path: str) -> dict[str, Any]:
    """Load an ORA file and return its structure and layer images."""
    root = parse_stack_xml(ora_path)
    width, height = get_image_size(root)
    groups = get_groups_and_layers(root)

    # Flatten layers with their group info
    layers = []
    for group in groups:
        for layer in group['layers']:
            layers.append({
                'name': layer['name'],
                'src': layer['src'],
                'group': group['name'] if group.get('is_group') else None,
                'visible': layer.get('visibility', 'visible') != 'hidden'
            })

    return {
        'width': width,
        'height': height,
        'layers': layers,
        'groups': groups
    }


def get_layer_image(ora_path: str, layer_name: str) -> Image.Image | None:
    """Get a specific layer image from the ORA."""
    root = parse_stack_xml(ora_path)
    groups = get_groups_and_layers(root)

    result = find_layer(groups, layer_name)
    if result is None:
        return None

    _, layer = result
    return extract_layer_image(ora_path, layer['src'])


def get_base_image(ora_path: str) -> Image.Image | None:
    """Get the base/merged image from the ORA."""
    try:
        with zipfile.ZipFile(ora_path, 'r') as zf:
            return Image.open(zf.open('mergedimage.png')).convert('RGBA')
    except Exception:
        return None


def save_ora(ora_path: str) -> bool:
    """Save the ORA structure back to disk (uses existing content)."""
    if not os.path.exists(ora_path):
        return False

    try:
        root = parse_stack_xml(ora_path)
        width, height = get_image_size(root)
        groups = get_groups_and_layers(root)

        # Extract all layer images from the current ORA
        layer_images = {}
        with zipfile.ZipFile(ora_path, 'r') as zf:
            for group in groups:
                for layer in group['layers']:
                    if layer['src'] not in layer_images:
                        try:
                            data = zf.read(layer['src'])
                            layer_images[layer['src']] = Image.open(io.BytesIO(data))
                        except KeyError:
                            pass

        create_ora_from_structure(groups, width, height, layer_images, ora_path)
        return True
    except Exception as e:
        print(f"Error saving ORA: {e}", file=sys.stderr)
        return False


def add_masked_layer(ora_path: str, entity_name: str, mask_path: str, source_layer: str = 'base') -> dict[str, Any]:
    """Add a new masked layer to the ORA."""
    # Check if ORA exists, create from PNG if not
    if not os.path.exists(ora_path):
        png_path = ora_path.replace('.ora', '.png')
        if os.path.exists(png_path):
            base = Image.open(png_path).convert('RGBA')
            width, height = base.size

            layer_src = 'data/base.png'
            groups = [{
                'name': 'base',
                'is_group': False,
                'layers': [{
                    'name': 'base',
                    'src': layer_src,
                    'opacity': '1.0',
                    'visibility': 'visible'
                }]
            }]

            layer_images = {layer_src: base}
            create_ora_from_structure(groups, width, height, layer_images, ora_path)
        else:
            return {'success': False, 'error': f"No ORA or PNG found at {ora_path}"}

    root = parse_stack_xml(ora_path)
    width, height = get_image_size(root)
    groups = get_groups_and_layers(root)

    # Find source layer
    layer_result = find_layer(groups, source_layer)
    if layer_result is None:
        return {'success': False, 'error': f"Layer '{source_layer}' not found"}

    source_group, source_layer_info = layer_result

    # Load source image
    with zipfile.ZipFile(ora_path, 'r') as zf:
        source_image_data = zf.read(source_layer_info['src'])
        source_image = Image.open(io.BytesIO(source_image_data)).convert('RGBA')

    # Load mask
    if not os.path.exists(mask_path):
        return {'success': False, 'error': f"Mask file not found: {mask_path}"}

    mask_image = Image.open(mask_path)
    if mask_image.mode not in ['L', 'RGBA', 'RGB']:
        mask_image = mask_image.convert('RGBA')

    # Apply mask
    masked_image = apply_mask(source_image, mask_image)

    # Find or create entity group
    entity_group = find_entity_group(groups, entity_name)
    if entity_group is None:
        entity_group = {
            'name': entity_name,
            'is_group': True,
            'layers': []
        }
        insert_idx = 0
        for i, g in enumerate(groups):
            if not g.get('is_group'):
                insert_idx = i
                break
            insert_idx = i + 1
        groups.insert(insert_idx, entity_group)

    # Generate new layer name
    next_index = get_next_layer_index(entity_group, entity_name)
    new_layer_name = f"{entity_name}_{next_index}"
    new_layer_src = f"data/{entity_name}/{new_layer_name}.png"

    # Add new layer to entity group
    entity_group['layers'].append({
        'name': new_layer_name,
        'src': new_layer_src,
        'opacity': '1.0',
        'visibility': 'visible'
    })

    # Extract all layer images from ORA
    layer_images = {}
    with zipfile.ZipFile(ora_path, 'r') as zf:
        for group in groups:
            for layer in group['layers']:
                if layer['src'] not in layer_images:
                    try:
                        data = zf.read(layer['src'])
                        layer_images[layer['src']] = Image.open(io.BytesIO(data))
                    except KeyError:
                        pass

    # Add the new masked image
    layer_images[new_layer_src] = masked_image

    create_ora_from_structure(groups, width, height, layer_images, ora_path)

    return {'success': True, 'layer_name': new_layer_name}


def rename_layer(ora_path: str, old_name: str, new_name: str) -> dict[str, Any]:
    """Rename a layer in the ORA."""
    root = parse_stack_xml(ora_path)
    width, height = get_image_size(root)
    groups = get_groups_and_layers(root)

    # Find the layer
    result = find_layer(groups, old_name)
    if result is None:
        return {'success': False, 'error': f"Layer '{old_name}' not found"}

    group, layer = result

    # Update name in both places
    old_src = layer['src']
    layer['name'] = new_name

    # Extract all layer images
    layer_images = {}
    with zipfile.ZipFile(ora_path, 'r') as zf:
        for g in groups:
            for l in g['layers']:
                if l['src'] not in layer_images:
                    try:
                        data = zf.read(l['src'])
                        layer_images[l['src']] = Image.open(io.BytesIO(data))
                    except KeyError:
                        pass

    # Get the image and create new entry with updated name
    if old_src in layer_images:
        img = layer_images.pop(old_src)
        ext = Path(old_src).suffix
        layer_dir = str(Path(old_src).parent)
        new_src = f"{layer_dir}/{new_name}{ext}"
        layer['src'] = new_src
        layer_images[new_src] = img

    create_ora_from_structure(groups, width, height, layer_images, ora_path)

    return {'success': True}


def delete_layer(ora_path: str, layer_name: str) -> dict[str, Any]:
    """Delete a layer from the ORA."""
    root = parse_stack_xml(ora_path)
    width, height = get_image_size(root)
    groups = get_groups_and_layers(root)

    # Find and remove the layer
    result = find_layer(groups, layer_name)
    if result is None:
        return {'success': False, 'error': f"Layer '{layer_name}' not found"}

    group, _ = result
    group['layers'] = [l for l in group['layers'] if l['name'] != layer_name]

    # Remove empty groups
    groups = [g for g in groups if len(g['layers']) > 0]

    # Extract remaining layer images
    layer_images = {}
    with zipfile.ZipFile(ora_path, 'r') as zf:
        for g in groups:
            for l in g['layers']:
                if l['src'] not in layer_images:
                    try:
                        data = zf.read(l['src'])
                        layer_images[l['src']] = Image.open(io.BytesIO(data))
                    except KeyError:
                        pass

    create_ora_from_structure(groups, width, height, layer_images, ora_path)

    return {'success': True}


def reorder_layer(ora_path: str, layer_name: str, direction: str) -> dict[str, Any]:
    """Move a layer up or down in the stack."""
    root = parse_stack_xml(ora_path)
    width, height = get_image_size(root)
    groups = get_groups_and_layers(root)

    # Find the layer
    result = find_layer(groups, layer_name)
    if result is None:
        return {'success': False, 'error': f"Layer '{layer_name}' not found"}

    current_group, _ = result

    # Flatten all layers with their group info for reordering
    flat_layers = []
    for g in groups:
        for l in g['layers']:
            flat_layers.append({'group': g, 'layer': l})

    # Find index of the layer
    idx = None
    for i, item in enumerate(flat_layers):
        if item['layer']['name'] == layer_name:
            idx = i
            break

    if idx is None:
        return {'success': False, 'error': f"Layer '{layer_name}' not found"}

    # Determine swap index
    swap_idx = None
    if direction == 'up' and idx > 0:
        swap_idx = idx - 1
    elif direction == 'down' and idx < len(flat_layers) - 1:
        swap_idx = idx + 1

    if swap_idx is None:
        return {'success': False, 'error': "Cannot move layer further in that direction"}

    # Swap the layers
    flat_layers[idx], flat_layers[swap_idx] = flat_layers[swap_idx], flat_layers[idx]

    # Rebuild groups
    new_groups = []
    current_group_obj = None

    for item in flat_layers:
        g_name = item['group']['name']
        is_group = item['group'].get('is_group')

        if is_group:
            # Entity layer - add to or create group
            existing_group = next((g for g in new_groups if g['name'] == g_name and g.get('is_group')), None)
            if existing_group is None:
                existing_group = {'name': g_name, 'is_group': True, 'layers': []}
                new_groups.append(existing_group)
            existing_group['layers'].append(item['layer'])
        else:
            # Base layer - standalone
            new_groups.append({
                'name': item['group']['name'],
                'is_group': False,
                'layers': [item['layer']]
            })

    groups = new_groups

    # Extract layer images
    layer_images = {}
    with zipfile.ZipFile(ora_path, 'r') as zf:
        for g in groups:
            for l in g['layers']:
                if l['src'] not in layer_images:
                    try:
                        data = zf.read(l['src'])
                        layer_images[l['src']] = Image.open(io.BytesIO(data))
                    except KeyError:
                        pass

    create_ora_from_structure(groups, width, height, layer_images, ora_path)

    return {'success': True}


def create_ora_from_png(png_path: str, output_ora: str, layer_name: str = 'base') -> dict[str, Any]:
    """Create an ORA file from a PNG."""
    if not os.path.exists(png_path):
        return {'success': False, 'error': f"File not found: {png_path}"}

    base = Image.open(png_path).convert('RGBA')
    width, height = base.size

    layer_src = f'data/{layer_name}.png'
    groups = [{
        'name': layer_name,
        'is_group': False,
        'layers': [{
            'name': layer_name,
            'src': layer_src,
            'opacity': '1.0',
            'visibility': 'visible'
        }]
    }]

    layer_images = {layer_src: base}
    create_ora_from_structure(groups, width, height, layer_images, output_ora)

    return {'success': True, 'ora_path': output_ora}


def get_layer_visibility(ora_path: str, layer_name: str) -> bool | None:
    """Get the visibility state of a layer."""
    root = parse_stack_xml(ora_path)
    groups = get_groups_and_layers(root)

    result = find_layer(groups, layer_name)
    if result is None:
        return None

    _, layer = result
    return layer.get('visibility', 'visible') != 'hidden'


def set_layer_visibility(ora_path: str, layer_name: str, visible: bool) -> dict[str, Any]:
    """Set the visibility of a layer."""
    root = parse_stack_xml(ora_path)
    width, height = get_image_size(root)
    groups = get_groups_and_layers(root)

    # Find and update the layer
    result = find_layer(groups, layer_name)
    if result is None:
        return {'success': False, 'error': f"Layer '{layer_name}' not found"}

    _, layer = result
    layer['visibility'] = 'visible' if visible else 'hidden'

    # Extract layer images
    layer_images = {}
    with zipfile.ZipFile(ora_path, 'r') as zf:
        for g in groups:
            for l in g['layers']:
                if l['src'] not in layer_images:
                    try:
                        data = zf.read(l['src'])
                        layer_images[l['src']] = Image.open(io.BytesIO(data))
                    except KeyError:
                        pass

    create_ora_from_structure(groups, width, height, layer_images, ora_path)

    return {'success': True}