use std::error::Error;
use std::str::FromStr;

use rosc::OscType;
use wesl::{Mangler, ModulePath};

use crate::uniform::{UniformCache, UniformError};
use crate::wesl::ShaderSources;

const CANVAS_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8UnormSrgb;

const VERTEX_SHADER: &str = "\
@group(0) @binding(0) var<uniform> _wgsl_resolution: vec2<f32>;

struct VertexOutput {
    @builtin(position) clip_position: vec4<f32>,
    @location(0) uv: vec2<f32>,
    @location(1) resolution: vec2<f32>,
};

@vertex
fn vs_main(@builtin(vertex_index) vi: u32) -> VertexOutput {
    var uv = vec2<f32>(f32(vi % 2u), f32(vi / 2u));
    var out: VertexOutput;
    out.clip_position = vec4<f32>(2.0 * uv - 1.0, 0.0, 1.0);
    out.uv = vec2<f32>(uv.x, 1.0 - uv.y);
    out.resolution = _wgsl_resolution;
    return out;
}
";

const DEFAULT_FRAGMENT: &str = "\
@fragment
fn fs_main(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
    let check = floor(uv * 10.0);
    return vec4<f32>(vec3<f32>(fract((check.x + check.y) / 2.0)), 1.0);
}
";

#[derive(Debug, thiserror::Error)]
pub enum ShaderError {
    #[error("error resolving module")]
    InvalidModulePath,

    #[error("error setting uniform '{0}': {1}")]
    UniformError(String, crate::uniform::UniformError),

    #[error("uniform '{0}' not found")]
    UniformNotFound(String),

    #[error("error binding resource to '{0}': {1}")]
    BindError(String, crate::uniform::BindError),

