#!/usr/bin/env python3
"""
ORA (OpenRaster) editing utility.

Supports creating ORA files from PNGs, extracting layers as PNGs, and masking 
elements with alpha channels.

Commands:
    create          Create an ORA file from a PNG
    mask_element    Create a masked layer for an entity in an ORA file
    inspect         Show the structure of an ORA file
    extract_png     Extract a layer from an ORA file as a PNG
"""

import argparse
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.
    Returns a list of group dicts with 'name', 'layers', and 'is_group' keys.
    """
    groups = []
    root_stack = root.find('stack')
    if root_stack is None:
        return groups

    for child in root_stack:
        if child.tag == 'stack':
            # This is a group
            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':
            # This is a standalone layer (base 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.
    Looks for pattern: entity_name_N
    """
    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.
    Mask can be grayscale or RGB; will be converted to grayscale for alpha.
    """
    # Ensure mask is grayscale
    if mask_image.mode != 'L':
        mask_image = mask_image.convert('L')

    # Resize mask to match source if needed
    if mask_image.size != source_image.size:
        mask_image = mask_image.resize(source_image.size, Image.Resampling.BILINEAR)

    # Apply mask to alpha channel
    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'):
            # Create group element
            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:
            # Standalone layer (base layer)
            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)

        # Save all layer images
        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)

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

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

        # Create mergedimage.png (use first base layer)
        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:
            # Create blank image if no base layer
            merged = Image.new('RGBA', (width, height), (0, 0, 0, 0))

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

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

        # Build ORA zip
        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 print_ora_summary(groups: list[dict]) -> None:
    """Print an itemized summary of the ORA structure."""
    print("\n" + "=" * 50)
    print("ORA STRUCTURE SUMMARY")
    print("=" * 50)

    for i, group in enumerate(groups, 1):
        if group.get('is_group'):
            print(f"\n📁 Group {i}: {group['name']}")
            for j, layer in enumerate(group['layers'], 1):
                print(f"   └─ Layer {j}: {layer['name']}")
        else:
            layer = group['layers'][0]
            print(f"\n🖼️  Base Layer {i}: {layer['name']}")

    print("\n" + "=" * 50)


def cmd_create(args: argparse.Namespace) -> int:
    """Handle 'create' subcommand."""
    input_png = args.input_png
    output_ora = args.output_ora
    layer_name = args.layer_name

    if output_ora is None:
        output_ora = str(Path(input_png).with_suffix('.ora'))

    # Load base image
    base = Image.open(input_png).convert('RGBA')
    width, height = base.size

    # Build structure
    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)

    print(f"✓ Created ORA: {output_ora}")
    print_ora_summary(groups)

    return 0


def cmd_mask_element(args: argparse.Namespace) -> int:
    """Handle 'mask_element' subcommand."""
    import io

    input_path = args.input
    mask_path = args.mask
    entity_name = args.entity
    layer_name = args.layer
    output_path = args.output

    # Handle PNG input: auto-create ORA first
    is_png_input = input_path.lower().endswith('.png')
    ora_path = input_path

    if is_png_input:
        ora_path = str(Path(input_path).with_suffix('.ora'))
        if output_path is None:
            output_path = ora_path

        # Create ORA from PNG with 'base' layer
        base = Image.open(input_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)
        print(f"✓ Converted PNG to ORA: {ora_path}")

    # Parse existing ORA
    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, layer_name)
    if layer_result is None:
        print(f"Error: Layer '{layer_name}' not found in ORA file", file=sys.stderr)
        print("Available layers:", file=sys.stderr)
        for group in groups:
            for layer in group['layers']:
                print(f"  - {layer['name']}", file=sys.stderr)
        return 1

    source_group, source_layer = layer_result

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

    # Load mask
    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:
        # Create new group before base layers
        entity_group = {
            'name': entity_name,
            'is_group': True,
            'layers': []
        }
        # Insert before first non-group (base layer)
        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 (auto-increment)
    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  # New layer, will be added later

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

    # Determine output path
    if output_path is None:
        output_path = ora_path

    # Create new ORA
    create_ora_from_structure(groups, width, height, layer_images, output_path)

    print(f"✓ Applied mask to create: {entity_name}/{new_layer_name}")
    print(f"  Source layer: {layer_name}")
    print(f"  Mask file: {mask_path}")
    print(f"  Output: {output_path}")
    print_ora_summary(groups)

    return 0


def cmd_inspect(args: argparse.Namespace) -> int:
    """Handle 'inspect' subcommand."""
    ora_path = args.ora_file

    if not os.path.exists(ora_path):
        print(f"Error: File not found: {ora_path}", file=sys.stderr)
        return 1

    # Parse existing ORA
    try:
        root = parse_stack_xml(ora_path)
    except (zipfile.BadZipFile, KeyError) as e:
        print(f"Error: Invalid ORA file: {ora_path}", file=sys.stderr)
        return 1

    groups = get_groups_and_layers(root)

    print(f"\nFile: {ora_path}")
    print_ora_summary(groups)

    return 0


def cmd_extract_png(args: argparse.Namespace) -> int:
    """Handle 'extract_png' subcommand."""
    ora_path = args.ora_file
    layer_name = args.layer
    output_path = args.output

    if not os.path.exists(ora_path):
        print(f"Error: File not found: {ora_path}", file=sys.stderr)
        return 1

    # Parse existing ORA
    try:
        root = parse_stack_xml(ora_path)
    except (zipfile.BadZipFile, KeyError) as e:
        print(f"Error: Invalid ORA file: {ora_path}", file=sys.stderr)
        return 1

    groups = get_groups_and_layers(root)

    # Find the layer
    layer_result = find_layer(groups, layer_name)
    if layer_result is None:
        print(f"Error: Layer '{layer_name}' not found in ORA file", file=sys.stderr)
        print("Available layers:", file=sys.stderr)
        for group in groups:
            for layer in group['layers']:
                print(f"  - {layer['name']}", file=sys.stderr)
        return 1

    _, layer = layer_result

    # Determine output path
    if output_path is None:
        output_path = f"{layer_name}.png"

    # Extract and save the layer image
    image = extract_layer_image(ora_path, layer['src'])
    image.save(output_path)

    print(f"✓ Extracted layer '{layer_name}' to: {output_path}")

    return 0


def main() -> int:
    parser = argparse.ArgumentParser(
        prog='ora_edit.py',
        description='ORA (OpenRaster) editing utility'
    )

    subparsers = parser.add_subparsers(dest='command', required=True)

    # Create subcommand
    create_parser = subparsers.add_parser(
        'create',
        help='Create an ORA file from a PNG'
    )
    create_parser.add_argument(
        'input_png',
        help='Input PNG file'
    )
    create_parser.add_argument(
        'output_ora',
        nargs='?',
        default=None,
        help='Output ORA file (default: same name as input with .ora extension)'
    )
    create_parser.add_argument(
        '--layer-name',
        default='base',
        help='Name for the root layer (default: base)'
    )
    create_parser.set_defaults(func=cmd_create)

    # Mask element subcommand
    mask_parser = subparsers.add_parser(
        'mask_element',
        help='Create a masked layer for an entity in an ORA file'
    )
    mask_parser.add_argument(
        'input',
        help='Input ORA or PNG file'
    )
    mask_parser.add_argument(
        '--mask',
        required=True,
        help='Mask file to use as alpha channel'
    )
    mask_parser.add_argument(
        '--entity',
        required=True,
        help='Entity name (group name)'
    )
    mask_parser.add_argument(
        '--layer',
        default='base',
        help='Source layer to copy and mask (default: base)'
    )
    mask_parser.add_argument(
        '--output',
        default=None,
        help='Output ORA file (default: same as input)'
    )
    mask_parser.set_defaults(func=cmd_mask_element)

    # Inspect subcommand
    inspect_parser = subparsers.add_parser(
        'inspect',
        help='Show the structure of an ORA file'
    )
    inspect_parser.add_argument(
        'ora_file',
        help='ORA file to inspect'
    )
    inspect_parser.set_defaults(func=cmd_inspect)

    # Extract PNG subcommand
    extract_parser = subparsers.add_parser(
        'extract_png',
        help='Extract a layer from an ORA file as a PNG'
    )
    extract_parser.add_argument(
        'ora_file',
        help='ORA file to extract from'
    )
    extract_parser.add_argument(
        '--layer',
        required=True,
        help='Name of the layer to extract'
    )
    extract_parser.add_argument(
        '--output',
        default=None,
        help='Output PNG file (default: <layer_name>.png)'
    )
    extract_parser.set_defaults(func=cmd_extract_png)

    args = parser.parse_args()
    return args.func(args)


if __name__ == '__main__':
    sys.exit(main())