#!/usr/bin/env python

# The basic area editor
#
# To edit an existing level, use
# editor levelname
#
# To create a new level:
#
# editor levelname <xsize> <ysiz>
# (size specified in pixels
#

import os
import sys

import pygame
import pygame.locals as pgl

sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

import pymunk

from albow.root import RootWidget
from albow.widget import Widget
from albow.controls import Button, Label, CheckBox
from albow.dialogs import alert, Dialog
from albow.layout import Row
from albow.fields import TextField
from albow.table_view import TableView, TableColumn

from nagslang.options import parse_args
from nagslang.constants import SCREEN
from nagslang.level import Level, POLY_COLORS, LINE_COLOR
from nagslang.yamlish import load_s
import nagslang.enemies as ne
import nagslang.game_object as ngo
import nagslang.puzzle as np

# layout constants
MENU_BUTTON_HEIGHT = 35
MENU_PAD = 6
MENU_HALF_PAD = MENU_PAD // 2
MENU_LEFT = SCREEN[0] + MENU_HALF_PAD
MENU_WIDTH = 200 - MENU_PAD

BUTTON_RECT = pygame.rect.Rect(0, 0, MENU_WIDTH, MENU_BUTTON_HEIGHT)
CHECK_RECT = pygame.rect.Rect(0, 0, MENU_BUTTON_HEIGHT // 2,
                              MENU_BUTTON_HEIGHT // 2)


class EditorLevel(Level):

    def __init__(self, name, x=800, y=600):
        super(EditorLevel, self).__init__(name)
        self.x = x
        self.y = y

    def round_point(self, pos):
        return (10 * (pos[0] // 10), 10 * (pos[1] // 10))

    def point_to_pymunk(self, pos):
        # inverse of point_to_pygame
        # (this is also the same as point_to_pygame, but a additional
        # function for sanity later in pyweek).
        return (pos[0], self.y - pos[1])

    def add_point(self, poly_index, pos):
        self.polygons.setdefault(poly_index, [])
        if not self.polygons[poly_index]:
            point = self.point_to_pymunk(self.round_point(pos))
            self.polygons[poly_index].append(point)
        else:
            add_pos = self.fix_poly_angle(poly_index, pos)
            self.polygons[poly_index].append(add_pos)

    def _fix_angle(self, point1, pos):
        # We want the line (point1 to pos) to be an angle of
        # 0, 45, 90, 135, 180, 225, 270, 305
        # However, we only need to consider half the circle
        # This is a hack to approximate the right thing
        pos0 = (pos[0], point1[1])
        pos90 = (point1[0], pos[1])
        dist = max(abs(point1[0] - pos[0]), abs(point1[1] - pos[1]))
        pos45 = (point1[0] + dist, point1[1] + dist)
        pos135 = (point1[0] + dist, point1[1] - dist)
        pos225 = (point1[0] - dist, point1[1] - dist)
        pos305 = (point1[0] - dist, point1[1] + dist)
        min_dist = 9999999
        new_pos = point1
        for cand in [pos0, pos90, pos45, pos135, pos225, pos305]:
            dist = (pos[0] - cand[0]) ** 2 + (pos[1] - cand[1]) ** 2
            if dist < min_dist:
                new_pos = cand
                min_dist = dist
        return self.point_to_pymunk(new_pos)

    def fix_line_angle(self, start_pos, pos):
        start_pos = self.round_point(start_pos)
        pos = self.round_point(pos)
        return self._fix_angle(start_pos, pos)

    def fix_poly_angle(self, index, pos):
        # Last point
        point1 = self.point_to_pygame(self.polygons[index][-1])
        pos = self.round_point(pos)
        return self._fix_angle(point1, pos)

    def delete_point(self, index):
        if index in self.polygons and len(self.polygons[index]) > 0:
            self.polygons[index].pop()

    def close_poly(self, index):
        """Attempts to close the current polygon.

           We allow a small additional step to close the polygon, but
           it's limited as it's a magic point addition"""
        if len(self.polygons[index]) < 2:
            # Too small
            return False
        first = self.polygons[index][0]
        if self.fix_poly_angle(index, self.point_to_pygame(first)) == first:
            self.add_point(index, self.point_to_pygame(first))
            return True
        candidates = [(first[0] + 10 * i, first[1]) for
                      i in (-3, -2, -1, 1, 2, 3)]
        candidates.extend([(first[0], first[1] + 10 * i) for
                          i in (-3, -2, -1, 1, 2, 3)])
        candidates.extend([(first[0] + 10 * i, first[1] + 10 * i) for
                          i in (-3, -2, -1, 1, 2, 3)])
        candidates.extend([(first[0] + 10 * i, first[1] - 10 * i) for
                          i in (-3, -2, -1, 1, 2, 3)])
        min_dist = 99999
        poss = None
        for cand in candidates:
            if self.fix_poly_angle(index, self.point_to_pygame(cand)) == cand:
                dist = (first[0] - cand[0]) ** 2 + (first[1] - cand[1]) ** 2
                if dist < min_dist:
                    poss = cand
        if poss is not None:
            self.add_point(index, self.point_to_pygame(poss))
            self.add_point(index, self.point_to_pygame(first))
            return True
        return False

    def add_line(self, start_pos, end_pos):
        endpoint = self.fix_line_angle(start_pos, end_pos)
        startpoint = self.point_to_pymunk(self.round_point(start_pos))
        self.lines.append([startpoint, endpoint])

    def draw(self, mouse_pos, mouse_poly, filled, draw_cand_line, start_pos):
        self._draw_background(True)
        # Draw polygons as needed for the editor
        if filled:
            self._draw_exterior(True)
        for index, polygon in self.polygons.items():
            color = POLY_COLORS[index]
            if len(polygon) > 1:
                pointlist = [self.point_to_pygame(p) for p in polygon]
                pygame.draw.lines(self._surface, color, False, pointlist, 2)
            if index == mouse_poly and mouse_pos:
                endpoint = self.fix_poly_angle(index, mouse_pos)
                pygame.draw.line(self._surface, color,
                                 self.point_to_pygame(polygon[-1]),
                                 self.point_to_pygame(endpoint))
        for line in self.lines:
            pointlist = [self.point_to_pygame(p) for p in line]
            pygame.draw.lines(self._surface, LINE_COLOR, False, pointlist, 2)
        if draw_cand_line and start_pos and mouse_pos:
            endpoint = self.fix_line_angle(start_pos, mouse_pos)
            pointlist = [self.round_point(start_pos),
                         self.point_to_pygame(endpoint)]
            pygame.draw.lines(self._surface, LINE_COLOR, False, pointlist, 1)
        return self._surface.copy()

    def reset_objs(self):
        # Reset the object state - needed when changing stuff
        self.drawables = []
        self.overlay_drawables = []
        self._glue = np.PuzzleGlue()
        for game_object_dict in self._game_objects:
            self._create_game_object(pymunk.Space(), **game_object_dict)
        for enemy_dict in self._enemies:
            self._create_enemy(pymunk.Space(), **enemy_dict)

    def get_class(self, classname, mod=None):
        # Get the class given the classname
        modules = {
            'game_object': ngo,
            'enemies': ne,
            'puzzle': np,
            }
        if '.' in classname:
            modname, classname = classname.split('.')
            mod = modules[modname]
        if mod is None:
            mod = ngo
        return getattr(mod, classname)

    def try_new_object(self, target, new, old=None):
        if old in target:
            target.remove(old)
        try:
            target.append(new)
            self.reset_objs()
            return True
        except Exception:
            target.remove(new)
            if old is not None:
                target.append(old)
            self.reset_objs()
        return False


class ObjectTable(TableView):

    columns = [TableColumn("Object", 690, 'l', '%r')]

    def __init__(self, data):
        super(ObjectTable, self).__init__(height=450)
        self.data = data
        self.selected_row = -1

    def num_rows(self):
        return len(self.data)

    def row_data(self, i):
        data = self.data[i]
        if 'name' in data:
            return ('%s (%s)' % (data['classname'], data['name']), )
        return (data['classname'], )

    def row_is_selected(self, i):
        return self.selected_row == i

    def click_row(self, i, ev):
        self.selected_row = i

    def get_selection(self):
        if self.selected_row >= 0:
            return self.data[self.selected_row]
        return None


class EditClassDialog(Dialog):

    def __init__(self, classname, cls, data):
        super(EditClassDialog, self).__init__()
        self.classname = classname
        self.rect = pygame.rect.Rect(0, 0, 800, 550)
        title = Label("Editing %s" % classname)
        title.rect = pygame.rect.Rect(100, 10, 600, 25)
        self.add(title)
        requires = cls.requires()
        y = 40
        self.fields = {}
        index = 0
        for requirement, hint in requires:
            label = Label(requirement)
            label.rect = pygame.rect.Rect(40, y, 200, 25)
            self.add(label)
            field = TextField()
            field.rect = pygame.rect.Rect(220, y, 400, 25)
            self.add(field)
            if data is not None:
                if requirement in data:
                    field.set_text('%s' % data[requirement])
                elif 'args' in data and requirement != 'name':
                    # NB: The ordering assumptions in requires should make
                    # this safe, but it's really, really, really fragile
                    try:
                        field.set_text('%s' % data['args'][index])
                        index += 1
                    except IndexError:
                        # Assumed to be arguments with the default value
                        pass
            self.fields[requirement] = field
            hintlabel = Label(hint)
            hintlabel.rect = pygame.rect.Rect(640, y, 100, 25)
            self.add(hintlabel)
            y += 30
        buttons = []
        for text in ['OK', 'Cancel']:
            but = Button(text, action=lambda x=text: self.dismiss(x))
            buttons.append(but)
        row = Row(buttons)
        row.rect = pygame.rect.Rect(250, 500, 700, 50)
        self.add(row)

    def get_data(self):
        result = {}
        result['classname'] = self.classname
        args = []
        # We arrange to bounce this through yaml'ish to convert
        # stuff to the expected type
        for val in self.fields:
            text = self.fields[val].get_text()
            if not text:
                # skip empty fields
                continue
            if val == 'name':
                result['name'] = text
            else:
                args.append(' - ' + text)
        data = "args:\n" + '\n'.join(args)
        result['args'] = load_s(data)['args']
        return result


class LevelWidget(Widget):

    def __init__(self, level):
        super(LevelWidget, self).__init__(pygame.rect.Rect(0, 0,
                                          SCREEN[0], SCREEN[1]))
        self.level = level
        self.pos = (0, 0)
        self.filled_mode = False
        self.mouse_pos = None
        self.cur_poly = None
        self._mouse_drag = False
        self._draw_objects = False
        self._draw_enemies = False
        self._draw_lines = False
        self._start_pos = None

    def _level_coordinates(self, pos):
        # Move positions to level values
        if not pos:
            return (0, 0)
        return pos[0] + self.pos[0], pos[1] + self.pos[1]

    def _move_view(self, offset):
        new_pos = [self.pos[0] + offset[0], self.pos[1] + offset[1]]
        if new_pos[0] < 0:
            new_pos[0] = self.pos[0]
        elif new_pos[0] > self.level.x - SCREEN[0]:
            new_pos[0] = self.pos[0]
        if new_pos[1] < 0:
            new_pos[1] = self.pos[1]
        elif new_pos[1] > self.level.y - SCREEN[1]:
            new_pos[1] = self.pos[1]
        self.pos = tuple(new_pos)

    def set_objects(self, value):
        if self._draw_objects != value:
            self._draw_objects = value
            self.invalidate()

    def set_enemies(self, value):
        if self._draw_enemies != value:
            self._draw_enemies = value
            self.invalidate()

    def draw(self, surface):
        if (self.cur_poly is not None and self.cur_poly in self.level.polygons
                and len(self.level.polygons[self.cur_poly])):
            # We have an active polygon
            mouse_pos = self._level_coordinates(self.mouse_pos)
        elif self._draw_lines:
            # Interior wall mode
            mouse_pos = self._level_coordinates(self.mouse_pos)
        else:
            mouse_pos = None
        level_surface = level.draw(mouse_pos, self.cur_poly, self.filled_mode,
                                   self._draw_lines, self._start_pos)
        if self._draw_objects:
            for thing in self.level.drawables:
                if not isinstance(thing, ne.Enemy):
                    thing.render(level_surface)
        if self._draw_enemies:
            for thing in self.level.drawables:
                if isinstance(thing, ne.Enemy):
                    thing.render(level_surface)
        surface_area = pygame.rect.Rect(self.pos, SCREEN)
        surface.blit(level_surface, (0, 0), surface_area)

    def change_poly(self, new_poly):
        self.cur_poly = new_poly
        self._draw_lines = False
        if self.cur_poly is not None:
            self.filled_mode = False

    def line_mode(self):
        self.cur_poly = None
        self._draw_lines = True
        self.filled_mode = False
        self._start_pos = None

    def key_down(self, ev):
        if ev.key == pgl.K_LEFT:
            self._move_view((-10, 0))
        elif ev.key == pgl.K_RIGHT:
            self._move_view((10, 0))
        elif ev.key == pgl.K_UP:
            self._move_view((0, -10))
        elif ev.key == pgl.K_DOWN:
            self._move_view((0, 10))
        elif ev.key in (pgl.K_1, pgl.K_2, pgl.K_3, pgl.K_4, pgl.K_5, pgl.K_6):
            self.change_poly(ev.key - pgl.K_0)
        elif ev.key == pgl.K_0:
            self.change_poly(None)
        elif ev.key == pgl.K_d and self.cur_poly:
            self.level.delete_point(self.cur_poly)
        elif ev.key == pgl.K_f:
            self.set_filled()
        elif ev.key == pgl.K_c:
            self.close_poly()

    def set_filled(self):
        closed, _ = self.level.all_closed()
        if closed:
            self.cur_poly = None
            self.filled_mode = True
            self._draw_lines = False
        else:
            alert('Not all polygons closed, so not filling')

    def mouse_move(self, ev):
        old_pos = self.mouse_pos
        self.mouse_pos = ev.pos
        if old_pos != self.mouse_pos and (self.cur_poly or self._draw_lines):
            self.invalidate()

    def mouse_drag(self, ev):
        if self._mouse_drag:
            old_pos = self.mouse_pos
            self.mouse_pos = ev.pos
            diff = (-self.mouse_pos[0] + old_pos[0],
                    -self.mouse_pos[1] + old_pos[1])
            self._move_view(diff)
            self.invalidate()

    def mouse_down(self, ev):
        if ev.button == 1:
            if self._draw_lines:
                if self._start_pos is None:
                    self._start_pos = ev.pos
                else:
                    self.level.add_line(self._start_pos, ev.pos)
                    self._start_pos = None
            else:
                print "Click: %r" % (
                    self.level.point_to_pymunk(
                        self._level_coordinates(ev.pos)),)
        if ev.button == 4:  # Scroll up
            self._move_view((0, -10))
        elif ev.button == 5:  # Scroll down
            self._move_view((0, 10))
        elif ev.button == 6:  # Scroll left
            self._move_view((-10, 0))
        elif ev.button == 7:  # Scroll right
            self._move_view((10, 0))
        elif self.cur_poly and ev.button == 1:
            # Add a point
            self.level.add_point(self.cur_poly,
                                 self._level_coordinates(ev.pos))
        elif ev.button == 3:
            self._mouse_drag = True

    def mouse_up(self, ev):
        if ev.button == 3:
            self._mouse_drag = False

    def close_poly(self):
        if self.cur_poly is None:
            return
        if self.level.close_poly(self.cur_poly):
            alert("Successfully closed the polygon")
            self.change_poly(None)
        else:
            alert("Failed to close the polygon")

    def _edit_class(self, classname, cls, data):
        # Dialog for class properties
        dialog = EditClassDialog(classname, cls, data)
        if dialog.present() == 'OK':
            return dialog.get_data()
        return None

    def _make_edit_dialog(self, entries):
        # Dialog to hold the editor
        edit_box = Dialog()
        edit_box.rect = pygame.rect.Rect(0, 0, 700, 500)
        table = ObjectTable(entries)
        edit_box.add(table)
        buttons = []
        for text in ['OK', 'Delete', 'Cancel']:
            but = Button(text, action=lambda x=text: edit_box.dismiss(x))
            buttons.append(but)
        row = Row(buttons)
        row.rect = pygame.rect.Rect(250, 450, 700, 50)
        edit_box.add(row)
        edit_box.get_selection = lambda: table.get_selection()
        return edit_box

    def edit_objects(self):
        edit_box = self._make_edit_dialog(self.level._game_objects)
        res = edit_box.present()
        choice = edit_box.get_selection()
        if choice is None:
            return
        if res == 'OK':
            cls = self.level.get_class(choice['classname'])
            edited = self._edit_class(choice['classname'], cls, choice)
            if edited is not None:
                if not self.level.try_new_object(self.level._game_objects,
                                                 edited, choice):
                    alert('Failed to update GameObject %s'
                          % choice['classname'])
        elif res == 'Delete':
            self.level._game_objects.remove(choice)
            self.level.reset_objs()

    def edit_enemies(self):
        edit_box = self._make_edit_dialog(self.level._enemies)
        res = edit_box.present()
        choice = edit_box.get_selection()
        if choice is None:
            return
        if res == 'OK':
            cls = self.level.get_class(choice['classname'], ne)
            edited = self._edit_class(choice['classname'], cls, choice)
            if edited is not None:
                if not self.level.try_new_object(self.level._enemies,
                                                 edited, choice):
                    alert('Failed to update Enemy %s'
                          % choice['classname'])
        elif res == 'Delete':
            self.level._enemies.remove(choice)
            self.level.reset_objs()

    def _make_choice_dialog(self, classes):
        # Dialog to hold the editor
        data = []
        for cls_name, cls in classes:
            data.append({"classname": cls_name, "class": cls})
        choice_box = Dialog()
        choice_box.rect = pygame.rect.Rect(0, 0, 700, 500)
        table = ObjectTable(data)
        choice_box.add(table)
        buttons = []
        for text in ['OK', 'Cancel']:
            but = Button(text, action=lambda x=text: choice_box.dismiss(x))
            buttons.append(but)
        row = Row(buttons)
        row.rect = pygame.rect.Rect(250, 450, 700, 50)
        choice_box.add(row)
        choice_box.get_selection = lambda: table.get_selection()
        return choice_box

    def add_game_object(self):
        classes = ngo.get_editable_game_objects()
        choose = self._make_choice_dialog(classes)
        res = choose.present()
        choice = choose.get_selection()
        if res == 'OK' and choice is not None:
            classname = choice['classname']
            cls = choice['class']
            new_cls = self._edit_class(classname, cls, None)
            if new_cls is not None:
                if not self.level.try_new_object(self.level._game_objects,
                                                 new_cls, None):
                    alert('Failed to add GameObject %s' % classname)

    def add_enemy(self):
        classes = ne.get_editable_enemies()
        choose = self._make_choice_dialog(classes)
        res = choose.present()
        choice = choose.get_selection()
        if res == 'OK' and choice is not None:
            classname = choice['classname']
            cls = choice['class']
            new_cls = self._edit_class(classname, cls, None)
            if new_cls is not None:
                if not self.level.try_new_object(self.level._enemies,
                                                 new_cls, None):
                    alert('Failed to add Enemy %s' % classname)

    def add_puzzler(self):
        classes = np.get_editable_puzzlers()
        choose = self._make_choice_dialog(classes)
        res = choose.present()
        choice = choose.get_selection()
        if res == 'OK' and choice is not None:
            classname = choice['classname']
            cls = choice['class']
            new_cls = self._edit_class(classname, cls, None)
            if new_cls is not None:
                if not self.level.try_new_object(self.level._game_objects,
                                                 new_cls, None):
                    alert('Failed to add Puzzler %s' % classname)


class PolyButton(Button):
    """Button for coosing the correct polygon"""

    def __init__(self, index, level_widget):
        if index is not None:
            text = "Draw: %s" % index
        else:
            text = 'Exit Draw Mode'
        super(PolyButton, self).__init__(text)
        self.index = index
        self.level_widget = level_widget

    def action(self):
        self.level_widget.change_poly(self.index)


class EditorApp(RootWidget):

    def __init__(self, level, surface):
        super(EditorApp, self).__init__(surface)
        self.level = level
        self.level_widget = LevelWidget(self.level)
        self.add(self.level_widget)

        self._dMenus = {}

        self._make_draw_menu()
        self._make_objects_menu()

        self._menu_mode = 'drawing'
        self._populate_menu()

    def _make_draw_menu(self):
        widgets = []

        # Add poly buttons
        y = 15
        for poly in range(1, 7):
            but = PolyButton(poly, self.level_widget)
            but.rect = pygame.rect.Rect(0, 0, MENU_WIDTH // 2 - MENU_PAD,
                                        MENU_BUTTON_HEIGHT)
            if poly % 2:
                but.rect.move_ip(MENU_LEFT, y)
            else:
                but.rect.move_ip(MENU_LEFT + MENU_WIDTH // 2 - MENU_HALF_PAD,
                                 y)
                y += MENU_BUTTON_HEIGHT + MENU_PAD
            widgets.append(but)

        end_poly_but = PolyButton(None, self.level_widget)
        end_poly_but.rect = BUTTON_RECT.copy()
        end_poly_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(end_poly_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        draw_line = Button("Draw interior wall", self.level_widget.line_mode)
        draw_line.rect = BUTTON_RECT.copy()
        draw_line.rect.move_ip(MENU_LEFT, y)
        widgets.append(draw_line)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        fill_but = Button('Fill exterior', action=self.level_widget.set_filled)
        fill_but.rect = BUTTON_RECT.copy()
        fill_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(fill_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        close_poly_but = Button('Close Polygon',
                                action=self.level_widget.close_poly)
        close_poly_but.rect = BUTTON_RECT.copy()
        close_poly_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(close_poly_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        white = pygame.color.Color("white")
        self.show_objs = CheckBox(fg_color=white)
        self.show_objs.rect = CHECK_RECT.copy()
        self.show_objs.rect.move_ip(MENU_LEFT, y)
        label = Label("Show Objects", fg_color=white)
        label.rect.move_ip(MENU_LEFT + MENU_BUTTON_HEIGHT // 2 + MENU_PAD, y)
        widgets.append(self.show_objs)
        widgets.append(label)
        y += label.rect.height + MENU_PAD

        self.show_enemies = CheckBox(fg_color=white)
        self.show_enemies.rect = CHECK_RECT.copy()
        self.show_enemies.rect.move_ip(MENU_LEFT, y)
        label = Label("Show enemy start pos", fg_color=white)
        label.rect.move_ip(MENU_LEFT + MENU_BUTTON_HEIGHT // 2 + MENU_PAD, y)
        widgets.append(self.show_enemies)
        widgets.append(label)
        y += label.rect.height + MENU_PAD

        y += MENU_PAD
        switch_but = Button('Switch to Objects', action=self.switch_to_objects)
        switch_but.rect = BUTTON_RECT.copy()
        switch_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(switch_but)
        y += switch_but.rect.height + MENU_PAD

        save_but = Button('Save Level', action=self.save)
        save_but.rect = BUTTON_RECT.copy()
        save_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(save_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        y += MENU_PAD
        quit_but = Button('Quit', action=self.quit)
        quit_but.rect = BUTTON_RECT.copy()
        quit_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(quit_but)

        self._dMenus['drawing'] = widgets

    def _make_objects_menu(self):
        widgets = []

        # Add poly buttons
        y = 15

        edit_objs_but = Button('Edit Objects',
                               action=self.level_widget.edit_objects)
        edit_objs_but.rect = BUTTON_RECT.copy()
        edit_objs_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(edit_objs_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        edir_enemies_but = Button('Edit Enemies',
                                  action=self.level_widget.edit_enemies)
        edir_enemies_but.rect = BUTTON_RECT.copy()
        edir_enemies_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(edir_enemies_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        add_obj_but = Button('Add Game Object',
                             action=self.level_widget.add_game_object)
        add_obj_but.rect = BUTTON_RECT.copy()
        add_obj_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(add_obj_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        add_puzzle_but = Button('Add Puzzler',
                                action=self.level_widget.add_puzzler)
        add_puzzle_but.rect = BUTTON_RECT.copy()
        add_puzzle_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(add_puzzle_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        add_enemy_but = Button('Add Enemy',
                               action=self.level_widget.add_enemy)
        add_enemy_but.rect = BUTTON_RECT.copy()
        add_enemy_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(add_enemy_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        y += MENU_PAD
        switch_but = Button('Switch to Drawing', action=self.switch_to_draw)
        switch_but.rect = BUTTON_RECT.copy()
        switch_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(switch_but)
        y += switch_but.rect.height + MENU_PAD

        save_but = Button('Save Level', action=self.save)
        save_but.rect = BUTTON_RECT.copy()
        save_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(save_but)
        y += MENU_BUTTON_HEIGHT + MENU_PAD

        y += MENU_PAD
        quit_but = Button('Quit', action=self.quit)
        quit_but.rect = BUTTON_RECT.copy()
        quit_but.rect.move_ip(MENU_LEFT, y)
        widgets.append(quit_but)

        self._dMenus['objects'] = widgets

    def key_down(self, ev):
        if ev.key == pgl.K_ESCAPE:
            self.quit()
        elif ev.key == pgl.K_s:
            self.save()
        else:
            self.level_widget.key_down(ev)

    def save(self):
        closed, messages = self.level.all_closed()
        if closed:
            self.level.save()
            # display success
            alert("Level %s saved successfully." % self.level.name)
        else:
            # display errors
            alert("Failed to save level.\n\n%s" % '\n'.join(messages))

    def switch_to_draw(self):
        if self._menu_mode != 'drawing':
            self._clear_menu()
            self._menu_mode = 'drawing'
            self._populate_menu()

    def switch_to_objects(self):
        if self._menu_mode != 'objects':
            self._clear_menu()
            self._menu_mode = 'objects'
            self._populate_menu()

    def _clear_menu(self):
        for widget in self._dMenus[self._menu_mode]:
            self.remove(widget)

    def _populate_menu(self):
        self.level_widget.change_poly(None)
        for widget in self._dMenus[self._menu_mode]:
            self.add(widget)
        self.invalidate()

    def mouse_move(self, ev):
        self.level_widget.mouse_move(ev)

    def draw(self, surface):
        # Update checkbox state
        if self._menu_mode == 'drawing':
            self.level_widget.set_objects(self.show_objs.value)
            self.level_widget.set_enemies(self.show_enemies.value)
        else:
            self.level_widget.set_objects(True)
            self.level_widget.set_enemies(True)
        super(EditorApp, self).draw(surface)


if __name__ == "__main__":
    if len(sys.argv) not in [2, 4]:
        print 'Please supply a levelname or levelname and level size'
        sys.exit()
    # Need to ensure we have defaults for rendering
    parse_args([])
    pygame.display.init()
    pygame.font.init()
    pygame.display.set_mode((SCREEN[0] + MENU_WIDTH, SCREEN[1]),
                            pgl.SWSURFACE)
    if len(sys.argv) == 2:
        level = EditorLevel(sys.argv[1])
        level.load(pymunk.Space())
    elif len(sys.argv) == 4:
        level = EditorLevel(sys.argv[1], int(sys.argv[2]), int(sys.argv[3]))
    pygame.display.set_caption('Nagslang Area Editor')
    pygame.key.set_repeat(200, 100)
    app = EditorApp(level, pygame.display.get_surface())
    app.run()