use bevy_app::{App, Plugin};
use bevy_asset::{Asset, AssetId, Handle, RenderAssetUsages, load_internal_asset, uuid_handle};
use bevy_camera::visibility::RenderLayers;
use bevy_core_pipeline::core_3d::{
    CORE_3D_DEPTH_FORMAT, Transparent3d, TransparentSortingInfo3d,
};
use bevy_core_pipeline::prepass::{
    DeferredPrepass, DepthPrepass, MotionVectorPrepass, NormalPrepass,
};
use bevy_core_pipeline::tonemapping::{DebandDither, Tonemapping};
use bevy_derive::{Deref, DerefMut};
use bevy_ecs::prelude::*;
use bevy_ecs::query::ROQueryItem;
use bevy_ecs::system::SystemParamItem;
use bevy_ecs::system::lifetimeless::{Read, SRes};
use bevy_mesh::{PrimitiveTopology, VertexBufferLayout};
use bevy_pbr::{
    DistanceFog, ExtractedAtmosphere, MeshPipeline, MeshPipelineKey, MeshPipelineSystems,
    ScreenSpaceAmbientOcclusion, SetMeshViewBindGroup, tonemapping_pipeline_key,
};
use bevy_platform::collections::{HashMap, HashSet};
use bevy_reflect::{Reflect, TypePath};
use bevy_render::extract_component::ExtractComponent;
use bevy_render::prelude::*;
use bevy_render::render_asset::{
    PrepareAssetError, RenderAsset, RenderAssetPlugin, RenderAssets, prepare_assets,
};
use bevy_render::render_phase::{
    AddRenderCommand, DrawFunctions, PhaseItem, PhaseItemExtraIndex, RenderCommand,
    RenderCommandResult, SetItemPipeline, TrackedRenderPass, ViewSortedRenderPhases,
};
use bevy_render::render_resource::{
    BlendState, Buffer, BufferInitDescriptor, BufferUsages, ColorTargetState, ColorWrites,
    CompareFunction, DepthBiasState, DepthStencilState, FragmentState, IndexFormat,
    MultisampleState, PipelineCache, PrimitiveState, RenderPipelineDescriptor,
    SpecializedRenderPipeline, SpecializedRenderPipelines, StencilState, VertexAttribute,
    VertexFormat, VertexState, VertexStepMode,
};
use bevy_render::renderer::RenderDevice;
use bevy_render::sync_world::TemporaryRenderEntity;
use bevy_render::view::ExtractedView;
use bevy_render::{Extract, Render, RenderApp, RenderStartup, RenderSystems};
use bevy_render::camera::ExtractedCamera;
use bevy_shader::Shader;
use bytemuck::cast_slice;
use uuid::Uuid;

use crate::GizmoCamera;

const GIZMO_SHADER_HANDLE: Handle<Shader> = uuid_handle!("e44be110-cb2b-4a8d-9c0c-965424e6a633");

pub(crate) struct TransformGizmoRenderPlugin;

impl Plugin for TransformGizmoRenderPlugin {
    fn build(&self, app: &mut App) {
        load_internal_asset!(app, GIZMO_SHADER_HANDLE, "gizmo.wgsl", Shader::from_wgsl);

        app.register_type::<DrawDataHandles>()
            .init_resource::<DrawDataHandles>()
            .add_plugins(RenderAssetPlugin::<GizmoBuffers>::default());

        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
            return;
        };

        render_app
            .add_render_command::<Transparent3d, DrawGizmo>()
            .init_resource::<SpecializedRenderPipelines<TransformGizmoPipeline>>()
            .add_systems(
                Render,
                queue_transform_gizmos
                    .in_set(RenderSystems::Queue)
                    .after(prepare_assets::<GizmoBuffers>),
            );
    }

    fn finish(&self, app: &mut App) {
        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
            return;
        };

        render_app
            .add_systems(ExtractSchedule, extract_gizmo_data)
            .add_systems(
                RenderStartup,
                init_transform_gizmo_pipeline.after(MeshPipelineSystems),
            );
    }
}

