use crate::scene::Scene;
use crate::workload::Workload;
use crate::{PreviewPresentMode, PreviewScenePreset, SimState};
use anyhow::Result;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use winit::{
    dpi::LogicalSize,
    event::{Event, WindowEvent},
    event_loop::{ControlFlow, EventLoop},
    window::Window,
};

const SHADER_SRC: &str = r#"
struct Camera {
    view_proj: mat4x4<f32>,
    time_pad: vec4<f32>,
}
@group(0) @binding(0) var<uniform> camera: Camera;

struct VertIn {
    @location(0) position: vec3<f32>,
    @location(1) color:    vec4<f32>,
    @location(2) uv:       vec2<f32>,
}

struct VertOut {
    @builtin(position) clip_pos: vec4<f32>,
    @location(0) color: vec4<f32>,
    @location(2) local_uv: vec2<f32>,
}

@vertex
fn vs_main(in: VertIn) -> VertOut {
    var out: VertOut;
    out.clip_pos = camera.view_proj * vec4<f32>(in.position, 1.0);
    out.color    = in.color;
    out.local_uv = in.uv;
    return out;
}

fn hash3(p: vec3<f32>) -> f32 {
    var p3 = fract(p * 0.1031);
    p3 = p3 + vec3<f32>(dot(p3, p3.yzx + vec3<f32>(33.33)));
    return fract((p3.x + p3.y) * p3.z);
}

fn value_noise(p: vec3<f32>) -> f32 {
    let i = floor(p);
    let f = fract(p);
    let u = f * f * (vec3<f32>(3.0) - 2.0 * f);

    return mix(
        mix(mix(hash3(i + vec3<f32>(0.0, 0.0, 0.0)), hash3(i + vec3<f32>(1.0, 0.0, 0.0)), u.x),
            mix(hash3(i + vec3<f32>(0.0, 1.0, 0.0)), hash3(i + vec3<f32>(1.0, 1.0, 0.0)), u.x), u.y),
        mix(mix(hash3(i + vec3<f32>(0.0, 0.0, 1.0)), hash3(i + vec3<f32>(1.0, 0.0, 1.0)), u.x),
            mix(hash3(i + vec3<f32>(0.0, 1.0, 1.0)), hash3(i + vec3<f32>(1.0, 1.0, 1.0)), u.x), u.y),
        u.z
    );
}

@fragment
fn fs_main(in: VertOut) -> @location(0) vec4<f32> {
    let uv = in.local_uv;
    let r2 = dot(uv, uv);
    if (r2 > 1.0) {
        discard;
    }

    let depth = sqrt(1.0 - r2);
    let ray_entry = -depth;
    let ray_exit = depth;
    let ray_length = 2.0 * depth;
    let step_size = ray_length / 64.0;
    var density = 0.0;

    for (var step = 0u; step < 64u; step = step + 1u) {
        let t = ray_entry + (f32(step) + 0.5) * step_size;
        let p = vec3<f32>(uv, t);
        let time = camera.time_pad.x;
        let n1 = value_noise(p * 3.0 + vec3<f32>(time * 0.3, 0.0, 0.0));
        let n2 = value_noise(p * 6.0 - vec3<f32>(0.0, time * 0.5, 0.0)) * 0.5;
        let n3 = value_noise(p * 12.0 + vec3<f32>(time * 0.7, time * 0.2, 0.0)) * 0.25;
        let sample = (n1 + n2 + n3) / 1.75;
        density = density + sample * (1.0 / 64.0);
    }

    return vec4<f32>(in.color.rgb, clamp(density, 0.0, 1.0));
}
"#;

type Mat4 = [f32; 16];

fn mat4_identity() -> Mat4 {
    [
        1., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1.,
    ]
}

fn mat4_mul(a: &Mat4, b: &Mat4) -> Mat4 {
    let mut out = [0f32; 16];
    for row in 0..4 {
        for col in 0..4 {
            out[col * 4 + row] = (0..4).map(|k| a[k * 4 + row] * b[col * 4 + k]).sum();
        }
    }
    out
}

