use std::num::NonZeroUsize;
use std::rc::Rc;

use glam::{Mat4, Vec2, Vec3Swizzles};
use hashbrown::HashMap;
use lru::LruCache;
use ordered_float::OrderedFloat;
use wgpu::util::DeviceExt;
use wgpu::{Device, Queue, RenderPass, TextureFormat};

use crate::commands::Draw;
use crate::renderer::traits::{Shader, SpritePipeline};
use crate::types::{DrawContent, DrawContentCacheKey};
use crate::utils::bytes::to_wgsl_bytes;
use crate::utils::rgb::to_linear_rgba;

use crate::platform::types::{CameraUniform, TextureResource, Vertex};

const PIPELINE_CACHE_SIZE: usize = 4096;
const MAX_SPRITES_PER_BATCH: usize = 1024;
const INDEX_PER_SPRITE: usize = 6;
#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
struct PipeLineCacheKey {
    shader: &'static str,
    transparent: bool,
}

impl From<BatchDrawCacheKey> for PipeLineCacheKey {
    fn from(key: BatchDrawCacheKey) -> Self {
        PipeLineCacheKey {
            shader: key.shader,
            transparent: key.transparent,
        }
    }
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct BatchDrawCacheKey {
    z_index: OrderedFloat<f32>,
    content: DrawContentCacheKey,
    shader: &'static str,
    transparent: bool,
}

impl BatchDrawCacheKey {
    fn from_draw(cmd: &Draw) -> Self {
        Self {
            content: cmd.content.clone().into(),
            z_index: OrderedFloat(cmd.transform.pos.z),
            shader: cmd.shader.key(),
            transparent: cmd.transparent,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PrimitiveKind {
    IndexedQuads,
    TriangleList,
}

struct BatchDraw {
    texture: Option<TextureResource>,
    vertices: Vec<Vertex>,
    count: usize,
    shader: Rc<dyn Shader>,
    shader_buffer: Vec<u8>,
    primitive: PrimitiveKind,
}

pub struct Render {
    device: Device,
    queue: Queue,
    color_format: TextureFormat,
    depth_format: TextureFormat,
    index_buffer: wgpu::Buffer,
    camera_buffer: wgpu::Buffer,
    camera_uniform: CameraUniform,
    pipeline_cache: LruCache<PipeLineCacheKey, Box<dyn SpritePipeline>>,
    batch_cache: HashMap<BatchDrawCacheKey, Vec<BatchDraw>>,
    transparent_batch_cache: HashMap<BatchDrawCacheKey, Vec<BatchDraw>>,
}

impl Render {
    pub fn setup(
        device: Device,
        queue: Queue,
        color_format: TextureFormat,
        depth_format: TextureFormat,
    ) -> Self {
        let index_buffer = Self::setup_index_buffer(&device);
        let pipeline_cache = LruCache::new(NonZeroUsize::new(PIPELINE_CACHE_SIZE).unwrap());
        let batch_cache = HashMap::default();
        let transparent_batch_cache = HashMap::default();
        let camera_buffer = Self::setup_camera_buffer(&device);
        let camera_uniform = CameraUniform::default();

        Self {
            device,
            queue,
            color_format,
            depth_format,
            index_buffer,
            camera_buffer,
            camera_uniform,
            pipeline_cache,
            batch_cache,
            transparent_batch_cache,
        }
    }

    fn setup_camera_buffer(device: &Device) -> wgpu::Buffer {
        device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Index Buffer"),
            contents: to_wgsl_bytes(&CameraUniform::default()).as_ref(),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        })
    }

    fn setup_index_buffer(device: &Device) -> wgpu::Buffer {
        let mut indexes = [0u16; MAX_SPRITES_PER_BATCH * INDEX_PER_SPRITE];
        for i in 0..(MAX_SPRITES_PER_BATCH as u16) {
            let start = i as usize * INDEX_PER_SPRITE;
            let end = start + INDEX_PER_SPRITE;
            indexes[start..end].copy_from_slice(&[
                i * 4,
                1 + i * 4,
                2 + i * 4,
                i * 4,
                2 + i * 4,
                3 + i * 4,
            ]);
        }
        device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Sprite Index Buffer"),
            contents: bytemuck::cast_slice(&indexes),
            usage: wgpu::BufferUsages::INDEX,
        })
    }

    pub fn draw(&mut self, cmd: Draw, texture: Option<TextureResource>) {
        match &cmd.content {
            DrawContent::Texture { .. } => self.draw_sprite(cmd, texture),
            DrawContent::Lines(_) => self.draw_lines_batch(cmd, texture),
            DrawContent::Circle { .. } => self.draw_circle(cmd, texture),
            DrawContent::Polygon { .. } | DrawContent::PolygonTextured { .. } => {
                self.draw_polygon(cmd, texture)
            }
        }
    }