#[derive(Resource, Default, Reflect)]
#[reflect(Resource)]
pub(crate) struct DrawDataHandles {
    pub(crate) handles: HashMap<Uuid, GizmoDrawDataHandle>,
}

#[derive(
    Component,
    Default,
    Clone,
    Debug,
    Deref,
    DerefMut,
    Reflect,
    PartialEq,
    Eq,
    ExtractComponent,
    Hash,
)]
#[reflect(Component)]
pub(crate) struct GizmoDrawDataHandle(pub(crate) Handle<GizmoDrawData>);

impl From<Handle<GizmoDrawData>> for GizmoDrawDataHandle {
    fn from(handle: Handle<GizmoDrawData>) -> Self {
        Self(handle)
    }
}

impl From<GizmoDrawDataHandle> for AssetId<GizmoDrawData> {
    fn from(handle: GizmoDrawDataHandle) -> Self {
        handle.0.id()
    }
}
impl From<&GizmoDrawDataHandle> for AssetId<GizmoDrawData> {
    fn from(handle: &GizmoDrawDataHandle) -> Self {
        handle.0.id()
    }
}

fn extract_gizmo_data(mut commands: Commands, handles: Extract<Res<DrawDataHandles>>) {
    let handle_refs = handles.handles.values().cloned().collect::<HashSet<_>>();

    for handle in handle_refs {
        commands.spawn((handle, TemporaryRenderEntity));
    }
}

#[derive(Asset, Debug, Default, Clone, TypePath)]
pub(crate) struct GizmoDrawData(pub(crate) transform_gizmo::GizmoDrawData);

#[derive(Debug, Clone)]
pub(crate) struct GizmoBuffers {
    position_buffer: Buffer,
    index_buffer: Buffer,
    color_buffer: Buffer,
    index_count: u32,
}

impl RenderAsset for GizmoBuffers {
    type SourceAsset = GizmoDrawData;
    type Param = SRes<RenderDevice>;

    fn asset_usage(_source_asset: &Self::SourceAsset) -> RenderAssetUsages {
        RenderAssetUsages::all()
    }

    fn prepare_asset(
        source_asset: Self::SourceAsset,
        _: AssetId<Self::SourceAsset>,
        render_device: &mut SystemParamItem<Self::Param>,
        _: Option<&Self>,
    ) -> std::result::Result<Self, PrepareAssetError<Self::SourceAsset>> {
        let position_buffer_data = cast_slice(&source_asset.0.vertices);
        let position_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
            usage: BufferUsages::VERTEX,
            label: Some("TransformGizmo Position Buffer"),
            contents: position_buffer_data,
        });

        let index_buffer_data = cast_slice(&source_asset.0.indices);
        let index_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
            usage: BufferUsages::INDEX,
            label: Some("TransformGizmo Index Buffer"),
            contents: index_buffer_data,
        });

        let color_buffer_data = cast_slice(&source_asset.0.colors);
        let color_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
            usage: BufferUsages::VERTEX,
            label: Some("TransformGizmo Color Buffer"),
            contents: color_buffer_data,
        });

        Ok(Self {
            index_buffer,
            position_buffer,
            color_buffer,
            index_count: source_asset.0.indices.len() as u32,
        })
    }
}

struct DrawTransformGizmo;

impl<P: PhaseItem> RenderCommand<P> for DrawTransformGizmo {
    type ViewQuery = ();
    type ItemQuery = Read<GizmoDrawDataHandle>;
    type Param = SRes<RenderAssets<GizmoBuffers>>;

