use bevy::prelude::*;
use rand::random;
use std::time::Duration;

// Constants
const ARENA_WIDTH: u32 = 50;
const ARENA_HEIGHT: u32 = 50;
const SNAKE_HEAD_COLOR: Color = Color::srgb(0.7, 0.7, 0.7);
const SNAKE_SEGMENT_COLOR: Color = Color::srgb(0.3, 0.3, 0.3);
const FOOD_COLOR: Color = Color::srgb(1.0, 0.0, 1.0);

#[derive(Component, Clone, Copy, PartialEq, Eq)]
struct Position {
    x: i32,
    y: i32,
}

#[derive(Component)]
struct Size {
    width: f32,
    height: f32,
}

impl Size {
    pub fn square(x: f32) -> Self {
        Self {
            width: x,
            height: x,
        }
    }
}

#[derive(Component)]
struct SnakeHead {
    direction: Direction,
}

#[derive(Component)]
struct SnakeSegment;

#[derive(Component)]
struct Food;

#[derive(Resource, Default)]
struct SnakeSegments(Vec<Entity>);

#[derive(Resource, Default)]
struct LastTailPosition(Option<Position>);

#[derive(PartialEq, Copy, Clone)]
enum Direction {
    Left,
    Up,
    Right,
    Down,
}

impl Direction {
    fn opposite(self) -> Self {
        match self {
            Self::Left => Self::Right,
            Self::Right => Self::Left,
            Self::Up => Self::Down,
            Self::Down => Self::Up,
        }
    }
}

fn main() {
    App::new()
        .add_plugins(DefaultPlugins.set(WindowPlugin {
            primary_window: Some(Window {
                title: "Snake!".to_string(),
                resolution: bevy::window::WindowResolution::new(500, 500),
                ..default()
            }),
            ..default()
        }))
        .insert_resource(ClearColor(Color::srgb(0.04, 0.04, 0.04)))
        .insert_resource(SnakeSegments::default())
        .insert_resource(LastTailPosition::default())
        .add_message::<GameOverEvent>()
        .add_message::<GrowthEvent>()
        .add_systems(Startup, (setup_camera, spawn_snake))
        .add_systems(
            Update,
            (
                snake_movement_input,
                game_over.after(snake_movement),
                food_spawner,
                snake_movement,
                snake_eating,
                snake_growth,
                size_scaling,
                position_translation,
            ),
        )
        .run();
}

fn setup_camera(mut commands: Commands) {
    commands.spawn(Camera2d::default());
}

fn spawn_snake(mut commands: Commands, mut segments: ResMut<SnakeSegments>) {
    *segments = SnakeSegments(vec![
        commands
            .spawn((
                SnakeHead {
                    direction: Direction::Up,
                },
                SnakeSegment,
                Position { x: 3, y: 3 },
                Size::square(0.8),
                Sprite {
                    color: SNAKE_HEAD_COLOR,
                    ..default()
                },
            ))
            .id(),
        spawn_segment(&mut commands, Position { x: 3, y: 2 }),
    ]);
}

fn spawn_segment(commands: &mut Commands, position: Position) -> Entity {
    commands
        .spawn((
            SnakeSegment,
            position,
            Size::square(0.65),
            Sprite {
                color: SNAKE_SEGMENT_COLOR,
                ..default()
            },
        ))
        .id()
}

#[derive(Message)]
struct GameOverEvent;

fn snake_movement_input(
    keyboard_input: Res<ButtonInput<KeyCode>>,
    mut heads: Query<&mut SnakeHead>,
) {
    if let Some(mut head) = heads.iter_mut().next() {
        let dir: Direction = if keyboard_input.pressed(KeyCode::ArrowLeft) {
            Direction::Left
        } else if keyboard_input.pressed(KeyCode::ArrowDown) {
            Direction::Down
        } else if keyboard_input.pressed(KeyCode::ArrowUp) {
            Direction::Up
        } else if keyboard_input.pressed(KeyCode::ArrowRight) {
            Direction::Right
        } else {
            head.direction
        };

        if dir != head.direction.opposite() {
            head.direction = dir;
        }
    }
}