fn perspective(fov_y_rad: f32, aspect: f32, near: f32, far: f32) -> Mat4 {
    let f = 1.0 / (fov_y_rad * 0.5).tan();
    [
        f / aspect,
        0.,
        0.,
        0.,
        0.,
        -f,
        0.,
        0.,
        0.,
        0.,
        far / (near - far),
        -1.,
        0.,
        0.,
        near * far / (near - far),
        0.,
    ]
}

fn look_at(eye: [f32; 3], center: [f32; 3], up: [f32; 3]) -> Mat4 {
    let sub = |a: [f32; 3], b: [f32; 3]| [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
    let norm = |v: [f32; 3]| {
        let l = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt();
        [v[0] / l, v[1] / l, v[2] / l]
    };
    let cross = |a: [f32; 3], b: [f32; 3]| {
        [
            a[1] * b[2] - a[2] * b[1],
            a[2] * b[0] - a[0] * b[2],
            a[0] * b[1] - a[1] * b[0],
        ]
    };
    let dot = |a: [f32; 3], b: [f32; 3]| a[0] * b[0] + a[1] * b[1] + a[2] * b[2];

    let f = norm(sub(center, eye));
    let s = norm(cross(f, up));
    let u = cross(s, f);

    let mut m = mat4_identity();
    m[0] = s[0];
    m[4] = s[1];
    m[8] = s[2];
    m[1] = u[0];
    m[5] = u[1];
    m[9] = u[2];
    m[2] = -f[0];
    m[6] = -f[1];
    m[10] = -f[2];
    m[12] = -dot(s, eye);
    m[13] = -dot(u, eye);
    m[14] = dot(f, eye);
    m
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct CameraUniform {
    view_proj: Mat4,
    time_pad: [f32; 4],
}

struct GpuState {
    surface: wgpu::Surface<'static>,
    device: wgpu::Device,
    queue: wgpu::Queue,
    config: wgpu::SurfaceConfiguration,
    supported_present_modes: Vec<wgpu::PresentMode>,
    pipeline: wgpu::RenderPipeline,
    camera_buffer: wgpu::Buffer,
    camera_bind_group: wgpu::BindGroup,
    scene: Scene,
    workload: Workload,
}

impl GpuState {
    async fn new(window: Arc<Window>) -> Result<Self> {
        let size = window.inner_size();
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
            backends: wgpu::Backends::VULKAN,
            flags: wgpu::InstanceFlags::empty(),
            ..Default::default()
        });

        let surface = instance.create_surface(Arc::clone(&window))?;
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                compatible_surface: Some(&surface),
                force_fallback_adapter: false,
            })
            .await
            .ok_or_else(|| anyhow::anyhow!("No suitable Vulkan adapter found"))?;

        let (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    label: Some("mangotune_preview"),
                    required_features: wgpu::Features::empty(),
                    required_limits: wgpu::Limits::default(),
                    memory_hints: wgpu::MemoryHints::default(),
                },
                None,
            )
            .await?;

        let surface_caps = surface.get_capabilities(&adapter);
        let format = surface_caps
            .formats
            .iter()
            .find(|f| f.is_srgb())
            .copied()
            .unwrap_or(surface_caps.formats[0]);

        let config = wgpu::SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format,
            width: size.width,
            height: size.height,
            present_mode: choose_present_mode(
                &surface_caps.present_modes,
                PreviewPresentMode::from_env(),
            ),
            alpha_mode: surface_caps.alpha_modes[0],
            view_formats: vec![],
            desired_maximum_frame_latency: 2,
        };
        surface.configure(&device, &config);

        let camera_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("camera_uniform"),
            size: std::mem::size_of::<CameraUniform>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

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

        let camera_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("camera_bg"),
            layout: &camera_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: camera_buffer.as_entire_binding(),
            }],
        });

        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("particle_shader"),
            source: wgpu::ShaderSource::Wgsl(SHADER_SRC.into()),
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("pipeline_layout"),
            bind_group_layouts: &[&camera_bind_group_layout],
            push_constant_ranges: &[],
        });

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("particle_pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: "vs_main",
                buffers: &[crate::scene::ParticleVertex::desc()],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: "fs_main",
                targets: &[Some(wgpu::ColorTargetState {
                    format: config.format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });

        let scene = Scene::new(&device);
        let workload = Workload::new(&device);

        Ok(Self {
            surface,
            device,
            queue,
            config,
            supported_present_modes: surface_caps.present_modes.clone(),
            pipeline,
            camera_buffer,
            camera_bind_group,
            scene,
            workload,
        })
    }

    fn resize(&mut self, new_size: winit::dpi::PhysicalSize<u32>) {
        if new_size.width > 0 && new_size.height > 0 {
            self.config.width = new_size.width;
            self.config.height = new_size.height;
            self.surface.configure(&self.device, &self.config);
        }
    }

    fn apply_runtime_mode(&mut self, state: &SimState) {
        let desired = choose_present_mode(&self.supported_present_modes, state.present_mode);
        if self.config.present_mode != desired {
            self.config.present_mode = desired;
            self.surface.configure(&self.device, &self.config);
        }
    }

    fn render(&mut self, state: &SimState) -> Result<(), wgpu::SurfaceError> {
        self.apply_runtime_mode(state);
        let aspect = self.config.width as f32 / self.config.height as f32;
        let t = self.scene.time;
        let eye = match state.scene_preset {
            PreviewScenePreset::DarkArena => [t.cos() * 5.8, 1.8, t.sin() * 5.8],
            PreviewScenePreset::BrightWash => [t.cos() * 4.6, 1.6, t.sin() * 4.6],
            PreviewScenePreset::MotionStress => [
                (t * 0.95).cos() * 6.8,
                2.2 + (t * 0.7).sin() * 0.8,
                (t * 0.95).sin() * 6.8,
            ],
            PreviewScenePreset::StaticInspection => [0.0, 2.4, 7.2],
            PreviewScenePreset::NoiseField => [
                (t * 0.45).cos() * 6.2,
                1.0 + (t * 0.25).sin() * 0.5,
                (t * 0.6).sin() * 6.2,
            ],
        };
        let view = look_at(eye, [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]);
        let proj = perspective(std::f32::consts::FRAC_PI_4, aspect, 0.1, 100.0);
        let view_proj = mat4_mul(&proj, &view);
        self.queue.write_buffer(
            &self.camera_buffer,
            0,
            bytemuck::bytes_of(&CameraUniform {
                view_proj,
                time_pad: [self.scene.time, 0.0, 0.0, 0.0],
            }),
        );

        let output = self.surface.get_current_texture()?;
        let view = output.texture.create_view(&Default::default());

        let clear = match state.scene_preset {
            PreviewScenePreset::DarkArena => wgpu::Color {
                r: 0.03,
                g: 0.04,
                b: 0.09,
                a: 1.0,
            },
            PreviewScenePreset::BrightWash => wgpu::Color {
                r: 0.78,
                g: 0.82,
                b: 0.82,
                a: 1.0,
            },
            PreviewScenePreset::MotionStress => wgpu::Color {
                r: 0.08,
                g: 0.05,
                b: 0.11,
                a: 1.0,
            },
            PreviewScenePreset::StaticInspection => wgpu::Color {
                r: 0.10,
                g: 0.11,
                b: 0.14,
                a: 1.0,
            },
            PreviewScenePreset::NoiseField => wgpu::Color {
                r: 0.04,
                g: 0.08,
                b: 0.09,
                a: 1.0,
            },
        };

        for pass_index in 0..state.gpu_passes {
            let mut encoder = self
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("frame_pass"),
                });

            {
                let load_op = if pass_index == 0 {
                    wgpu::LoadOp::Clear(clear)
                } else {
                    wgpu::LoadOp::Load
                };

                let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                    label: Some("particle_pass"),
                    color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                        view: &view,
                        resolve_target: None,
                        ops: wgpu::Operations {
                            load: load_op,
                            store: wgpu::StoreOp::Store,
                        },
                    })],
                    depth_stencil_attachment: None,
                    timestamp_writes: None,
                    occlusion_query_set: None,
                });
                rpass.set_pipeline(&self.pipeline);
                rpass.set_bind_group(0, &self.camera_bind_group, &[]);
                rpass.set_vertex_buffer(0, self.scene.vertex_buffer.slice(..));
                rpass.draw(0..self.scene.vertex_count(), 0..1);
            }

            self.queue.submit(std::iter::once(encoder.finish()));
        }
        output.present();
        self.workload.update(&self.device, state);
        Ok(())
    }
}