    #[inline]
    fn render<'w>(
        _item: &P,
        _view: ROQueryItem<'w, '_, Self::ViewQuery>,
        handle: Option<ROQueryItem<'w, '_, Self::ItemQuery>>,
        gizmos: SystemParamItem<'w, '_, Self::Param>,
        pass: &mut TrackedRenderPass<'w>,
    ) -> RenderCommandResult {
        let Some(handle) = handle else {
            return RenderCommandResult::Failure("No GizmoDrawDataHandle component found");
        };

        let Some(gizmo) = gizmos.into_inner().get(handle) else {
            return RenderCommandResult::Failure("No GizmoDrawDataHandle inner found");
        };

        if gizmo.index_buffer.size() > 0 {
            pass.set_index_buffer(gizmo.index_buffer.slice(..), IndexFormat::Uint32);
            pass.set_vertex_buffer(0, gizmo.position_buffer.slice(..));
            pass.set_vertex_buffer(1, gizmo.color_buffer.slice(..));

            pass.draw_indexed(0..gizmo.index_count, 0, 0..1);
        }

        RenderCommandResult::Success
    }
}

#[derive(Clone, Resource)]
struct TransformGizmoPipeline {
    mesh_pipeline: MeshPipeline,
}

fn init_transform_gizmo_pipeline(mut commands: Commands, mesh_pipeline: Res<MeshPipeline>) {
    commands.insert_resource(TransformGizmoPipeline {
        mesh_pipeline: mesh_pipeline.clone(),
    });
}

#[derive(PartialEq, Eq, Hash, Clone)]
struct TransformGizmoPipelineKey {
    view_key: MeshPipelineKey,
    perspective: bool,
}

impl SpecializedRenderPipeline for TransformGizmoPipeline {
    type Key = TransformGizmoPipelineKey;

    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
        let mut shader_defs = vec![
            // TODO: When is this flag actually used?
            // #[cfg(feature = "webgl")]
            // "SIXTEEN_BYTE_ALIGNMENT".into(),
        ];

        if key.perspective {
            shader_defs.push("PERSPECTIVE".into());
        }

        let format = key.view_key.target_format();

        let view_layout = self
            .mesh_pipeline
            .get_view_layout(key.view_key.into())
            .clone();
        let layout = vec![view_layout.main_layout.clone()];

        RenderPipelineDescriptor {
            label: Some("TransformGizmo Pipeline".into()),
            zero_initialize_workgroup_memory: true, // ?
            vertex: VertexState {
                shader: GIZMO_SHADER_HANDLE,
                entry_point: Some("vertex".into()),
                shader_defs: shader_defs.clone(),
                buffers: vec![
                    VertexBufferLayout {
                        array_stride: VertexFormat::Float32x2.size(),
                        step_mode: VertexStepMode::Vertex,
                        attributes: vec![VertexAttribute {
                            format: VertexFormat::Float32x2,
                            offset: 0,
                            shader_location: 0,
                        }],
                    },
                    VertexBufferLayout {
                        array_stride: VertexFormat::Float32x4.size(),
                        step_mode: VertexStepMode::Vertex,
                        attributes: vec![VertexAttribute {
                            format: VertexFormat::Float32x4,
                            offset: 0,
                            shader_location: 1,
                        }],
                    },
                ],
            },
            fragment: Some(FragmentState {
                shader: GIZMO_SHADER_HANDLE,
                shader_defs,
                entry_point: Some("fragment".into()),
                targets: vec![Some(ColorTargetState {
                    format,
                    blend: Some(BlendState::PREMULTIPLIED_ALPHA_BLENDING),
                    write_mask: ColorWrites::ALL,
                })],
            }),
            layout,
            primitive: PrimitiveState {
                topology: PrimitiveTopology::TriangleList,
                cull_mode: None,
                ..PrimitiveState::default()
            },
            depth_stencil: Some(DepthStencilState {
                format: CORE_3D_DEPTH_FORMAT,
                depth_write_enabled: Some(true),
                depth_compare: Some(CompareFunction::Always),
                stencil: StencilState::default(),
                bias: DepthBiasState::default(),
            }),
            multisample: MultisampleState {
                count: key.view_key.msaa_samples(),
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            immediate_size: 0,
        }
    }
}

type DrawGizmo = (SetItemPipeline, SetMeshViewBindGroup<0>, DrawTransformGizmo);

