import pygame

# Simple Infinite State Machine
class State:
    def __init__(self, control):
        self.control = control

    def on_draw(self, surface): pass
    def on_event(self, event): pass
    def on_update(self, delta): pass

class StateHandler:
    def __init__(self):
        self.states = {}
        self.current = None
        self.next_ = None

    def add(self, name, state):
        self.states[name] = state

    def get(self, name):
        return self.states[name]

    def push(self, state):
        self.next_ = self.states[state]

    def update(self):
        if self.next_:
            self.current = self.next_
            self.next_ = None

class StateControl:
    def __init__(self, caption, size, flags=0, fps=0):
        pygame.display.set_caption(caption)
        self.display = pygame.display.set_mode(size, flags)
        self.rect = self.display.get_rect()
        self.clock = pygame.time.Clock()
        self.running = True
        self.delta = 0
        self.fps = fps

        self.state = StateHandler()

    def loop(self):
        while self.running:
            self.state.update()
            for event in pygame.event.get():
                if event.type != pygame.QUIT:
                    self.state.current.on_event(event)
                else:
                    self.running = False

            self.state.current.on_update(self.delta)
            self.state.current.on_draw(self.display)
            pygame.display.flip()
            self.delta = self.clock.tick(self.fps)

    def run(self, state=None):
        if state:
            self.state.current = state
        else:
            if self.state.current is None:
                self.state.current = State(self)
                print("You forgot to set state")

        self.loop()

# States
class Intro(State):
    def __init__(self, control):
        super().__init__(control)
        font = pygame.font.Font(None, 30)
        self.image = font.render("Game Intro", 1, "dodgerblue")
        self.rect = self.image.get_rect()
        self.rect.top = self.control.rect.top + 20
        self.rect.centerx = self.control.rect.centerx

    def on_draw(self, surface):
        surface.fill("black")
        surface.blit(self.image, self.rect)

    def on_event(self, event):
        if event.type == pygame.KEYDOWN:
            self.control.state.push("Game")

class Game(State):
    def __init__(self, control):
        super().__init__(control)
        self.images = self.create_images()
        self.positions = (40, 30), (500, 200), (150, 320)

    def create_images(self):
        box = pygame.Surface((100, 100))
        box.fill("blue")

        rect = pygame.Rect(0, 0, 50, 50)
        circle = pygame.Surface(rect.size, pygame.SRCALPHA)
        pygame.draw.circle(circle, "lawngreen", rect.center, rect.width / 2 - 1)

        rect = pygame.Rect(0, 0, 60, 60)
        triangle = pygame.Surface(rect.size, pygame.SRCALPHA)
        points = rect.midtop, rect.bottomright, rect.bottomleft
        pygame.draw.polygon(triangle, "firebrick", points)

        return box, circle, triangle

    def on_draw(self, surface):
        surface.fill("black")
        for (image, pos) in zip(self.images, self.positions):
            surface.blit(image, pos)

    def on_event(self, event):
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                self.control.state.push("Pause")

class Pause(State):
    def __init__(self, control):
        super().__init__(control)
        self.image = pygame.Surface(self.control.rect.size, pygame.SRCALPHA)
        self.fill_color = pygame.Color(250, 250, 0, 100)
        self.font = pygame.font.Font(None, 24)
        self.alpha_sprite = QuickSprite()
        self.update_alpha()

    def on_draw(self, surface):
        game = self.control.state.get("Game")
        game.on_draw(surface)
        surface.blit(self.image, (0, 0))
        self.alpha_sprite.draw(surface)

    def on_event(self, event):
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP:
                if self.fill_color.a < 250:
                    self.fill_color.a += 10
                    self.update_alpha()
            elif event.key == pygame.K_DOWN:
                if self.fill_color.a > 10:
                    self.fill_color.a -= 10
                    self.update_alpha()
            else:
                self.control.state.push("Game")

    def update_alpha(self):
        self.image.fill(self.fill_color)
        image = self.font.render(str(self.fill_color.a), 1, 'navy')
        self.alpha_sprite.set_image(image)
        self.alpha_sprite.rect.top = 20
        self.alpha_sprite.rect.right = self.control.rect.right - 20

# Other
class QuickSprite:
    def __init__(self, image=None):
        if image:
            self.set_image(image)

    def draw(self, surface):
        surface.blit(self.image, self.rect)

    def set_image(self, image):
        self.image = image
        self.rect = image.get_rect()

def main():
    control = StateControl("Example", (800, 600))
    control.state.add("Game", Game(control))
    control.state.add("Pause", Pause(control))
    control.run(Intro(control))

if __name__ == "__main__":
    pygame.init()
    main()
    pygame.quit()