use crate::tile::{FixtureTile, FloorTile, TileState};
use crate::tilemap::Z_BELOW;
use crate::{constants::*, tilemap};
use bevy::asset::RenderAssetUsages;
use bevy::prelude::*;
use bevy::render::render_resource;
use bevy_platform::collections::hash_map::HashMap;
use rayon::prelude::*;

#[derive(Resource)]
pub struct Textures {
    pub handles: HashMap<String, Handle<Image>>,
}

#[derive(Resource)]
pub struct TextureIDs {
    pub refs: HashMap<u32, String>,
}

const FLOOR_ID_OFFSET: u32 = 0;
const FIXTURE_ID_OFFSET: u32 = 500000;

pub const DEFAULT_TEXTURE: &str = "default.png";

const GRASS_FLOOR_PATH: &str = "grass_floor.png";
const DIRT_FLOOR_PATH: &str = "dirt_floor.png";
const ROCK_FLOOR_PATH: &str = "rock_floor.png";
const BEDROCK_FLOOR_PATH: &str = "bedrock_floor.png";
const SKY_PATH: &str = "sky.png";
const DIRT_WALL_PATH: &str = "dirt_wall.png";
const ROCK_WALL_PATH: &str = "rock_wall.png";
const BEDROCK_WALL_PATH: &str = "bedrock_wall.png";
const LOG_PATH: &str = "log.png";
const LEAVES_PATH: &str = "leaves.png";

pub fn initialize_textures(mut commands: Commands, asset_server: Res<AssetServer>) {
    println!("");
    println!("####################");
    let start = Instant::now();
    let mut textures: HashMap<String, Handle<Image>> = HashMap::new();
    let mut texture_ids: HashMap<u32, String> = HashMap::new();
    texture_ids.insert(u32::MAX, DEFAULT_TEXTURE.to_string());

    texture_ids.insert(FLOOR_ID_OFFSET, SKY_PATH.to_string());
    textures.insert(SKY_PATH.to_string(), asset_server.load(SKY_PATH));
    texture_ids.insert(FLOOR_ID_OFFSET + 1, GRASS_FLOOR_PATH.to_string());
    textures.insert(
        GRASS_FLOOR_PATH.to_string(),
        asset_server.load(GRASS_FLOOR_PATH),
    );
    texture_ids.insert(FLOOR_ID_OFFSET + 2, DIRT_FLOOR_PATH.to_string());
    textures.insert(
        DIRT_FLOOR_PATH.to_string(),
        asset_server.load(DIRT_FLOOR_PATH),
    );
    texture_ids.insert(FLOOR_ID_OFFSET + 3, ROCK_FLOOR_PATH.to_string());
    textures.insert(
        ROCK_FLOOR_PATH.to_string(),
        asset_server.load(ROCK_FLOOR_PATH),
    );
    texture_ids.insert(FLOOR_ID_OFFSET + 4, BEDROCK_FLOOR_PATH.to_string());
    textures.insert(
        BEDROCK_FLOOR_PATH.to_string(),
        asset_server.load(BEDROCK_FLOOR_PATH),
    );

    texture_ids.insert(FIXTURE_ID_OFFSET + 1, DIRT_WALL_PATH.to_string());
    textures.insert(
        DIRT_WALL_PATH.to_string(),
        asset_server.load(DIRT_WALL_PATH),
    );
    texture_ids.insert(FIXTURE_ID_OFFSET + 2, ROCK_WALL_PATH.to_string());
    textures.insert(
        ROCK_WALL_PATH.to_string(),
        asset_server.load(ROCK_WALL_PATH),
    );
    texture_ids.insert(FIXTURE_ID_OFFSET + 3, BEDROCK_WALL_PATH.to_string());
    textures.insert(
        BEDROCK_WALL_PATH.to_string(),
        asset_server.load(BEDROCK_WALL_PATH),
    );
    texture_ids.insert(FIXTURE_ID_OFFSET + 4, LOG_PATH.to_string());
    textures.insert(LOG_PATH.to_string(), asset_server.load(LOG_PATH));
    texture_ids.insert(FIXTURE_ID_OFFSET + 5, LEAVES_PATH.to_string());
    textures.insert(LEAVES_PATH.to_string(), asset_server.load(LEAVES_PATH));

    commands.insert_resource(Textures { handles: textures });
    commands.insert_resource(TextureIDs { refs: texture_ids });
    println!("Textures initialized in {:.2?}", start.elapsed());
}

#[derive(Debug, PartialEq)]
pub enum TerrainSpriteState {
    Inactive,
    WaitingForRender,
    InProgress,
    RenderReady,
}

#[derive(Resource)]
pub struct CurrentWorldSpriteState {
    pub state: TerrainSpriteState,
}