fn choose_present_mode(
    supported: &[wgpu::PresentMode],
    requested: PreviewPresentMode,
) -> wgpu::PresentMode {
    match requested {
        PreviewPresentMode::Fifo => supported
            .iter()
            .copied()
            .find(|mode| *mode == wgpu::PresentMode::Fifo)
            .unwrap_or(wgpu::PresentMode::Fifo),
        PreviewPresentMode::Immediate => supported
            .iter()
            .copied()
            .find(|mode| {
                matches!(
                    mode,
                    wgpu::PresentMode::Immediate
                        | wgpu::PresentMode::Mailbox
                        | wgpu::PresentMode::AutoNoVsync
                )
            })
            .unwrap_or(wgpu::PresentMode::Fifo),
    }
}

#[allow(deprecated)]
pub fn run(state: Arc<Mutex<SimState>>) -> Result<()> {
    let event_loop = EventLoop::new()?;
    event_loop.set_control_flow(ControlFlow::Poll);

    let width = std::env::var("MANGOTUNE_PREVIEW_WIDTH")
        .ok()
        .and_then(|value| value.parse::<u32>().ok())
        .unwrap_or(800);
    let height = std::env::var("MANGOTUNE_PREVIEW_HEIGHT")
        .ok()
        .and_then(|value| value.parse::<u32>().ok())
        .unwrap_or(600);

    let window = Arc::new(
        event_loop.create_window(
            Window::default_attributes()
                .with_title("MangoTune Preview")
                .with_inner_size(LogicalSize::new(width, height)),
        )?,
    );

    let mut gpu = pollster::block_on(GpuState::new(Arc::clone(&window)))?;
    let mut last_frame = Instant::now();

    event_loop.run(move |event, elwt| {
        {
            let sim = state.lock().expect("preview state");
            if sim.should_quit {
                elwt.exit();
                return;
            }
        }

        match event {
            Event::WindowEvent {
                event: WindowEvent::CloseRequested,
                ..
            } => {
                elwt.exit();
            }
            Event::WindowEvent {
                event: WindowEvent::Resized(size),
                ..
            } => {
                gpu.resize(size);
            }
            Event::AboutToWait => {
                let now = Instant::now();
                let dt = now.duration_since(last_frame).as_secs_f32().min(0.05);
                last_frame = now;

                let sim = state.lock().expect("preview state").clone();

                gpu.scene.update(dt, sim.particle_size, &sim, &gpu.queue);
                match gpu.render(&sim) {
                    Ok(_) => {}
                    Err(wgpu::SurfaceError::Lost | wgpu::SurfaceError::Outdated) => {
                        gpu.resize(window.inner_size());
                    }
                    Err(wgpu::SurfaceError::OutOfMemory) => elwt.exit(),
                    Err(err) => log::warn!("Surface error: {}", err),
                }
            }
            _ => {}
        }
    })?;

    let _ = std::fs::remove_file(crate::socket_api::socket_path());
    Ok(())
}