"""Habitat for editing levels."""

import pygame.display
from pygame.locals import SWSURFACE

from mamba.engine import Habitat
from mamba.widgets.level import LevelWidget
from mamba.level import Level
from mamba.constants import SCREEN, EDIT_SCREEN, NAME


class EditorHabitat(Habitat):
    def __init__(self, level_name):
        super(EditorHabitat, self).__init__()
        self.level = Level(level_name)
        self.container.add(LevelWidget(self.level))

    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_caption('%s Level editor' % NAME)

    def on_exit(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(SCREEN, SWSURFACE)
        pygame.display.set_caption(NAME)