"""Habitat for editing levels."""

import pygame.display
from pygame.locals import (SWSURFACE, KEYDOWN, K_1, K_2, MOUSEBUTTONDOWN,
        K_LEFT, K_RIGHT, K_DOWN, K_UP, K_SPACE)
import sys
import traceback
from StringIO import StringIO
import urllib
import urllib2

from mamba.engine import (Habitat, NewHabitatEvent, SnakeDiedEvent,
        LevelCompletedEvent)
from mamba.sprites import find_special_sprites
from mamba.widgets.level import EditLevelWidget
from mamba.widgets.text import TextWidget, TextButton
from mamba.widgets.imagebutton import ImageButtonWidget
from mamba.widgets.messagebox import MessageBox
from mamba.widgets.entrybox import EntryBox
from mamba.widgets.listbox import ListBox
from mamba.widgets.toollist import ToolListWidget
from mamba.widgets.editsprite import EditSpriteBox
from mamba.widgets.editlevel import EditLevelBox
from mamba.level import Level, Tileset, TILE_MAP, THING_MAP, InvalidMapError
from mamba.data import (check_level_exists, get_level_list, load_file,
                        load_image, load_tile_image)
from mamba.constants import (SCREEN, EDIT_SCREEN, NAME, ESCAPE_KEYS,
        RESERVED_NAMES, WINDOW_ICON, LEVEL_SERVER, UP, DOWN, LEFT, RIGHT)

MAX_TOOLS = 6
MODE_HEIGHT = 370
LOAD_SAVE_HEIGHT = 500
HELP = 'editor_help.txt'