    fn get_or_create_batch(
        &mut self,
        cmd: &Draw,
        texture: Option<TextureResource>,
    ) -> &mut BatchDraw {
        let key = BatchDrawCacheKey::from_draw(cmd);
        let cache = if cmd.transparent {
            &mut self.transparent_batch_cache
        } else {
            &mut self.batch_cache
        };
        let batches = cache.entry(key).or_default();
        if batches.is_empty()
            || batches
                .last()
                .is_some_and(|last| last.count >= MAX_SPRITES_PER_BATCH)
            || cmd.shader_data.is_some()
        {
            batches.push(BatchDraw {
                texture,
                vertices: Default::default(),
                count: 0,
                shader: cmd.shader.clone(),
                shader_buffer: Vec::new(),
                primitive: PrimitiveKind::IndexedQuads,
            });
        }

        batches.last_mut().unwrap()
    }

    fn draw_sprite(&mut self, cmd: Draw, texture: Option<TextureResource>) {
        let batch = self.get_or_create_batch(&cmd, texture);
        let size = cmd.size * cmd.transform.scale;
        let color = match &cmd.content {
            DrawContent::Texture { color, .. } => to_linear_rgba(*color),
            DrawContent::Circle { color, .. } => to_linear_rgba(*color),
            _ => to_linear_rgba(crate::color::WHITE),
        };
        let pos = cmd.transform.pos;
        let anchor = cmd.anchor;

        let min = pos.xy() - size * anchor;
        let max: Vec2 = pos.xy() + size * (Vec2::ONE - anchor);

        let (min_tex, max_tex) = match (cmd.src.clone(), &batch.texture) {
            (Some(src), Some(texture)) => {
                let size = texture.texture.size();
                let tex_size = Vec2::new(size.width as f32, size.height as f32);
                let half_pixel = Vec2::splat(0.5);
                (
                    (src.min + half_pixel) / tex_size,
                    (src.max - half_pixel) / tex_size,
                )
            }
            _ => (Vec2::ZERO, Vec2::ONE),
        };

        let mut vertices = [
            Vertex {
                pos: [min.x, min.y, pos.z],
                tex_coord: [min_tex.x, max_tex.y],
                color,
            },
            Vertex {
                pos: [max.x, min.y, pos.z],
                tex_coord: max_tex.to_array(),
                color,
            },
            Vertex {
                pos: [max.x, max.y, pos.z],
                tex_coord: [max_tex.x, min_tex.y],
                color,
            },
            Vertex {
                pos: [min.x, max.y, pos.z],
                tex_coord: min_tex.to_array(),
                color,
            },
        ];

        let mut swap_tex = |a: usize, b: usize, t: usize| {
            let v = vertices[a].tex_coord[t];
            vertices[a].tex_coord[t] = vertices[b].tex_coord[t];
            vertices[b].tex_coord[t] = v;
        };
        if cmd.flip_x {
            swap_tex(0, 1, 0);
            swap_tex(3, 2, 0);
        }
        if cmd.flip_y {
            swap_tex(0, 3, 1);
            swap_tex(2, 1, 1);
        }

        let angle = cmd.transform.angle();
        if angle != 0.0 {
            let angle = Vec2::from_angle(angle);
            for v in &mut vertices {
                let p = Vec2::from_slice(&v.pos[0..2]);
                let rotated = pos.xy() + angle.rotate(p - pos.xy());
                v.pos[0] = rotated.x;
                v.pos[1] = rotated.y;
            }
        }

        batch.vertices.extend_from_slice(&vertices);
        if let Some(data) = cmd.shader_data {
            batch.shader_buffer.extend(data);
        }
        batch.count += 1;
    }

    fn draw_lines_batch(&mut self, cmd: Draw, texture: Option<TextureResource>) {
        if let DrawContent::Lines(lines) = &cmd.content {
            let batch = self.get_or_create_batch(&cmd, texture.clone());
            for line in lines {
                let color = to_linear_rgba(line.color);
                let line_dir = (line.end - line.start).normalize_or_zero();
                let perp = Vec2::new(-line_dir.y, line_dir.x) * (line.thickness * 0.5);

                let start_pos = cmd.transform.pos;
                let vertices = [
                    Vertex {
                        pos: (line.start - perp).extend(start_pos.z).to_array(),
                        tex_coord: [0.0, 0.0],
                        color,
                    },
                    Vertex {
                        pos: (line.end - perp).extend(start_pos.z).to_array(),
                        tex_coord: [1.0, 0.0],
                        color,
                    },
                    Vertex {
                        pos: (line.end + perp).extend(start_pos.z).to_array(),
                        tex_coord: [1.0, 1.0],
                        color,
                    },
                    Vertex {
                        pos: (line.start + perp).extend(start_pos.z).to_array(),
                        tex_coord: [0.0, 1.0],
                        color,
                    },
                ];
                batch.vertices.extend_from_slice(&vertices);
                batch.count += 1;
            }
        }
    }

    fn draw_circle(&mut self, cmd: Draw, texture: Option<TextureResource>) {
        self.draw_sprite(cmd, texture)
    }

