#include <SDL.h>
#include <iostream>
#include <vector>

using namespace std;

const int WINDOW_WIDTH = 640;
const int WINDOW_HEIGHT = 480;
const int CELL_SIZE = 20;

struct SnakeSegment {
    int x, y;
};

enum class Direction { UP, DOWN, LEFT, RIGHT };

class SnakeGame {
public:
    SnakeGame(SDL_Renderer* renderer);
    ~SnakeGame();

    void Run();

private:
    void HandleInput();
    void Update();
    void Render();
    void SpawnFruit();
    bool CheckCollision();

    SDL_Renderer* renderer;
    vector<SnakeSegment> snake;
    Direction direction;
    SnakeSegment fruit;
};

SnakeGame::SnakeGame(SDL_Renderer* renderer) : renderer(renderer) {
    direction = Direction::RIGHT;
    snake.push_back({2, 2});
    SpawnFruit();
}

SnakeGame::~SnakeGame() {
    SDL_DestroyRenderer(renderer);
    SDL_Quit();
}

void SnakeGame::Run() {
    SDL_Event e;
    bool quit = false;

    while (!quit) {
        while (SDL_PollEvent(&e) != 0) {
            if (e.type == SDL_QUIT) {
                quit = true;
            }
        }

        HandleInput();
        Update();
        Render();
        SDL_Delay(100);
    }
}

void SnakeGame::HandleInput() {
    const Uint8* keys = SDL_GetKeyboardState(nullptr);

    if (keys[SDL_SCANCODE_UP] && direction != Direction::DOWN) {
        direction = Direction::UP;
    } else if (keys[SDL_SCANCODE_DOWN] && direction != Direction::UP) {
        direction = Direction::DOWN;
    } else if (keys[SDL_SCANCODE_LEFT] && direction != Direction::RIGHT) {
        direction = Direction::LEFT;
    } else if (keys[SDL_SCANCODE_RIGHT] && direction != Direction::LEFT) {
        direction = Direction::RIGHT;
    }
}

void SnakeGame::Update() {
    SnakeSegment newHead = snake.front();

    switch (direction) {
        case Direction::UP:
            newHead.y--;
            break;
        case Direction::DOWN:
            newHead.y++;
            break;
        case Direction::LEFT:
            newHead.x--;
            break;
        case Direction::RIGHT:
            newHead.x++;
            break;
    }

    snake.insert(snake.begin(), newHead);

    if (newHead.x == fruit.x && newHead.y == fruit.y) {
        SpawnFruit();
    } else {
        snake.pop_back();
    }

    if (CheckCollision()) {
        cout << "Game Over!" << endl;
        exit(0);
    }
}

void SnakeGame::Render() {
    SDL_SetRenderDrawColor(renderer, 0, 0, 0, 255);
    SDL_RenderClear(renderer);

    SDL_SetRenderDrawColor(renderer, 255, 255, 255, 255);

    for (const auto& segment : snake) {
        SDL_Rect rect = {segment.x * CELL_SIZE, segment.y * CELL_SIZE, CELL_SIZE, CELL_SIZE};
        SDL_RenderFillRect(renderer, &rect);
    }

    SDL_SetRenderDrawColor(renderer, 255, 0, 0, 255);
    SDL_Rect fruitRect = {fruit.x * CELL_SIZE, fruit.y * CELL_SIZE, CELL_SIZE, CELL_SIZE};
    SDL_RenderFillRect(renderer, &fruitRect);

    SDL_RenderPresent(renderer);
}

void SnakeGame::SpawnFruit() {
    fruit.x = rand() % (WINDOW_WIDTH / CELL_SIZE);
    fruit.y = rand() % (WINDOW_HEIGHT / CELL_SIZE);
}

bool SnakeGame::CheckCollision() {
    // Check collision with walls
    if (snake.front().x < 0 || snake.front().x >= WINDOW_WIDTH / CELL_SIZE ||
        snake.front().y < 0 || snake.front().y >= WINDOW_HEIGHT / CELL_SIZE) {
        return true;
    }

    // Check collision with itself
    for (size_t i = 1; i < snake.size(); ++i) {
        if (snake.front().x == snake[i].x && snake.front().y == snake[i].y) {
            return true;
        }
    }

    return false;
}

int main() {
    if (SDL_Init(SDL_INIT_VIDEO) < 0) {
        cout << "SDL could not initialize! SDL_Error: " << SDL_GetError() << endl;
        return 1;
    }

    SDL_Window* window = SDL_CreateWindow("Snake Game", SDL_WINDOWPOS_UNDEFINED, SDL_WINDOWPOS_UNDEFINED,
                                          WINDOW_WIDTH, WINDOW_HEIGHT, SDL_WINDOW_SHOWN);

    if (!window) {
        cout << "Window could not be created! SDL_Error: " << SDL_GetError() << endl;
        return 1;
    }

    SDL_Renderer* renderer = SDL_CreateRenderer(window, -1, SDL_RENDERER_ACCELERATED);
    if (!renderer) {
        cout << "Renderer could not be created! SDL_Error: " << SDL_GetError() << endl;
        return 1;
    }

    SnakeGame game(renderer);
    game.Run();

    SDL_DestroyWindow(window);

    return 0;
}