import collections

import pygame
from pygame.constants import K_UP, K_DOWN, K_RETURN
from pygame.locals import MOUSEMOTION, MOUSEBUTTONUP, MOUSEBUTTONDOWN, KEYDOWN


class Widget(object):

    def __init__(self, rect):
        if not isinstance(rect, pygame.Rect):
            rect = pygame.Rect(rect, (0, 0))
        self.rect = rect
        self.focussable = False
        self.focussed = False
        self.parent = None
        self.callbacks = collections.defaultdict(list)

    def add_callback(self, eventtype, callback, *args):
        self.callbacks[eventtype].append((callback, args))

    def event(self, ev):
        for callback, args in self.callbacks[ev.type]:
            if callback(ev, self, *args):
                return True
        return False

    def draw(self, surface):
        "Override me"
        pass


class Button(Widget):

    def event(self, ev):
        if super(Button, self).event(ev):
            return True
        if (ev.type == MOUSEBUTTONDOWN
                or (ev.type == KEYDOWN and ev.key == K_RETURN)):
            for callback, args in self.callbacks['clicked']:
                if callback(ev, self, *args):
                    return True
            return False


class Container(Widget):

    def __init__(self, rect):
        super(Container, self).__init__(rect)
        self.children = []
        self.focussed_child = None

    def event(self, ev):
        """Push an event down through the tree, and fire our own event as a
        last resort
        """
        if ev.type in (MOUSEMOTION, MOUSEBUTTONUP, MOUSEBUTTONDOWN):
            for child in self.children:
                if child.rect.collidepoint(ev.pos):
                    if ev.type == MOUSEBUTTONDOWN and child.focussable:
                        root = self
                        while root.parent is not None:
                            root = root.parent
                        root.defocus()
                        widget = child
                        while widget.parent is not None:
                            parent = widget.parent
                            if isinstance(parent, Container):
                                idx = parent.children.index(widget)
                                parent.focussed_child = idx
                            widget = parent
                        child.focussed = True

                    if child.event(ev):
                        return True

        if ev.type == KEYDOWN:
            for child in self.children:
                if child.focussed:
                    if child.event(ev):
                        return True
        if super(Container, self).event(ev):
            return True
        if (self.parent is None and ev.type == KEYDOWN
                and ev.key in (K_UP, K_DOWN)):
            return self.adjust_focus(1 if ev.key == K_DOWN else -1)

    def add(self, widget):
        widget.parent = self
        self.children.append(widget)
        self.rect = self.rect.union(widget.rect)

    def defocus(self):
        if self.focussed_child is not None:
            child = self.children[self.focussed_child]
            if isinstance(child, Container):
                child.defocus()
            child.focussed = False

    def adjust_focus(self, direction):
        """Try and adjust focus in direction (integer)
        """
        if self.focussed_child is not None:
            child = self.children[self.focussed_child]
            if isinstance(child, Container):
                if child.adjust_focus(direction):
                    return True
            else:
                child.focussed = False

        current = self.focussed_child
        if current is None:
            current = -1 if direction > 0 else len(self.children)
        if direction > 0:
            possibles = list(enumerate(self.children))[current + 1:]
        else:
            possibles = list(enumerate(self.children))[:current]
            possibles.reverse()
        for i, child in possibles:
            if child.focussable:
                child.focussed = True
                self.focussed_child = i
                return True
            if isinstance(child, Container):
                if child.adjust_focus(direction):
                    self.focussed_child = i
                    return True
        else:
            if self.parent is None:
                if self.focussed_child is not None:
                    # At the end, mark the last one as focussed, again
                    child = self.children[self.focussed_child]
                    if isinstance(child, Container):
                        if child.adjust_focus(-direction):
                            return True
                    else:
                        child.focussed = True
                        return True
            else:
                self.focussed_child = None
            return False

    def draw(self, surface):
        if self.parent is None and not self.focussed:
            self.focussed = True
            self.adjust_focus(1)
        for child in self.children:
            child.draw(surface)