    fn draw_polygon(&mut self, cmd: Draw, texture: Option<TextureResource>) {
        let key = BatchDrawCacheKey::from_draw(&cmd);
        let cache = if cmd.transparent {
            &mut self.transparent_batch_cache
        } else {
            &mut self.batch_cache
        };
        let batches = cache.entry(key).or_default();
        let Draw {
            shader,
            mut shader_data,
            content,
            transform,
            ..
        } = cmd;

        match content {
            DrawContent::Polygon { vertices, color } => {
                if vertices.len() < 3 {
                    return;
                }

                let mut batch = BatchDraw {
                    texture: None,
                    vertices: Vec::with_capacity((vertices.len() - 2) * 3),
                    count: 0,
                    shader,
                    shader_buffer: shader_data.take().unwrap_or_default(),
                    primitive: PrimitiveKind::TriangleList,
                };

                let color = to_linear_rgba(color);
                let pivot = vertices[0];
                for i in 1..vertices.len() - 1 {
                    for point in [pivot, vertices[i], vertices[i + 1]] {
                        batch.vertices.push(Vertex {
                            pos: point.extend(transform.pos.z).to_array(),
                            tex_coord: [0.0, 0.0],
                            color,
                        });
                    }
                }
                batches.push(batch);
            }
            DrawContent::PolygonTextured {
                vertices,
                uvs,
                color,
                ..
            } => {
                if vertices.len() < 3 || uvs.len() != vertices.len() {
                    return;
                }
                let Some(texture) = texture else {
                    return;
                };
                let mut batch = BatchDraw {
                    texture: Some(texture),
                    vertices: Vec::with_capacity((vertices.len() - 2) * 3),
                    count: 0,
                    shader,
                    shader_buffer: shader_data.take().unwrap_or_default(),
                    primitive: PrimitiveKind::TriangleList,
                };

                let color = to_linear_rgba(color);
                let pivot = vertices[0];
                let pivot_uv = uvs[0];
                for i in 1..vertices.len() - 1 {
                    for (point, uv) in [
                        (pivot, pivot_uv),
                        (vertices[i], uvs[i]),
                        (vertices[i + 1], uvs[i + 1]),
                    ] {
                        batch.vertices.push(Vertex {
                            pos: point.extend(transform.pos.z).to_array(),
                            tex_coord: uv.to_array(),
                            color,
                        });
                    }
                }
                batches.push(batch);
            }
            _ => {}
        }
    }

    pub fn frame_pass(&mut self, view_proj: Mat4) {
        self.camera_uniform.update(view_proj);
        self.queue
            .write_buffer(&self.camera_buffer, 0, &to_wgsl_bytes(&self.camera_uniform));
    }

    pub fn frame_end(&mut self, pass: &mut RenderPass) {
        // Opaque objects
        let batch_cache: Vec<_> = std::mem::take(&mut self.batch_cache).into_iter().collect();
        for (key, batches) in batch_cache {
            let pipeline_key = PipeLineCacheKey::from(key.clone());
            for batch in batches {
                let sp = self
                    .pipeline_cache
                    .get_or_insert(pipeline_key.clone(), || {
                        batch.shader.new_pipeline(
                            self.color_format,
                            Some(self.depth_format),
                            &self.camera_buffer,
                            false,
                        )
                    })
                    .as_ref();
                write_batch(&self.device, &self.index_buffer, pass, &batch, sp);
            }
        }

        // Transparent objects
        let mut transparent_batch_cache: Vec<_> = std::mem::take(&mut self.transparent_batch_cache)
            .into_iter()
            .collect();
        transparent_batch_cache.sort_by(|(a, _), (b, _)| b.z_index.cmp(&a.z_index));

        for (key, batches) in transparent_batch_cache {
            let pipeline_key = PipeLineCacheKey::from(key.clone());
            for batch in batches {
                let sp = self
                    .pipeline_cache
                    .get_or_insert(pipeline_key.clone(), || {
                        batch.shader.new_pipeline(
                            self.color_format,
                            Some(self.depth_format),
                            &self.camera_buffer,
                            true,
                        )
                    })
                    .as_ref();
                write_batch(&self.device, &self.index_buffer, pass, &batch, sp);
            }
        }
    }
}

fn write_batch(
    device: &Device,
    index_buffer: &wgpu::Buffer,
    pass: &mut RenderPass,
    batch: &BatchDraw,
    sp: &dyn SpritePipeline,
) {
    let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("Vertex Buffer"),
        contents: bytemuck::cast_slice(&batch.vertices),
        usage: wgpu::BufferUsages::VERTEX,
    });

    pass.set_vertex_buffer(0, vertex_buffer.slice(..));
    let shader_data_opt = if batch.shader_buffer.is_empty() {
        None
    } else {
        Some(batch.shader_buffer.as_slice())
    };
    sp.render(pass, batch.texture.as_ref(), shader_data_opt);
    match batch.primitive {
        PrimitiveKind::IndexedQuads => {
            pass.set_index_buffer(index_buffer.slice(..), wgpu::IndexFormat::Uint16);
            let index_count = (batch.count * INDEX_PER_SPRITE) as u32;
            pass.draw_indexed(0..index_count, 0, 0..1);
        }
        PrimitiveKind::TriangleList => {
            pass.draw(0..batch.vertices.len() as u32, 0..1);
        }
    }
}