fn snake_movement(
    mut heads: Query<(Entity, &mut SnakeHead)>,
    mut positions: Query<&mut Position>,
    segments: Res<SnakeSegments>,
    mut last_tail_position: ResMut<LastTailPosition>,
    mut game_over_writer: MessageWriter<GameOverEvent>,
    time: Res<Time>,
    mut timer: Local<Timer>,
) {
    if timer.duration() == Duration::ZERO {
        *timer = Timer::from_seconds(0.15, TimerMode::Repeating);
    }
    
    if !timer.tick(time.delta()).just_finished() {
        return;
    }

    if let Some((head_entity, head)) = heads.iter_mut().next() {
        let segment_positions: Vec<Position> = segments
            .0
            .iter()
            .map(|e| *positions.get(*e).unwrap())
            .collect();

        // Check for collision with tail *before* moving
        if segment_positions.iter().skip(1).any(|p| *p == segment_positions[0]) {
             game_over_writer.write(GameOverEvent);
        }

        let mut head_pos = positions.get_mut(head_entity).unwrap();
        
        // Remember tail position for growth
        *last_tail_position = LastTailPosition(Some(*segment_positions.last().unwrap()));

        match &head.direction {
            Direction::Left => head_pos.x -= 1,
            Direction::Right => head_pos.x += 1,
            Direction::Up => head_pos.y += 1,
            Direction::Down => head_pos.y -= 1,
        };

        // Wall collision
        if head_pos.x < 0 || head_pos.y < 0 || head_pos.x as u32 >= ARENA_WIDTH || head_pos.y as u32 >= ARENA_HEIGHT {
            game_over_writer.write(GameOverEvent);
        }

        // Move the rest of the body
        segment_positions
            .iter()
            .zip(segments.0.iter().skip(1))
            .for_each(|(pos, segment)| {
                *positions.get_mut(*segment).unwrap() = *pos;
            });
    }
}

fn snake_eating(
    mut commands: Commands,
    mut growth_writer: MessageWriter<GrowthEvent>,
    food_positions: Query<(Entity, &Position), With<Food>>,
    head_positions: Query<&Position, With<SnakeHead>>,
) {
    for head_pos in head_positions.iter() {
        for (ent, food_pos) in food_positions.iter() {
            if head_pos == food_pos {
                commands.entity(ent).despawn();
                growth_writer.write(GrowthEvent);
            }
        }
    }
}

#[derive(Message)]
struct GrowthEvent;

fn snake_growth(
    mut commands: Commands,
    last_tail_position: Res<LastTailPosition>,
    mut segments: ResMut<SnakeSegments>,
    mut growth_reader: MessageReader<GrowthEvent>,
) {
    if growth_reader.read().next().is_some() {
        if let Some(pos) = last_tail_position.0 {
            segments.0.push(spawn_segment(&mut commands, pos));
        }
    }
}

fn size_scaling(primary_query: Query<&Window, With<bevy::window::PrimaryWindow>>, mut q: Query<(&Size, &mut Transform)>) {
    let window = primary_query.single().unwrap();
    for (sprite_size, mut transform) in q.iter_mut() {
        transform.scale = Vec3::new(
            sprite_size.width / ARENA_WIDTH as f32 * window.width() as f32,
            sprite_size.height / ARENA_HEIGHT as f32 * window.height() as f32,
            1.0,
        );
    }
}

fn position_translation(primary_query: Query<&Window, With<bevy::window::PrimaryWindow>>, mut q: Query<(&Position, &mut Transform)>) {
    fn convert(pos: f32, bound_window: f32, bound_game: f32) -> f32 {
        let tile_size = bound_window / bound_game;
        pos / bound_game * bound_window - (bound_window / 2.) + (tile_size / 2.)
    }
    let window = primary_query.single().unwrap();
    for (pos, mut transform) in q.iter_mut() {
        transform.translation = Vec3::new(
            convert(pos.x as f32, window.width() as f32, ARENA_WIDTH as f32),
            convert(pos.y as f32, window.height() as f32, ARENA_HEIGHT as f32),
            0.0,
        );
    }
}

fn food_spawner(mut commands: Commands, food: Query<&Food>) {
    // Check if food exists
    if food.iter().len() == 0 {
         commands.spawn((
            Food,
            Position {
                x: (random::<f32>() * ARENA_WIDTH as f32) as i32,
                y: (random::<f32>() * ARENA_HEIGHT as f32) as i32,
            },
            Size::square(0.8),
            Sprite {
                color: FOOD_COLOR,
                ..default()
            },
        ));
    }
}

fn game_over(mut reader: MessageReader<GameOverEvent>) {
    if reader.read().next().is_some() {
        println!("Game Over!");
        std::process::exit(0); 
    }
}