# Import the pygame module
import pygame

# Import pygame.locals for easier access to key coordinates
# Updated to conform to flake8 and black standards
from pygame.locals import (
    K_UP,
    K_DOWN,
    K_LEFT,
    K_RIGHT,
    K_ESCAPE,
    K_SPACE,
    KEYDOWN,
    QUIT,
)

# Initialize pygame
pygame.init()

# Grid size
GSIZE = 10

# This sets the WIDTH and HEIGHT of each grid location
WIDTH = 35
HEIGHT = 35

# This sets the margin between each cell
MARGIN = 5

# Window size
SCREEN_WIDTH = GSIZE * (WIDTH + MARGIN) + MARGIN 
SCREEN_HEIGHT = GSIZE * (HEIGHT + MARGIN) + MARGIN

# Set some colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BROWN = (140, 95, 65)
DBROWN = (65, 50, 20)
YELLOW = (255, 255, 0)
BLUE = (0, 0, 255)


# Define a Tractor object by extending pygame.sprite.Sprite
# The surface drawn on the screen is now an attribute of 'tractor'
class Tractor(pygame.sprite.Sprite):
    def __init__(self):
        super(Tractor, self).__init__()
        self.surf = pygame.Surface((WIDTH, HEIGHT))
        self.surf.fill(RED)
        self.rect = self.surf.get_rect(topleft=(MARGIN, MARGIN))
        self.position = [0, 0]

    def update(self, pressed_keys):
        if pressed_keys[K_UP]:
            self.rect.move_ip(0, -(HEIGHT + MARGIN))
            self.position[1] -= 1
        if pressed_keys[K_DOWN]:
            self.rect.move_ip(0, HEIGHT + MARGIN)
            self.position[1] += 1 
        if pressed_keys[K_LEFT]:
            self.rect.move_ip(-(WIDTH + MARGIN), 0)
            self.position[0] -= 1 
        if pressed_keys[K_RIGHT]:
            self.rect.move_ip(WIDTH + MARGIN, 0)
            self.position[0] += 1

        if self.rect.left < MARGIN:
            self.rect.left = MARGIN
        if self.position[0] < 0:
            self.position[0] = 0
        if self.rect.right > SCREEN_WIDTH-MARGIN:
            self.rect.right = SCREEN_WIDTH-MARGIN
        if self.position[0] >= GSIZE-1:
            self.position[0] = GSIZE-1
        if self.rect.top <= MARGIN:
            self.rect.top = MARGIN
        if self.position[1] < 0:
            self.position[1] = 0
        if self.rect.bottom >= SCREEN_HEIGHT-MARGIN:
            self.rect.bottom = SCREEN_HEIGHT-MARGIN
        if self.position[1] >= GSIZE-1:
            self.position[1] = GSIZE-1

    def hydrate(self, field):
        field[self.position[0]][self.position[1]].hydrate()

# Define a Field object by extending pygame.sprite.Sprite
# The surface drawn on the screen is now an attribute of 'field'
class Field(pygame.sprite.Sprite):
    def __init__(self, row, column):
        super(Field, self).__init__()
        self.surf = pygame.Surface((WIDTH, HEIGHT))
        self.surf.fill(BROWN)
        self.rect = self.surf.get_rect(topleft=((MARGIN + WIDTH) * row + MARGIN, (MARGIN + HEIGHT) * column + MARGIN))
        self.hydration = 0

    def hydrate(self):
        if self.hydration <= 3: 
            self.hydration += 1
            if self.hydration == 1:
                self.surf.fill(YELLOW)
            if self.hydration == 2:
                self.surf.fill(GREEN)
            if self.hydration == 3:
                self.surf.fill(BLUE)

# Name the window
pygame.display.set_caption("Inteligentny Traktor")

# Create the screen object
# The size is determined by the constant SCREEN_WIDTH and SCREEN_HEIGHT
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))

# Instantiate tractor. Right now, this is just a rectangle.
tractor = Tractor()

# Create field array
field = []

for row in range(GSIZE):
    field.append([])
    for column in range(GSIZE):
        fieldbit = Field(row, column)
        field[row].append(fieldbit)

# Variable to keep the main loop running
running = True

clock = pygame.time.Clock()

# Main loop
while running:
    # Look at every event in the queue
    for event in pygame.event.get():
        # Did the user hit a key?
        if event.type == KEYDOWN:
            # Was it the Escape key? If so, stop the loop.
            if event.key == K_ESCAPE:
                running = False

        # Did the user click the window close button? If so, stop the loop.
        elif event.type == QUIT:
            running = False

    pressed_keys = pygame.key.get_pressed()

    # Set the screen background
    tractor.update(pressed_keys)
    if pressed_keys[K_SPACE]:
        tractor.hydrate(field)

    screen.fill(DBROWN)
    
    # Draw the field
    for row in range(GSIZE):
        for column in range(GSIZE):
            screen.blit(field[row][column].surf, field[row][column].rect)

    # Draw the player on the screen
    screen.blit(tractor.surf, tractor.rect)

    # Update the screen                          
    pygame.display.flip()

    # Ensure program maintains a rate of 30 frames per second
    clock.tick(10)