use iced::widget::shader;
use iced::wgpu;
use iced::mouse;
use iced::{Event, Rectangle};
use iced::Point;

use cord_trig::TrigGraph;
use cord_shader::generate_wgsl_from_trig;

pub struct SdfViewport {
    wgsl: String,
    generation: u64,
    camera: Camera,
    bounding_radius: f64,
    drag_state: DragState,
    pub render_flags: RenderFlags,
}

#[derive(Debug, Clone, Copy)]
pub struct RenderFlags {
    pub shadows: bool,
    pub ao: bool,
    pub ground: bool,
}

impl Default for RenderFlags {
    fn default() -> Self {
        Self { shadows: true, ao: true, ground: true }
    }
}

struct Camera {
    yaw: f32,
    pitch: f32,
    distance: f32,
    target: [f32; 3],
    fov: f32,
}

#[derive(Default)]
struct DragState {
    active: bool,
    last: Option<Point>,
}

impl Camera {
    fn position(&self) -> [f32; 3] {
        let (sy, cy) = self.yaw.sin_cos();
        let (sp, cp) = self.pitch.sin_cos();
        [
            self.target[0] + self.distance * cp * cy,
            self.target[1] + self.distance * cp * sy,
            self.target[2] + self.distance * sp,
        ]
    }
}

impl SdfViewport {
    pub fn new() -> Self {
        let default = default_graph();
        Self {
            wgsl: generate_wgsl_from_trig(&default),
            generation: 0,
            camera: Camera {
                yaw: 0.6,
                pitch: 0.4,
                distance: 8.0,
                target: [0.0, 0.0, 0.0],
                fov: 1.0,
            },
            bounding_radius: 2.0,
            drag_state: DragState::default(),
            render_flags: RenderFlags::default(),
        }
    }

    pub fn set_graph(&mut self, graph: &TrigGraph) {
        self.wgsl = generate_wgsl_from_trig(graph);
        self.generation += 1;
    }

    pub fn set_bounds(&mut self, radius: f64) {
        let radius = radius.max(0.5);
        self.bounding_radius = radius;
        self.camera.distance = (radius as f32) * 3.0;
    }
}

fn default_graph() -> TrigGraph {
    use cord_trig::ir::{TrigGraph, TrigOp};
    let mut g = TrigGraph::new();
    let x = g.push(TrigOp::InputX);
    let y = g.push(TrigOp::InputY);
    let z = g.push(TrigOp::InputZ);
    let xy = g.push(TrigOp::Hypot(x, y));
    let mag = g.push(TrigOp::Hypot(xy, z));
    let r = g.push(TrigOp::Const(2.0));
    let out = g.push(TrigOp::Sub(mag, r));
    g.set_output(out);
    g
}

#[derive(Debug)]
pub struct ViewportPrimitive {
    wgsl: String,
    generation: u64,
    time: f32,
    camera_pos: [f32; 3],
    camera_target: [f32; 3],
    fov: f32,
    render_flags: [f32; 4],
    scene_scale: f32,
}

#[derive(Default)]
pub struct ViewportState {
    start: Option<std::time::Instant>,
}

impl<Message: 'static> shader::Program<Message> for SdfViewport {
    type State = ViewportState;
    type Primitive = ViewportPrimitive;

    fn draw(
        &self,
        state: &Self::State,
        _cursor: mouse::Cursor,
        _bounds: Rectangle,
    ) -> Self::Primitive {
        let elapsed = state.start
            .map(|s| s.elapsed().as_secs_f32())
            .unwrap_or(0.0);

        let rf = &self.render_flags;
        ViewportPrimitive {
            wgsl: self.wgsl.clone(),
            generation: self.generation,
            time: elapsed,
            camera_pos: self.camera.position(),
            camera_target: self.camera.target,
            fov: self.camera.fov,
            render_flags: [
                if rf.shadows { 1.0 } else { 0.0 },
                if rf.ao { 1.0 } else { 0.0 },
                if rf.ground { 1.0 } else { 0.0 },
                0.0,
            ],
            scene_scale: self.bounding_radius as f32,
        }
    }

    fn update(
        &self,
        state: &mut Self::State,
        _event: &Event,
        _bounds: Rectangle,
        _cursor: mouse::Cursor,
    ) -> Option<shader::Action<Message>> {
        if state.start.is_none() {
            state.start = Some(std::time::Instant::now());
        }
        None
    }

    fn mouse_interaction(
        &self,
        _state: &Self::State,
        bounds: Rectangle,
        cursor: mouse::Cursor,
    ) -> mouse::Interaction {
        if cursor.is_over(bounds) {
            if self.drag_state.active {
                mouse::Interaction::Grabbing
            } else {
                mouse::Interaction::Grab
            }
        } else {
            mouse::Interaction::default()
        }
    }
}