use std::process::exit;
use std::sync::Mutex;
use std::time::Instant;

#[derive(Component)]
pub struct TerrainSprite {
    pub z_index: usize, // Store the z-index this sprite belongs to (as world_Z)
}

#[derive(Resource)]
pub struct QuiltCache {
    pub dimensions: HashMap<usize, (u32, u32)>,
    pub dirty_indices: Vec<usize>,
}

impl Default for QuiltCache {
    fn default() -> Self {
        Self {
            dimensions: HashMap::new(),
            dirty_indices: Vec::new(),
        }
    }
}

// here be dragons :(
pub fn build_quilted_terrain_sprites(
    query_tiles: Query<(&FloorTile, &Transform)>,
    commands: ParallelCommands<'_, '_>,
    mut cwss: ResMut<CurrentWorldSpriteState>,
    textures: Res<Textures>,
    texture_ids: Res<TextureIDs>,
    query_terrain_sprites: Query<Entity, With<TerrainSprite>>,
    mut images: ResMut<Assets<Image>>,
    mut quilt_cache: ResMut<QuiltCache>,
) {
    if cwss.state != TerrainSpriteState::WaitingForRender {
        return;
    }
    let now = Instant::now();
    cwss.state = TerrainSpriteState::InProgress;

    // Despawn existing terrain sprites
    let despawn_entities: Vec<Entity> = query_terrain_sprites.iter().collect();
    for entity in despawn_entities {
        commands.command_scope(|mut cmd| {
            cmd.entity(entity).despawn();
        });
    }

    // Collect tiles by z-index first
    let mut tiles_by_z: HashMap<usize, Vec<(Vec2, &FloorTile)>> = HashMap::new();

    // Calculate bounds for all visible tiles
    for (floortile, transform) in query_tiles.iter() {
        let position = Vec2::new(transform.translation.x, transform.translation.y);

        for z_index in 0..=tilemap::Z_TOTAL as usize {
            if (floortile.visible_range[z_index / 32] & (1 << (z_index % 32) as u32)) != 0 {
                tiles_by_z
                    .entry(z_index)
                    .or_default()
                    .push((position, floortile));
            }
        }
    }

    // Thread-safe collections to store results
    let dimensions_mutex = Mutex::new(HashMap::new());
    let texture_handles_mutex = Mutex::new(HashMap::new());

    // Process each z-level in parallel
    let z_indices: Vec<usize> = tiles_by_z.keys().cloned().collect();

    z_indices.into_par_iter().for_each(|z_index| {
        let tiles = if let Some(tiles) = tiles_by_z.get(&z_index) {
            tiles
        } else {
            return; // Skip empty z-levels
        };

        if tiles.is_empty() {
            return;
        }

        // Calculate bounds
        let min_x_aligned: f32 = ((tiles.iter().map(|(pos, _)| pos.x).reduce(f32::min).unwrap()
            - TILE_SIZE / 2.0)
            / TILE_SIZE)
            .floor()
            * TILE_SIZE;
        let min_y_aligned: f32 = ((tiles.iter().map(|(pos, _)| pos.y).reduce(f32::min).unwrap()
            - TILE_SIZE / 2.0)
            / TILE_SIZE)
            .floor()
            * TILE_SIZE;
        let max_x_aligned: f32 = ((tiles.iter().map(|(pos, _)| pos.x).reduce(f32::max).unwrap()
            + TILE_SIZE / 2.0)
            / TILE_SIZE)
            .ceil()
            * TILE_SIZE;
        let max_y_aligned: f32 = ((tiles.iter().map(|(pos, _)| pos.y).reduce(f32::max).unwrap()
            + TILE_SIZE / 2.0)
            / TILE_SIZE)
            .ceil()
            * TILE_SIZE;

        // Calculate terrain texture dimensions in pixels
        let width_tiles = ((max_x_aligned - min_x_aligned) / TILE_SIZE) as u32;
        let height_tiles = ((max_y_aligned - min_y_aligned) / TILE_SIZE) as u32;
        let width_px = width_tiles * TILE_PIXELS;
        let height_px = height_tiles * TILE_PIXELS;

        // Store dimensions in our thread-safe map
        dimensions_mutex
            .lock()
            .unwrap()
            .insert(z_index, (width_px, height_px));

        let mut texture_data = vec![0u8; (width_px * height_px * 4) as usize];

        // Process tiles for this z-level
        // We can't parallelize this inner loop without more complex locking on texture_data
        for (pos, floortile) in tiles {
            if let Some(texture_id) = texture_ids.refs.get(&floortile.id) {
                if let Some(texture) = textures.handles.get(texture_id) {
                    let rel_x = pos.x - min_x_aligned;
                    let rel_y = pos.y - min_y_aligned;

                    let tile_x = (rel_x / TILE_SIZE).round() as u32;
                    let tile_y = (height_tiles as f32 - 1.0 - (rel_y / TILE_SIZE).round()) as u32;

                    let target_x = tile_x * TILE_PIXELS;
                    let target_y = tile_y * TILE_PIXELS;
                    let mut data: Vec<&[u8]> = vec![];

                    let mut base_texture: &Image = &Default::default();

                    // Use a thread-safe approach to access images
                    // In a full implementation, this would require a more sophisticated
                    // thread-safe access pattern to Assets<Image>
                    if let Some(_base_texture) = images.get(texture) {
                        if let Some(_data) = &_base_texture.data {
                            data.push(_data);
                            base_texture = _base_texture;
                        }
                    }

                    blit_texture_with_alpha(
                        data,              // tile
                        &mut texture_data, // terrain
                        base_texture.size().x as u32,
                        base_texture.size().y as u32,
                        width_px,
                        height_px,
                        target_x,
                        target_y,
                    );
                }
            }
        }

        // Create the quilted texture
        let quilted_texture = Image::new_fill(
            render_resource::Extent3d {
                width: width_px,
                height: height_px,
                depth_or_array_layers: 1,
            },
            render_resource::TextureDimension::D2,
            &texture_data,
            render_resource::TextureFormat::Rgba8UnormSrgb,
            RenderAssetUsages::RENDER_WORLD,
        );

        // In a real implementation, we would need thread-safe access to images
        // For now, we'll collect the textures and add them after parallel processing
        let center_x = min_x_aligned + (max_x_aligned - min_x_aligned) / 2.0;
        let center_y = min_y_aligned + (max_y_aligned - min_y_aligned) / 2.0;

        // Store texture and position data for later spawning
        texture_handles_mutex
            .lock()
            .unwrap()
            .insert(z_index, (quilted_texture, center_x, center_y));
    });

    // Process collected textures and spawn entities
    let collected_dimensions = dimensions_mutex.into_inner().unwrap();
    let collected_textures = texture_handles_mutex.into_inner().unwrap();

    // Update quilt_cache dimensions
    for (z_index, dimensions) in collected_dimensions {
        quilt_cache.dimensions.insert(z_index, dimensions);
    }

    // Add images and spawn entities with the collected data
    for (z_index, (quilted_texture, center_x, center_y)) in collected_textures {
        let texture_handle = images.add(quilted_texture);

        commands.command_scope(|mut cmd| {
            cmd.spawn((
                Sprite {
                    image: texture_handle,
                    ..Default::default()
                },
                Transform::from_xyz(
                    center_x - TILE_SIZE / 2.0,
                    center_y - TILE_SIZE / 2.0,
                    -Z_BELOW * TILE_SIZE,
                )
                .with_scale(Vec3::splat(PIXEL_RATIO)),
                Visibility::Hidden,
                TerrainSprite { z_index },
            ));
        });
    }

    quilt_cache.dirty_indices.clear();
    cwss.state = TerrainSpriteState::RenderReady;
    println!("Terrain sprites baked in: {:.2?}", now.elapsed());
}