class EditorHabitat(Habitat):
    def __init__(self, level):
        super(EditorHabitat, self).__init__(EDIT_SCREEN)
        self.toolbar = {}
        self.level = level
        self.container.paused = False
        self.edit_widget = EditLevelWidget(self.level)
        self.container.add(self.edit_widget)
        self.container.add_callback(KEYDOWN, self.keydown_event)
        self.container.add_callback(MOUSEBUTTONDOWN, self.mouse_event)
        self.edit_widget.add_callback(SnakeDiedEvent, self.halt_test)
        self.edit_widget.add_callback(LevelCompletedEvent, self.halt_test)
        self.mode = 'Tiles'
        self.sprite_mode = 'Add'
        self.sprite_cls = None
        # Map for test mode lookups
        self.replays = {}
        self.action_map = {
                K_UP: UP,
                K_DOWN: DOWN,
                K_LEFT: LEFT,
                K_RIGHT: RIGHT,
                K_SPACE: None,
                }
        helpfile = load_file(HELP)
        self.help_msg = ''.join(helpfile.readlines())

    def on_enter(self):
        # We need to juggle the display to the correct size
        # This is a horrible hack
        pygame.display.quit()
        pygame.display.init()
        pygame.display.set_mode(EDIT_SCREEN, SWSURFACE)
        pygame.display.set_icon(load_image(WINDOW_ICON))
        pygame.display.set_caption('%s Level editor' % NAME)
        super(EditorHabitat, self).on_enter()
        self.setup_toolbar()

    def on_exit(self):
        # We need to juggle the display to the correct size
        # This is a horrible hack
        super(EditorHabitat, self).on_exit()
        pygame.display.quit()
        pygame.display.init()
        pygame.display.set_mode(SCREEN, SWSURFACE)
        pygame.display.set_icon(load_image(WINDOW_ICON))
        pygame.display.set_caption(NAME)

    def keydown_event(self, ev, widget):
        if ev.key in ESCAPE_KEYS:
            from mamba.habitats.mainmenu import MainMenu
            NewHabitatEvent.post(MainMenu())
            return True
        elif ev.key == K_1:
            # Activate floor button
            self.floor_button.forced_click()
            return True
        elif self.mode == 'Test' and ev.key in self.action_map:
            self.edit_widget.apply_action(self.action_map[ev.key])
            return True

    def halt_test(self, ev, widget):
        if SnakeDiedEvent.matches(ev):
            text = 'Snake died: %s' % ev.reason
        else:
            text = 'Level completed'
        if self.edit_widget.snake_alive:
            self.edit_widget.kill_snake()
            message = MessageBox((300, 200), text,
                    post_callback=self.test_restart)
            self.display_dialog(message)
        return True

    def test_restart(self):
        self.unpause()
        self.edit_widget.restart()
        return True

    def setup_toolbar(self):
        """Draw the editor toolbar"""
        button_height = 5
        button_left = 820
        button_padding = 2

        level_name = TextWidget((button_left, button_height),
                                'Level: %s' % self.level.level_name)
        self.container.add(level_name)
        button_height += level_name.rect.height + button_padding

        edit_level = TextButton((button_left, button_height),
                                'Edit Level Metadata')
        edit_level.add_callback('clicked', self.edit_level_data)
        self.container.add(edit_level)
        button_height += edit_level.rect.height + button_padding

        # TODO: Add Image widget for the current tool
        if self.mode != 'Sprites':
            self.current_tool = TextWidget((button_left, button_height),
                    'Tool: Floor', color='white')
        else:
            self.current_tool = TextWidget((button_left, button_height),
                    '%s Sprite' % self.sprite_mode, color='white')
        self.container.add(self.current_tool)
        button_height += self.current_tool.surface.get_height()
        button_height += button_padding
        if self.mode != 'Sprites':
            self.floor_button = ImageButtonWidget(
                    (button_left, button_height), self.level.tileset.floor,
                    'Floor', color='white')
            self.container.add(self.floor_button)
            self.floor_button.add_callback('clicked', self.change_tool,
                    '.', 'Floor')
            self.edit_widget.set_tool('.')
            button_height += (self.floor_button.surface.get_height()
                    + button_padding)
        if self.mode in self.toolbar:
            # FIXME: This needs to be recreated on tileset changes
            self.tool_widget = self.toolbar[self.mode]
        else:
            if self.mode == 'Tiles':
                tile_map = TILE_MAP
            elif self.mode == 'Things':
                tile_map = THING_MAP
            else:
                tile_map = []
            tool_list = []
            for tile_char in sorted(tile_map):
                try:
                    tile = self.level.tileset[tile_char]
                except pygame.error:
                    # Ignore stuff we can't load for now
                    continue
                if tile is None:
                    continue
                if tile.name:
                    text = tile.name
                else:
                    text = 'Tile'
                tile_button = ImageButtonWidget((0, 0), tile.image, text,
                        color='white')
                tile_button.add_callback('clicked', self.change_tool,
                        tile_char, text)
                tool_list.append(tile_button)
            if self.mode == "Sprites":
                for cls_name, sprite_cls in find_special_sprites():
                    image = load_tile_image(sprite_cls.image_name,
                            self.level.tileset.name)
                    name = sprite_cls.name
                    tile_button = ImageButtonWidget((0, 0), image, name,
                            color='white')
                    tile_button.add_callback('clicked', self.sprite_tool,
                            'Add', cls_name, sprite_cls)
                    tool_list.append(tile_button)
                for name in ['Edit', 'Delete']:
                    tile_button = TextButton((0, 0), '%s Sprite' % name)
                    tile_button.add_callback('clicked', self.sprite_tool,
                            name, None, None)
                    tool_list.append(tile_button)
            elif self.mode == 'Test':
                for replay in [1, 2]:
                    store_button = TextButton((0, 0),
                            'Store run as %s' % replay)
                    store_button.add_callback('clicked', self.store_replay,
                            replay)
                    tool_list.append(store_button)
                    replay_button = TextButton((0, 0),
                            'Replay run %s' % replay)
                    replay_button.add_callback('clicked', self.do_replay,
                            replay)
                    tool_list.append(replay_button)
                last_button = TextButton((0, 0),
                        'Replay last run')
                last_button.add_callback('clicked', self.do_replay, None)
                tool_list.append(last_button)
            self.tool_widget = ToolListWidget((button_left, button_height),
                    tool_list, MAX_TOOLS, start_key=K_2)
            self.toolbar[self.mode] = self.tool_widget
        self.container.add(self.tool_widget)

        button_height = self.container.rect.top + MODE_HEIGHT
        tile_button = TextButton((button_left, button_height),
                                 'Tiles')
        tile_button.add_callback('clicked', self.change_toolbar, 'Tiles')
        thing_button = TextButton((button_left + tile_button.rect.width +
                                   button_padding, button_height),
                                  'Things')
        thing_button.add_callback('clicked', self.change_toolbar, 'Things')
        sprite_button = TextButton((button_left,
                                    thing_button.rect.bottom +
                                    button_padding),
                                   'Sprites')
        sprite_button.add_callback('clicked', self.change_toolbar,
                                   'Sprites')
        test_button = TextButton((sprite_button.rect.right + button_padding,
                                  thing_button.rect.bottom + button_padding),
                                  'Test')
        test_button.add_callback('clicked', self.change_toolbar, 'Test')
        help_button = TextButton((button_left,
                                  test_button.rect.bottom + button_padding),
                                 'Help')
        help_button.add_callback('clicked', self.show_help)

        if self.mode == "Tiles":
            tile_button.disable()
        elif self.mode == "Things":
            thing_button.disable()
        elif self.mode == "Sprites":
            sprite_button.disable()
        elif self.mode == 'Test':
            test_button.disable()
        self.container.add(tile_button)
        self.container.add(thing_button)
        self.container.add(sprite_button)
        self.container.add(test_button)
        self.container.add(help_button)

        button_height = LOAD_SAVE_HEIGHT

        upload = TextButton((button_left, button_height), "Upload Level")
        upload.add_callback('clicked', self.upload)
        self.container.add(upload)

        mainmenu = TextButton((upload.rect.right + button_padding,
                               button_height), 'Exit')
        mainmenu.add_callback('clicked', self.go_mainmenu)
        self.container.add(mainmenu)

        button_height = upload.rect.bottom + button_padding

        new = TextButton((button_left, button_height), "New")
        new.add_callback('clicked', self.new)
        self.container.add(new)

        load = TextButton((new.rect.right + button_padding, button_height),
                          "Load")
        load.add_callback('clicked', self.load)
        self.container.add(load)

        save = TextButton((load.rect.right + button_padding, button_height),
                          "Save")
        save.add_callback('clicked', self.save)
        self.container.add(save)

        if self.mode == 'Test':
            save.disable()
            load.disable()
            new.disable()

    def change_tool(self, ev, widget, new_tool, text):
        self.edit_widget.set_tool(new_tool)
        self.current_tool.text = 'Tool: %s' % text
        self.current_tool.prepare()
        return True

    def show_help(self, ev, widget):
        message = MessageBox((20, 20), self.help_msg, color="black",
                             fontsize=12)
        self.display_dialog(message)
        return True

    def go_mainmenu(self, ev, widget):
        from mamba.habitats.mainmenu import MainMenu
        NewHabitatEvent.post(MainMenu())
        return True

    def check_level(self):
        message = None
        if not self.level.level_name:
            message = MessageBox((300, 300), 'Please enter a name')
        elif self.level.level_name in RESERVED_NAMES:
            message = MessageBox((300, 300), 'Reserved level name')
        elif '/' in self.level.level_name:
            message = MessageBox((300, 300), 'Illegal level name')

        if message is None:
            try:
                self.level.validate_level()
            except InvalidMapError, error:
                message = MessageBox((300, 300), "Map isn't valid\n%s" % error)

        return message

    def save(self, ev, widget):
        message = self.check_level()
        if message:
            self.display_dialog(message)
            return True
        self.level.save_level('user_levels', is_user_dir=True)
        self.refresh_display()
        message = MessageBox((300, 300),
                'Success!\nYou have saved a user level')
        self.display_dialog(message)
        return True

    def new(self, ev, widget):
        return self.load(ev, widget, 'levels', subdir='templates')

    def upload(self, ev, widget):
        message = self.check_level()
        if message:
            self.display_dialog(message)
            return True
        save_file = StringIO()
        self.level.save_level(save_file=save_file)

        url = "%s/save/%s" % (LEVEL_SERVER, self.level.level_name)
        args = urllib.urlencode([('data', save_file.getvalue())])
        try:
            result = urllib2.urlopen(url, args)
            mtxt = result.read()
        except:
            mtxt = "Failed to upload level. :("
        else:
            mtxt = "Success! %s\n%s" % (mtxt,
                   "Your level is now awaiting curation.")

        message = MessageBox((300, 300), mtxt)
        self.display_dialog(message)
        return True

    def load(self, ev, widget, level_dir=None, is_user_dir=False, subdir=''):
        if level_dir is None:
            level_dir = 'user_levels'
            is_user_dir = True
        self.container.paused = True
        levels = get_level_list('/'.join([level_dir, subdir]), is_user_dir)
        load_list = []
        for level_name in levels:
            if level_name in RESERVED_NAMES:
                continue
            if subdir:
                level_name = '/'.join([subdir, level_name])
            load_button = TextButton((0, 0), level_name)
            load_button.add_callback(
                'clicked', self.load_level, level_name, level_dir, is_user_dir)
            load_list.append(load_button)
        self.display_dialog(ListBox((200, 200), 'Select Level', load_list, 6))
        return True

    def load_level(self, ev, widget, level_name, level_dir, is_user_dir):
        try:
            source = load_file("%s/%s.txt" % (level_dir, level_name),
                               is_user_dir=is_user_dir)
            new_level = Level(level_name, 'user', source.read())
        except (IOError, InvalidMapError, pygame.error), error:
            message = MessageBox((300, 300),
                    'Loading Level Failed: %s' % error, color='red')
            self.display_dialog(message)
            return False
        self.container.remove(self.edit_widget)
        self.level = new_level
        if level_name in RESERVED_NAMES:
            self.level.level_name = ''  # Special case for new level
        self.container.paused = False
        self.edit_widget = EditLevelWidget(self.level)
        self.edit_widget.add_callback(SnakeDiedEvent, self.halt_test)
        self.edit_widget.add_callback(LevelCompletedEvent, self.halt_test)
        self.replays = {}
        self.container.add(self.edit_widget)
        self.clear_toolbar()
        self.setup_toolbar()
        return True

    def change_toolbar(self, ev, widget, new_mode):
        old_mode = self.mode
        if new_mode == 'Test':
            try:
                self.level.validate_level()
            except InvalidMapError, error:
                # Fail to change mode on invalid maps
                message = MessageBox((300, 300), "Map isn't valid\n%s" % error)
                self.display_dialog(message)
                return True
        self.mode = new_mode
        self.edit_widget.tile_mode = (self.mode in ('Tiles', 'Things'))
        self.clear_toolbar()
        self.setup_toolbar()
        if self.mode == 'Test':
            self.edit_widget.start_test()
        elif old_mode == 'Test':
            self.edit_widget.stop_test()
        return True

    def clear_toolbar(self):
        """Remove every non-edit widget from the container"""
        for widget in self.container.children[:]:
            if widget is not self.edit_widget:
                self.container.remove(widget)

    def do_edit(self, ev, widget, message, init_value, callback):
        self.display_dialog(
            EntryBox((200, 200), message, init_value, callback))
        return True

    def unpause(self):
        self.container.paused = False

    def update_level_data(self, filename, name, author, tileset, track):

        self.level.name = name
        self.level.author = author
        self.level.background_track = track
        # err_ts = self.change_tileset(tileset)
        err_fn = self.check_file(filename)

        self.clear_toolbar()
        self.setup_toolbar()

        if err_fn:
            err_fn.post_callback = self.unpause
            self.display_dialog(err_fn)
        else:
            self.container.paused = False

    def refresh_display(self):
        self.level.restart()
        self.clear_toolbar()
        self.setup_toolbar()
        return True

    def change_tileset(self, new_tileset):
        print "changing tileset"
        self.level.update_tiles_ascii()
        old_tileset = self.level.tileset
        try:
            self.level.tileset = Tileset(new_tileset)
            self.level.restart()
            return None
        except pygame.error, error:
            self.level.tileset = old_tileset
            return MessageBox(
                (300, 300), 'Unable to change tileset:: %s' % error,
                color='red')

    def check_file(self, new_name):
        message = None
        if new_name == self.level.level_name:
            return None  # No-op change
        if not new_name:
            message = MessageBox((300, 300), 'Please enter a name')
        if new_name in RESERVED_NAMES:
            # This case is caught by the existance check, but the
            # importance of the reserved names means we use a different
            # message
            message = MessageBox((300, 300), 'Reserved level name')
        elif check_level_exists(new_name, 'user_levels', is_user_dir=True):
            message = MessageBox((300, 300), 'Name already in use')
        if message:
            return message
        self.level.level_name = new_name
        return None

    def sprite_tool(self, ev, widget, sprite_mode, cls_name, sprite_cls):
        """Handle sprite stuff"""
        self.sprite_mode = sprite_mode
        self.sprite_cls_name = cls_name
        self.sprite_cls = sprite_cls
        self.level.update_tiles_ascii()  # commit any changes
        self.clear_toolbar()
        self.setup_toolbar()
        return True

    def mouse_event(self, ev, widget):
        """Handle mouse clicks when we are in sprite mode"""
        if self.mode != 'Sprites':
            return False
        if self.container.paused:
            return False
        tile_pos = self.edit_widget.convert_pos(ev.pos)
        sprite, sprite_ascii = self.level.get_sprite_at(tile_pos)
        if self.sprite_mode == 'Delete' and sprite:
            self.level.remove_sprite(sprite_ascii)
            self.level.restart()
        elif self.sprite_mode == 'Edit' and sprite:
            self.edit_sprite(tile_pos, sprite)
        elif (self.sprite_mode == 'Add' and sprite is None
              and self.sprite_cls is not None):
            self.edit_sprite(tile_pos,
                    (self.sprite_cls_name, self.sprite_cls, None, []))
        return True

    def edit_sprite(self, tile_pos, sprite_info):
        sprite_editor = EditSpriteBox((200, 100), tile_pos, sprite_info,
                post_callback=self.commit_line)
        self.display_dialog(sprite_editor)

    def edit_level_data(self, ev, widget):
        elb = EditLevelBox((200, 100), self.level, self.update_level_data)
        self.display_dialog(elb)

    def commit_line(self, sprite):
        try:
            self.level.validate_sprite(sprite)
        except:
            # We don't know what errors thwe constructor may show, so
            # we catch everything
            # We use sys.exc_info to get slight neater info
            exc_type, info, _ = sys.exc_info()
            info = traceback.format_exception_only(exc_type, info)[0]
            message = MessageBox((300, 300),
                    'Validation failed:\n%s' % info)
            self.display_dialog(message)
            return False
        # Validation successful, so add to level
        if self.sprite_mode == 'Add':
            self.level.add_sprite(sprite)
        elif self.sprite_mode == 'Edit':
            self.level.replace_sprite(sprite)
        self.level.restart()
        return True

    def store_replay(self, ev, widget, number):
        self.replays[number] = self.edit_widget.get_replay()

    def do_replay(self, ev, widget, number):
        if number:
            # We do nothing if the user tries to replay something thay
            # isn't there
            if number in self.replays:
                # Replay stored replay
                self.edit_widget.replay(self.replays[number])
        else:
            # Replaying last run
            self.edit_widget.replay()