    #[error("error compiling shader: {0}")]
    CompileError(#[from] wesl::Error),
}

pub struct Renderer {
    canvas: wgpu::Texture,
    canvas_view: wgpu::TextureView,
    vertex_module: wgpu::ShaderModule,
    vertex_bgl: wgpu::BindGroupLayout,
    vertex_bg: wgpu::BindGroup,
    render_pipeline: wgpu::RenderPipeline,
    shader_sources: ShaderSources,
    uniforms: UniformCache,
}

impl Renderer {
    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue, width: u32, height: u32) -> Self {
        let canvas = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("canvas"),
            size: wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: CANVAS_FORMAT,
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT
                | wgpu::TextureUsages::TEXTURE_BINDING
                | wgpu::TextureUsages::COPY_SRC,
            view_formats: &[],
        });
        let canvas_view = canvas.create_view(&Default::default());

        let vertex_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("vertex_shader"),
            source: wgpu::ShaderSource::Wgsl(VERTEX_SHADER.into()),
        });

        let resolution_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("resolution"),
            size: 8,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let vertex_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: None,
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        });

        let vertex_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &vertex_bgl,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: resolution_buffer.as_entire_binding(),
            }],
        });

        queue.write_buffer(
            &resolution_buffer,
            0,
            bytemuck::cast_slice(&[width as f32, height as f32]),
        );

        let mut uniforms = UniformCache::new();
        let mut shader_sources = ShaderSources::default();
        shader_sources.set_module(wesl::ModulePath::new_root(), DEFAULT_FRAGMENT);

        let initial_fragment = shader_sources
            .compile(&wesl::ModulePath::new_root())
            .expect("default shader");
        log::info!("frag source:\n {initial_fragment}");

        let render_pipeline = build_pipeline(
            device,
            &vertex_module,
            &vertex_bgl,
            &initial_fragment,
            &mut uniforms,
        )
        .expect("default shader");

        Self {
            canvas,
            canvas_view,
            vertex_module,
            vertex_bgl,
            vertex_bg,
            render_pipeline,
            shader_sources,
            uniforms,
        }
    }

    pub fn canvas_texture(&self) -> &wgpu::Texture {
        &self.canvas
    }

    pub fn uniforms(&mut self) -> &mut UniformCache {
        &mut self.uniforms
    }

    pub fn set_uniform(
        &mut self,
        module: &str,
        item: &str,
        rest: &[&str],
        args: &[OscType],
    ) -> Result<(), ShaderError> {
        let path = ModulePath::from_str(module).map_err(|_| ShaderError::InvalidModulePath)?;
        let name = self.shader_sources.mangler().mangle(&path, item);

        let mut uref = self
            .uniforms
            .get(&name)
            .ok_or(ShaderError::UniformNotFound(name))?;

        let lift_err = |e| ShaderError::UniformError(item.to_string(), e);

        for component in rest.iter() {
            uref = uref.field(component).map_err(lift_err)?;
        }

        let scalar = uref.leaf_scalar().map_err(lift_err)?;
        match scalar.kind {
            naga::ScalarKind::Float => args
                .iter()
                .map(|a| match a {
                    OscType::Float(f) => Ok(*f),
                    OscType::Double(d) => Ok(*d as f32),
                    OscType::Int(i) => Ok(*i as f32),
                    OscType::Bool(b) => Ok(if *b { 1.0 } else { 0.0 }),
                    _ => Err(UniformError::TypeMismatch),
                })
                .collect::<Result<Vec<_>, _>>()
                .and_then(|values| uref.set_f32(&values)),
            naga::ScalarKind::Sint => args
                .iter()
                .map(|a| match a {
                    OscType::Int(i) => Ok(*i),
                    OscType::Float(f) => Ok(*f as i32),
                    OscType::Double(d) => Ok(*d as i32),
                    OscType::Bool(b) => Ok(if *b { 1 } else { 0 }),
                    _ => Err(UniformError::TypeMismatch),
                })
                .collect::<Result<Vec<_>, _>>()
                .and_then(|values| uref.set_i32(&values)),
            naga::ScalarKind::Uint => args
                .iter()
                .map(|a| match a {
                    OscType::Int(i) => Ok(*i as u32),
                    OscType::Float(f) => Ok(*f as u32),
                    OscType::Double(d) => Ok(*d as u32),
                    OscType::Bool(b) => Ok(if *b { 1 } else { 0 }),
                    _ => Err(UniformError::TypeMismatch),
                })
                .collect::<Result<Vec<_>, _>>()
                .and_then(|values| uref.set_u32(&values)),
            _ => Err(UniformError::TypeMismatch),
        }
        .map_err(lift_err)
    }

    pub fn set_binding(
        &mut self,
        device: &wgpu::Device,
        module: &str,
        item: &str,
        target: &str,
    ) -> Result<(), ShaderError> {
        let path = ModulePath::from_str(module).map_err(|_| ShaderError::InvalidModulePath)?;
        let name = self.shader_sources.mangler().mangle(&path, item);

        if let Some(id) = target.strip_prefix("/texture/") {
            self.uniforms.bind_texture(device, &name, id)
        } else if let Some(id) = target.strip_prefix("/sampler/") {
            self.uniforms.bind_sampler(device, &name, id)
        } else {
            Err(crate::uniform::BindError::ResourceNotFound(
                target.to_string(),
            ))
        }
        .map_err(|e| ShaderError::BindError(name, e))
    }

    pub fn load_module(&mut self, module_name: &str, source: &str) -> Result<(), ShaderError> {
        let path = ModulePath::from_str(module_name).map_err(|_| ShaderError::InvalidModulePath)?;
        self.shader_sources.set_module(path, source);
        Ok(())
    }

    pub fn compile(
        &mut self,
        device: &wgpu::Device,
        module_name: &str,
    ) -> Result<(), Box<dyn Error>> {
        let path = ModulePath::from_str(module_name).map_err(|_| ShaderError::InvalidModulePath)?;
        let fragment_source = self.shader_sources.compile(&path)?;
        log::info!("frag source:\n {fragment_source}");

        self.render_pipeline = build_pipeline(
            device,
            &self.vertex_module,
            &self.vertex_bgl,
            &fragment_source,
            &mut self.uniforms,
        )?;
        Ok(())
    }

    pub fn render(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
        self.uniforms.flush(queue);

        let mut encoder =
            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });

        {
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: None,
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &self.canvas_view,
                    depth_slice: None,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                occlusion_query_set: None,
                timestamp_writes: None,
                multiview_mask: None,
            });

            pass.set_pipeline(&self.render_pipeline);
            pass.set_bind_group(0, &self.vertex_bg, &[]);
            if let Some(ref bg) = self.uniforms.bind_group {
                pass.set_bind_group(1, bg, &[]);
            }
            pass.draw(0..4, 0..1);
        }

        queue.submit([encoder.finish()]);
    }
}

fn build_pipeline(
    device: &wgpu::Device,
    vertex_module: &wgpu::ShaderModule,
    vertex_bgl: &wgpu::BindGroupLayout,
    fragment_source: &str,
    cache: &mut UniformCache,
) -> Result<wgpu::RenderPipeline, Box<dyn Error>> {
    let module = naga::front::wgsl::parse_str(fragment_source)?;

    naga::valid::Validator::new(
        naga::valid::ValidationFlags::all(),
        naga::valid::Capabilities::all(),
    )
    .validate(&module)?;

    cache.refresh(module, device);
    cache.rebuild_bind_group(device);

    let fragment_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("fragment_shader"),
        source: wgpu::ShaderSource::Wgsl(fragment_source.into()),
    });

    let mut bind_group_layouts: Vec<Option<&wgpu::BindGroupLayout>> = vec![Some(vertex_bgl)];
    if let Some(ref bgl) = cache.bind_group_layout {
        bind_group_layouts.push(Some(bgl));
    }

    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: None,
        bind_group_layouts: &bind_group_layouts,
        immediate_size: 0,
    });

    Ok(
        device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: None,
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: vertex_module,
                entry_point: None,
                buffers: &[],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &fragment_module,
                entry_point: None,
                targets: &[Some(wgpu::ColorTargetState {
                    format: CANVAS_FORMAT,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleStrip,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: Default::default(),
            multiview_mask: None,
            cache: None,
        }),
    )
}