// Helper function to blit a texture onto another texture
fn blit_texture_with_alpha(
    sources: Vec<&[u8]>,
    target: &mut [u8],
    source_width: u32,
    source_height: u32,
    target_width: u32,
    target_height: u32,
    offset_x: u32,
    offset_y: u32,
) {
    for (_, source_data) in sources.iter().rev().enumerate() {
        for y in 0..source_height {
            if y + offset_y >= target_height {
                continue;
            }
            for x in 0..source_width {
                if x + offset_x >= target_width {
                    continue;
                }

                let source_pixel_idx = ((y * source_width) + x) as usize * 4;
                let target_pixel_idx =
                    (((y + offset_y) * target_width) + (x + offset_x)) as usize * 4;

                let src_a = source_data[source_pixel_idx + 3];

                if src_a == 0 {
                    continue;
                }

                let src_r = source_data[source_pixel_idx];
                let src_g = source_data[source_pixel_idx + 1];
                let src_b = source_data[source_pixel_idx + 2];

                if src_a == 255 {
                    target[target_pixel_idx] = src_r;
                    target[target_pixel_idx + 1] = src_g;
                    target[target_pixel_idx + 2] = src_b;
                    target[target_pixel_idx + 3] = 255;
                } else {
                    let dst_r = target[target_pixel_idx];
                    let dst_g = target[target_pixel_idx + 1];
                    let dst_b = target[target_pixel_idx + 2];

                    let alpha_factor = src_a as f32 / 255.0;
                    let inv_alpha = 1.0 - alpha_factor;

                    target[target_pixel_idx] =
                        (src_r as f32 * alpha_factor + dst_r as f32 * inv_alpha) as u8;
                    target[target_pixel_idx + 1] =
                        (src_g as f32 * alpha_factor + dst_g as f32 * inv_alpha) as u8;
                    target[target_pixel_idx + 2] =
                        (src_b as f32 * alpha_factor + dst_b as f32 * inv_alpha) as u8;
                    target[target_pixel_idx + 3] = 255;
                }
            }
        }
    }
}

