diff --git a/a2/snake.c b/a2/snake.c
index 9d58be1..064bb80 100644
--- a/a2/snake.c
+++ b/a2/snake.c
@@ -1,142 +1,55 @@
-#include <stdio.h>
+#include "snake.h"
 #include <stdlib.h>
-#include <time.h>
 
-#define FOOD_COUNT 3
-#define MAX_SNAKE_LENGTH 100
-
-// Define the direction constants
-#define UP    0
-#define LEFT  1
-#define DOWN  2
-#define RIGHT 3
-
-struct state {
-    int width;
-    int height;
-    int snake_length;
-    int snake_x[MAX_SNAKE_LENGTH];
-    int snake_y[MAX_SNAKE_LENGTH];
-    int foodx[FOOD_COUNT];
-    int foody[FOOD_COUNT];
-    int sx;
-    int sy;
-};
-
-// Function prototypes
-void init_state(struct state* state, int width, int height);
-void draw_state(struct state* state);
-void process_input(struct state* state);
-int step_state(struct state* state);
-void generate_food(struct state* state);
-int is_snake(struct state* state, int x, int y);
-void free_snake(struct state* state);
-
-// Add a new segment to the snake
-void add_snake(struct state* state, int x, int y) {
-    state->snake_x[state->snake_length] = x;
-    state->snake_y[state->snake_length] = y;
-    state->snake_length++;
-}
-
-// Check if the given position is occupied by the snake
-int is_snake(struct state* state, int x, int y) {
-    for (int i = 0; i < state->snake_length; i++) {
-        if (state->snake_x[i] == x && state->snake_y[i] == y) {
-            return 1;
+// Pridá novú odmenu na náhodnú voľnú pozíciu v rámci hracej plochy
+void add_food(struct state* state) {
+    int placed = 0;
+    while (!placed) {
+        int x = rand() % state->width;
+        int y = rand() % state->height;
+        if (!is_snake(state->snake, x, y)) {
+            state->foodx[state->foods] = x;
+            state->foody[state->foods] = y;
+            state->foods++;
+            placed = 1;
         }
     }
-    return 0;
 }
 
-// Free the memory allocated for the snake
-void free_snake(struct state* state) {
-    // No dynamic memory allocation for the snake in this implementation
-}
-
-// Move the snake by one step
 int step_state(struct state* state) {
-    // Move the snake body
-    for (int i = state->snake_length - 1; i > 0; i--) {
-        state->snake_x[i] = state->snake_x[i - 1];
-        state->snake_y[i] = state->snake_y[i - 1];
+    // Vypočíta novú pozíciu hlavy hada
+    int new_x = state->snake->x + state->sx;
+    int new_y = state->snake->y + state->sy;
+
+    // Overí, či je nová pozícia mimo hracej plochy
+    if (new_x < 0 || new_x >= state->width || new_y < 0 || new_y >= state->height) {
+        return END_WALL;
     }
 
-    // Move the snake head
-    state->snake_x[0] += state->sx;
-    state->snake_y[0] += state->sy;
-
-    // Check for collisions with walls
-    if (state->snake_x[0] < 0 || state->snake_x[0] >= state->width ||
-        state->snake_y[0] < 0 || state->snake_y[0] >= state->height) {
-        return 1; // Collision with wall
+    // Overí, či je nová pozícia na tele hada
+    if (is_snake(state->snake, new_x, new_y)) {
+        return END_SNAKE;
     }
 
-    // Check for collisions with itself
-    for (int i = 1; i < state->snake_length; i++) {
-        if (state->snake_x[i] == state->snake_x[0] && state->snake_y[i] == state->snake_y[0]) {
-            return 2; // Collision with itself
-        }
-    }
-
-    // Check for collisions with food
+    // Overí, či je na novej pozícii odmena
     for (int i = 0; i < FOOD_COUNT; i++) {
-        if (state->snake_x[0] == state->foodx[i] && state->snake_y[0] == state->foody[i]) {
-            // Increase the snake length
-            state->snake_length++;
-            if (state->snake_length >= MAX_SNAKE_LENGTH) {
-                return 3; // Snake length exceeds the maximum
-            }
-            // Move the food to a new random position
-            state->foodx[i] = rand() % state->width;
-            state->foody[i] = rand() % state->height;
-            return 0; // Food eaten
-        }
-    }
-
-    return 0; // No collision
-}
-
-void init_state(struct state* state, int width, int height) {
-    state->width = width;
-    state->height = height;
-    state->snake_length = 1;
-    state->snake_x[0] = width / 2;
-    state->snake_y[0] = height / 2;
-    state->sx = 0;
-    state->sy = -1; // Start moving upward
-    for (int i = 0; i < FOOD_COUNT; i++) {
-        state->foodx[i] = rand() % width;
-        state->foody[i] = rand() % height;
-    }
-}
-
-void generate_food(struct state* state) {
-    for (int i = 0; i < FOOD_COUNT; i++) {
-        state->foodx[i] = rand() % state->width;
-        state->foody[i] = rand() % state->height;
-    }
-}
-
-void draw_state(struct state* state) {
-    for (int y = 0; y < state->height; y++) {
-        for (int x = 0; x < state->width; x++) {
-            if (x == state->snake_x[0] && y == state->snake_y[0]) {
-                printf("H"); // Snake head
-            } else if (is_snake(state, x, y)) {
-                printf("S"); // Snake body
+        if (state->foodx[i] == new_x && state->foody[i] == new_y) {
+            state->foodx[i] = -1;
+            state->foody[i] = -1;
+            state->foods_eaten++;
+            add_food(state); // Pridá novú odmenu
+            state->snake = add_snake(state->snake, new_x, new_y);
+            if (state->foods_eaten == FOOD_COUNT) {
+                return END_FOOD;
             } else {
-                int is_food = 0;
-                for (int i = 0; i < FOOD_COUNT; i++) {
-                    if (state->foodx[i] == x && state->foody[i] == y) {
-                        is_food = 1;
-                        break;
-                    }
-                }
-                printf(is_food ? "F" : ".");
+                return END_CONTINUE;
             }
         }
-        printf("\n");
     }
-}
 
+    // Odstráni poslednú časť hada
+    state->snake = add_snake(state->snake, new_x, new_y);
+    state->snake = remove_snake(state->snake);
+
+    return END_CONTINUE;
+}