#[allow(clippy::too_many_arguments, clippy::type_complexity)]
fn queue_transform_gizmos(
    draw_functions: Res<DrawFunctions<Transparent3d>>,
    pipeline: Res<TransformGizmoPipeline>,
    mut pipelines: ResMut<SpecializedRenderPipelines<TransformGizmoPipeline>>,
    pipeline_cache: Res<PipelineCache>,
    msaa_q: Query<Option<&Msaa>, With<GizmoCamera>>,
    transform_gizmos: Query<(Entity, &GizmoDrawDataHandle)>,
    transform_gizmo_assets: Res<RenderAssets<GizmoBuffers>>,
    mut views: Query<(
        Entity,
        &ExtractedView,
        Option<&ExtractedCamera>,
        Option<&Msaa>,
        (Option<&Tonemapping>, Option<&DebandDither>),
        Option<&RenderLayers>,
        (
            Has<NormalPrepass>,
            Has<DepthPrepass>,
            Has<MotionVectorPrepass>,
            Has<DeferredPrepass>,
            Has<ExtractedAtmosphere>,
            Has<ScreenSpaceAmbientOcclusion>,
            Has<DistanceFog>,
        ),
    )>,
    mut transparent_render_phases: ResMut<ViewSortedRenderPhases<Transparent3d>>,
) {
    let draw_function = draw_functions.read().get_id::<DrawGizmo>().unwrap();
    let camera_msaa = msaa_q.single().ok().flatten();
    for (
        view_entity,
        view,
        camera,
        entity_msaa,
        (tonemapping, dither),
        _render_layers,
        (
            normal_prepass,
            depth_prepass,
            motion_vector_prepass,
            deferred_prepass,
            atmosphere,
            ssao,
            distance_fog,
        ),
    ) in &mut views
    {
        let Some(transparent_phase) = transparent_render_phases.get_mut(&view.retained_view_entity)
        else {
            continue;
        };

        // entity_msaa > camera_msaa > default
        let msaa_sample_count = entity_msaa.map_or(
            camera_msaa.unwrap_or(&Msaa::default()).samples(),
            Msaa::samples,
        );

        let mut view_key = MeshPipelineKey::from_msaa_samples(msaa_sample_count)
            | MeshPipelineKey::from_target_format(view.target_format);

        if normal_prepass {
            view_key |= MeshPipelineKey::NORMAL_PREPASS;
        }

        if depth_prepass {
            view_key |= MeshPipelineKey::DEPTH_PREPASS;
        }

        if motion_vector_prepass {
            view_key |= MeshPipelineKey::MOTION_VECTOR_PREPASS;
        }

        if deferred_prepass {
            view_key |= MeshPipelineKey::DEFERRED_PREPASS;
        }

        if atmosphere {
            view_key |= MeshPipelineKey::ATMOSPHERE;
        }

        if ssao {
            view_key |= MeshPipelineKey::SCREEN_SPACE_AMBIENT_OCCLUSION;
        }

        if distance_fog {
            view_key |= MeshPipelineKey::DISTANCE_FOG;
        }

        if !camera.is_some_and(|camera| camera.hdr) {
            if let Some(tonemapping) = tonemapping {
                view_key |= MeshPipelineKey::TONEMAP_IN_SHADER;
                view_key |= tonemapping_pipeline_key(*tonemapping);
            }
            if let Some(DebandDither::Enabled) = dither {
                view_key |= MeshPipelineKey::DEBAND_DITHER;
            }
        }

        for (entity, handle) in &transform_gizmos {
            let Some(_) = transform_gizmo_assets.get(handle.0.id()) else {
                continue;
            };

            let pipeline = pipelines.specialize(
                &pipeline_cache,
                &pipeline,
                TransformGizmoPipelineKey {
                    view_key,
                    perspective: true,
                },
            );

            transparent_phase.add_transient(Transparent3d {
                sorting_info: TransparentSortingInfo3d::AlwaysOnTop,
                entity: (entity, view_entity.into()),
                draw_function,
                pipeline,
                distance: 0.,
                batch_range: 0..1,
                extra_index: PhaseItemExtraIndex::None,
                indexed: false,
            });
        }
    }
}