// This system handles updating quilted sprites when the world changes
// pub fn update_quilts_on_world_change(
//     mut commands: Commands,
//     mut quilt_cache: ResMut<QuiltCache>,
//     mut cwss: ResMut<CurrentWorldSpriteState>,
// Add any resources or queries that indicate world changes
// For example, if you have a WorldChangeEvent:
// mut world_changes: EventReader<WorldChangeEvent>,
// ) {
// Example: Check for world changes
// if !world_changes.is_empty() {
//     for event in world_changes.iter() {
//         quilt_cache.dirty_indices.push(event.z_index);
//     }
//     cwss.state = WorldSpriteState::WaitingForRender;
// }

// Alternatively, if you have specific systems that modify the world,
// you could have them set cwss.state = WorldSpriteState::WaitingForRender
// and add affected z-indices to quilt_cache.dirty_indices
// }

#[derive(Bundle)]
pub struct FloorTilePrefab {
    transform: Transform,
    tile: FloorTile,
    tile_state: TileState,
    visibility: Visibility,
}

impl FloorTilePrefab {
    pub fn grass(position: Vec3) -> Self {
        FloorTilePrefab {
            transform: Transform::from_translation(position),
            tile: FloorTile {
                id: 1,
                astar_weight: 100,
                ..Default::default()
            },
            tile_state: TileState {
                timer: Timer::from_seconds(1.0, TimerMode::Repeating),
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn dirt(position: Vec3) -> Self {
        FloorTilePrefab {
            transform: Transform::from_translation(position),
            tile: FloorTile {
                id: 2,
                astar_weight: 85,
                ..Default::default()
            },
            tile_state: TileState {
                timer: Timer::from_seconds(1.0, TimerMode::Repeating),
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn rock(position: Vec3) -> Self {
        FloorTilePrefab {
            transform: Transform::from_translation(position),
            tile: FloorTile {
                id: 3,
                astar_weight: 50,
                ..Default::default()
            },
            tile_state: TileState {
                timer: Timer::from_seconds(1.0, TimerMode::Repeating),
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn air(position: Vec3) -> Self {
        FloorTilePrefab {
            transform: Transform::from_translation(position),
            tile: FloorTile {
                id: 0,
                opaque: false,
                walkable: false,
                ..Default::default()
            },
            tile_state: TileState {
                timer: Timer::from_seconds(1.0, TimerMode::Repeating),
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn bedrock(position: Vec3) -> Self {
        FloorTilePrefab {
            transform: Transform::from_translation(position),
            tile: FloorTile {
                id: 4,
                astar_weight: 150,
                ..Default::default()
            },
            tile_state: TileState {
                timer: Timer::from_seconds(1.0, TimerMode::Repeating),
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn spawn(self, commands: &mut Commands) {
        commands.spawn((self.tile, self.transform, self.tile_state, self.visibility));
    }
}

#[derive(Bundle)]
pub struct FixtureTilePrefab {
    transform: Transform,
    tile: FixtureTile,
    visibility: Visibility,
}

impl FixtureTilePrefab {
    pub fn dirt_wall(position: Vec3) -> Self {
        FixtureTilePrefab {
            transform: Transform::from_translation(position),
            tile: FixtureTile {
                id: FIXTURE_ID_OFFSET + 1,
                ..Default::default()
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn rock_wall(position: Vec3) -> Self {
        FixtureTilePrefab {
            transform: Transform::from_translation(position),
            tile: FixtureTile {
                id: FIXTURE_ID_OFFSET + 2,
                ..Default::default()
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn log(position: Vec3) -> Self {
        FixtureTilePrefab {
            transform: Transform::from_translation(position),
            tile: FixtureTile {
                id: FIXTURE_ID_OFFSET + 4,
                ..Default::default()
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn leaves(position: Vec3) -> Self {
        FixtureTilePrefab {
            transform: Transform::from_translation(position),
            tile: FixtureTile {
                id: FIXTURE_ID_OFFSET + 5,
                ..Default::default()
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn bedrock_wall(position: Vec3) -> Self {
        FixtureTilePrefab {
            transform: Transform::from_translation(position),
            tile: FixtureTile {
                id: 0,
                ..Default::default()
            },
            visibility: Visibility::Hidden,
        }
    }

    pub fn spawn(self, commands: &mut Commands) -> Entity {
        commands
            .spawn((self.tile, self.transform, self.visibility))
            .id()
    }
}