impl SdfViewport {
    pub fn on_drag(&mut self, dx: f32, dy: f32) {
        let sensitivity = 0.005;
        self.camera.yaw -= dx * sensitivity;
        self.camera.pitch = (self.camera.pitch - dy * sensitivity)
            .clamp(-1.4, 1.4);
    }

    pub fn on_scroll(&mut self, delta: f32) {
        let factor = (-delta * 0.08).exp();
        let max_dist = (self.bounding_radius as f32 * 20.0).max(100.0);
        let min_dist = (self.bounding_radius as f32 * 0.05).max(0.1);
        self.camera.distance = (self.camera.distance * factor)
            .clamp(min_dist, max_dist);
    }

    pub fn start_drag(&mut self, pos: Point) {
        self.drag_state.active = true;
        self.drag_state.last = Some(pos);
    }

    pub fn drag_to(&mut self, pos: Point) -> bool {
        if let Some(last) = self.drag_state.last {
            let dx = pos.x - last.x;
            let dy = pos.y - last.y;
            self.on_drag(dx, dy);
            self.drag_state.last = Some(pos);
            true
        } else {
            false
        }
    }

    pub fn end_drag(&mut self) {
        self.drag_state.active = false;
        self.drag_state.last = None;
    }

    pub fn is_dragging(&self) -> bool {
        self.drag_state.active
    }

    pub fn reset_camera(&mut self) {
        self.camera.target = [0.0, 0.0, 0.0];
        self.camera.yaw = 0.6;
        self.camera.pitch = 0.4;
        self.camera.distance = (self.bounding_radius as f32) * 3.0;
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct Uniforms {
    resolution: [f32; 2],
    viewport_offset: [f32; 2],
    camera_pos: [f32; 3],
    time: f32,
    camera_target: [f32; 3],
    fov: f32,
    render_flags: [f32; 4],
    scene_scale: f32,
    _pad: [f32; 7],
}

pub struct SdfPipeline {
    render_pipeline: wgpu::RenderPipeline,
    uniform_buffer: wgpu::Buffer,
    bind_group_layout: wgpu::BindGroupLayout,
    bind_group: wgpu::BindGroup,
    format: wgpu::TextureFormat,
    generation: u64,
}

impl SdfPipeline {
    fn build_pipeline(
        device: &wgpu::Device,
        format: wgpu::TextureFormat,
        layout: &wgpu::BindGroupLayout,
        wgsl: &str,
    ) -> wgpu::RenderPipeline {
        let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("sdf_shader"),
            source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(wgsl)),
        });

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

        device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("sdf_rp"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader_module,
                entry_point: Some("vs_main"),
                buffers: &[],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            fragment: Some(wgpu::FragmentState {
                module: &shader_module,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            multiview: None,
            cache: None,
        })
    }
}

impl shader::Pipeline for SdfPipeline {
    fn new(
        device: &wgpu::Device,
        _queue: &wgpu::Queue,
        format: wgpu::TextureFormat,
    ) -> Self {
        let default = default_graph();
        let wgsl = generate_wgsl_from_trig(&default);

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("sdf_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 uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("sdf_uniforms"),
            size: std::mem::size_of::<Uniforms>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("sdf_bg"),
            layout: &bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: uniform_buffer.as_entire_binding(),
            }],
        });

        let render_pipeline = Self::build_pipeline(device, format, &bind_group_layout, &wgsl);

        SdfPipeline {
            render_pipeline,
            uniform_buffer,
            bind_group_layout,
            bind_group,
            format,
            generation: 0,
        }
    }
}

impl shader::Primitive for ViewportPrimitive {
    type Pipeline = SdfPipeline;

    fn prepare(
        &self,
        pipeline: &mut SdfPipeline,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        bounds: &Rectangle,
        viewport: &shader::Viewport,
    ) {
        if self.generation != pipeline.generation {
            pipeline.render_pipeline = SdfPipeline::build_pipeline(
                device,
                pipeline.format,
                &pipeline.bind_group_layout,
                &self.wgsl,
            );
            pipeline.generation = self.generation;
        }

        let scale = viewport.scale_factor() as f32;
        let uniforms = Uniforms {
            resolution: [bounds.width * scale, bounds.height * scale],
            viewport_offset: [bounds.x * scale, bounds.y * scale],
            camera_pos: self.camera_pos,
            time: self.time,
            camera_target: self.camera_target,
            fov: self.fov,
            render_flags: self.render_flags,
            scene_scale: self.scene_scale,
            _pad: [0.0; 7],
        };
        queue.write_buffer(&pipeline.uniform_buffer, 0, bytemuck::bytes_of(&uniforms));
    }

    fn draw(
        &self,
        pipeline: &SdfPipeline,
        render_pass: &mut wgpu::RenderPass<'_>,
    ) -> bool {
        render_pass.set_pipeline(&pipeline.render_pipeline);
        render_pass.set_bind_group(0, &pipeline.bind_group, &[]);
        render_pass.draw(0..3, 0..1);
        true
    }
}
