import time
import pygame
from .obj.Goal import Goal
from .obj.Object import Object
from .obj.Waiter import Waiter
from .obj.Kitchen import Kitchen
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, kitchen: Kitchen, user: UserController, state: StateController):
        pygame.display.set_caption('Waiter Agent')

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

        self.kitchen: Kitchen = kitchen
        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 self.user.obj.chechNeighbor(o):
                    o.updateState(self.action_clock)
                if o.compare_pos(self.user.obj.position):
                    o.action(self.user.obj, self.action_clock)

            if self.kitchen.compare_pos(self.user.obj.position):
                self.kitchen.action(self.user.obj, self.action_clock)

            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.kitchen.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)

        if self.goals:
            self.goals[-1].blit(self.screen)

        pygame.display.flip()

    def appendGoalPosition(self, position):
        self.goals.append(Goal(position, self.square_size, self.screen_size))

    def predict(self):

        goal_queue = PriorityQueue()

        for o in self.objects:

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

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

            medium_dist = (self.screen_size[0] // self.square_size) // 2

            dataset = [
                # battery
                self.user.obj.battery_status(),
                # high | low |

                # distance between kitchen and object
                0 if o.distance_to(self.kitchen.position) > medium_dist else 1,
                # far | close |

                # mood 
                o.get_mood(self.action_clock),
                # undefined | good | bad |

                # basket is empty
                1 if self.user.obj.basket_is_empty() else 0,
                # yes | no |

                # dish is ready
                1 if o.dish_is_ready(self.action_clock) else 0,
                # yes | no |

                # dish in basket
                1 if self.user.obj.dish_in_basket(o) else 0,
                # yes | no |

                # status
                o.get_state_number(),
                # empty | new order | waiting for dish | have a dish |

                # is actual
                1 if o.isActual() else 0,
                # yes | no |
            ]

            p = self.tree.make_predict(dataset)
            goal_queue.put((p, o.position))

        if goal_queue.queue:
            priority, goal = goal_queue.queue[0]
            if priority == 2:
                self.appendGoalPosition(self.kitchen.position)
            else:
                self.appendGoalPosition(goal)