use crate::tile::{FixtureTile, FloorTile, NeedsOccluded, 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;

#[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";

pub fn initialize_textures(mut commands: Commands, asset_server: Res<AssetServer>) {
    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),
    );

    commands.insert_resource(Textures { handles: textures });
    commands.insert_resource(TextureIDs { refs: texture_ids });
}

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

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

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(),
        }
    }
}

pub fn build_quilted_terrain_sprites(
    query: Query<(&FloorTile, &Transform)>,
    mut commands: Commands,
    mut cwss: ResMut<CurrentWorldSpriteState>,
    textures: Res<Textures>,
    texture_ids: Res<TextureIDs>,
    query_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;

    for entity in query_sprites.iter() {
        commands.entity(entity).despawn();
    }

    let mut tiles_by_z: HashMap<usize, Vec<(Vec2, &FloorTile)>> = HashMap::new();
    // Calculate bounds for all visible tiles
    for z_index in 0..=tilemap::Z_TOTAL as usize {
        for (floortile, transform) in query.iter() {
            let is_visible =
                (floortile.visible_range[z_index / 32] & (1 << (z_index % 32) as u32)) != 0;
            if is_visible {
                let position = Vec2::new(transform.translation.x, transform.translation.y);
                tiles_by_z
                    .entry(z_index)
                    .or_default()
                    .push((position, floortile));
            }
        }
    }

    // For each z-index, create a new quilted texture, baked and composited.
    for (z_index, tiles) in tiles_by_z.iter() {
        if tiles.is_empty() {
            continue;
        }

        // Snap min/max values to tile grid - accounting for half tile on each edge
        // We need to extend the boundaries by half a tile in each direction to ensure full coverage
        // Don't ask, dragons.
        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;

        quilt_cache
            .dimensions
            .insert(*z_index, (width_px, height_px));

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

        for (pos, floortile) in tiles {
            let texture = textures
                .handles
                .get(texture_ids.refs.get(&floortile.id).unwrap())
                .unwrap();

            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![];

            if let Some(source_texture) = images.get(texture) {
                if let Some(_data) = &source_texture.data {
                    // replace this one texture push with loop
                    // if the texture contains any transparent or semi-transparent pixels we should also add the tile below it (z-1)
                    data.push(_data);
                }
                blit_texture(
                    data,              // tile
                    &mut texture_data, // terrain
                    source_texture.size().x as u32,
                    source_texture.size().y as u32,
                    width_px,
                    height_px,
                    target_x,
                    target_y,
                );
            }
        }

        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,
        );
        let texture_handle = images.add(quilted_texture);

        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;

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

    quilt_cache.dirty_indices.clear();
    cwss.state = TerrainSpriteState::RenderReady;
    println!(
        "Quilted world sprites built. Elapsed: {:.2?}",
        now.elapsed()
    );
}

// Helper function to blit a texture onto another texture
fn blit_texture(
    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 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_idx = ((y * source_width) + x) as usize * 4;
            let target_idx = (((y + offset_y) * target_width) + (x + offset_x)) as usize * 4;

            // New colour
            let mut r: u64 = 0;
            let mut g: u64 = 0;
            let mut b: u64 = 0;

            for source in sources.iter() {
                r += source[source_idx] as u64;
                g += source[source_idx + 1] as u64;
                b += source[source_idx + 2] as u64;
            }
            r /= sources.len() as u64;
            g /= sources.len() as u64;
            b /= sources.len() as u64;

            target[target_idx] = r as u8;
            target[target_idx + 1] = g as u8;
            target[target_idx + 2] = b as u8;
            target[target_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,
    needs_occluded: NeedsOccluded,
}

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,
            needs_occluded: NeedsOccluded {
                has_been_occluded: false,
            },
        }
    }

    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,
            needs_occluded: NeedsOccluded {
                has_been_occluded: false,
            },
        }
    }

    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,
            needs_occluded: NeedsOccluded {
                has_been_occluded: false,
            },
        }
    }

    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,
            needs_occluded: NeedsOccluded {
                has_been_occluded: false,
            },
        }
    }

    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,
            needs_occluded: NeedsOccluded {
                has_been_occluded: false,
            },
        }
    }

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

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

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,
            needs_occluded: NeedsOccluded {
                has_been_occluded: false,
            },
        }
    }

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

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

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