import time
import pygame
from .obj.Object import Object
from .obj.Waiter import Waiter
from .UserController import UserController
from .StateController import StateController
from .decisionTree.TreeConcept import TreeEngine
from queue import PriorityQueue


class Engine:

    def __init__(self, screen_size, square_size, user: UserController, state: StateController):
        pygame.display.set_caption('Waiter Agent')

        self.action_clock = 0
        self.tree = TreeEngine()

        self.user: Waiter = user
        self.state: StateController = state
        self.screen_size: list[int] = screen_size
        self.screen = pygame.display.set_mode(self.screen_size)

        self.square_size: int = square_size
        self.num_squares = self.screen_size[0] // self.square_size
        self.squares = self.__init_squares_field__(
            self.num_squares, self.square_size)

        self.objects: list[Object] = []
        self.goals: list = []

        self.runnin: bool = False

    def __init_squares_field__(self, num_squares, square_size):
        squares = []
        for i in range(num_squares):
            row = []
            for j in range(num_squares):
                square_rect = pygame.Rect(
                    j * square_size, i * square_size,
                    square_size, square_size)
                row.append(square_rect)
            squares.append(row)

        return squares

    def subscribe(self, object: Object):
        self.objects.append(object)

    def loop(self):
        self.running = True
        while self.running:

            self.action()
            self.redraw()

    def quit(self):
        self.running = False

    def action(self):
        if not self.state.path:
            if self.goals:
                self.state.graphsearch(self)
            self.user.handler(self)
            self.predict()
        else:
            # went path

            state = self.user.obj.changeState(self.state.path.pop())
            self.clock_increment(state.cost)

            print("Action:\t{0}\tCost:\t{1}\tCost so far: {2}\tBattery: {3}".format(
                state.agent_role,
                state.cost,
                state.cost_so_far,
                self.user.obj.battery)
            )

            # waiter interaction

            for o in self.objects:
                if o.compare_pos(self.user.obj.position):
                    o.action(self.user.obj)

            time.sleep(0.5)

    def clock_increment(self, action_time):
        self.action_clock += action_time

    def redraw(self):
        self.screen.fill((255, 255, 255))

        for row in self.squares:
            for square_rect in row:
                pygame.draw.rect(self.screen, (0, 0, 0), square_rect, 1)

        for o in self.objects:
            o.blit(self.screen)

        self.user.obj.blit(self.screen)

        for f in self.state.fringe.queue:
            f.blit(self.screen)

        for s in self.state.path:
            s.blit(self.screen)

        pygame.display.flip()

    def appendGoalPosition(self, position):
        self.goals.append(position)

    def predict(self):

        goal_queue = PriorityQueue()

        for o in self.objects:

            condition = o.agent_role in [
                "order",
                "wait",
                "done"
            ]

            if not condition or o.compare_pos(self.user.obj.position):
                continue

            battery = self.user.obj.battary_status()[0]
            distance = o.distance_to(
                self.user.obj.position) // (self.num_squares // 2)
            mood = o.get_mood(self.action_clock)[0] if condition else 0
            memory = self.user.obj.memory_size
            dishes_held = self.user.obj.basket_size
            customers = o.customers if condition else 0
            waiting_for_order = o.is_order() if condition else 0
            waiting_for_dish = o.is_done() if condition else 0

            p = self.tree.make_predict(
                battery,
                distance,
                mood,
                memory,
                dishes_held,
                customers,
                waiting_for_order,
                waiting_for_dish
            )

            goal_queue.put((p, o.position))

        if len(goal_queue.queue):
            priority, goal = goal_queue.queue.pop()
            if priority:
                print(goal, priority, end='\r')
                self.appendGoalPosition(goal)
            else:
                print(goal, priority, end='\r')