use image::{Rgba, Rgba32FImage};
use intuicio_core::{
    IntuicioStruct, context::Context, function::FunctionQuery, host::HostProducer,
    registry::Registry,
};
use intuicio_derive::*;
use intuicio_frontend_simpleton::{
    Array, Boolean, Integer, Map, Real, Reference,
    library::{closure::Closure, jobs::Jobs},
};
use std::{
    collections::HashMap,
    sync::Arc,
    thread::{available_parallelism, spawn},
};

use crate::library::{color::Color, image::Image};

#[derive(Default, Clone)]
enum Primitive {
    #[default]
    Null,
    Integer(Integer),
    Real(Real),
    Array(Vec<Primitive>),
    Map(HashMap<String, Primitive>),
}

impl Primitive {
    fn from_value(value: &Reference) -> Self {
        if let Some(value) = value.read::<Integer>() {
            Self::Integer(*value)
        } else if let Some(value) = value.read::<Real>() {
            Self::Real(*value)
        } else if let Some(value) = value.read::<Array>() {
            Self::Array(value.iter().map(Self::from_value).collect())
        } else if let Some(value) = value.read::<Map>() {
            Self::Map(
                value
                    .iter()
                    .map(|(key, value)| (key.to_owned(), Self::from_value(value)))
                    .collect(),
            )
        } else {
            Self::Null
        }
    }

    fn to_value(&self, registry: &Registry) -> Reference {
        match self {
            Self::Null => Reference::null(),
            Self::Integer(value) => Reference::new_integer(*value, registry),
            Self::Real(value) => Reference::new_real(*value, registry),
            Self::Array(value) => Reference::new_array(
                value.iter().map(|value| value.to_value(registry)).collect(),
                registry,
            ),
            Self::Map(value) => Reference::new_map(
                value
                    .iter()
                    .map(|(key, value)| (key.to_owned(), value.to_value(registry)))
                    .collect(),
                registry,
            ),
        }
    }
}

#[derive(IntuicioStruct, Default, Clone)]
#[intuicio(
    name = "Fragment",
    module_name = "image_pipeline",
    override_send = true
)]
pub struct Fragment {
    pub index: Reference,
    pub col: Reference,
    pub row: Reference,
    pub u: Reference,
    pub v: Reference,
    pub width: Reference,
    pub height: Reference,
    pub start: Reference,
    pub end: Reference,
}

#[derive(IntuicioStruct, Default, Clone)]
#[intuicio(
    name = "Pipeline",
    module_name = "image_pipeline",
    override_send = true
)]
pub struct Pipeline {
    pub width: Reference,
    pub height: Reference,
    pub samplers: Reference,
}

#[intuicio_methods(module_name = "image_pipeline")]
impl Pipeline {
    #[intuicio_method(use_context, use_registry)]
    pub fn process_single_thread(
        context: &mut Context,
        registry: &Registry,
        pipeline: Reference,
        closure: Reference,
    ) -> Reference {
        let pipeline = pipeline.read::<Pipeline>().unwrap();
        let width = *pipeline.width.read::<Integer>().unwrap() as u32;
        let height = *pipeline.height.read::<Integer>().unwrap() as u32;
        {
            let samplers = pipeline.samplers.read::<Map>().unwrap();
            for sampler in samplers.values() {
                if !sampler.type_of().unwrap().is::<Sampler>() {
                    panic!("All items of `samplers` must be of Sampler type!");
                }
            }
        }
        let closure = closure.read::<Closure>().unwrap();
        let pixels_count = width * height;
        let output_width = width.saturating_sub(1) as Real;
        let output_height = height.saturating_sub(1) as Real;
        let result = (0..pixels_count)
            .map(|index| {
                let x = (index % width) as Integer;
                let y = (index / width) as Integer;
                let u = x as Real / output_width;
                let v = y as Real / output_height;
                let fragment = Reference::new(
                    Fragment {
                        index: Reference::new_integer(index as Integer, registry),
                        col: Reference::new_integer(x, registry),
                        row: Reference::new_integer(y, registry),
                        u: Reference::new_real(u, registry),
                        v: Reference::new_real(v, registry),
                        width: Reference::new_integer(width as Integer, registry),
                        height: Reference::new_integer(height as Integer, registry),
                        start: Reference::new_integer(0, registry),
                        end: Reference::new_integer(pixels_count as Integer, registry),
                    },
                    registry,
                );
                let args = [fragment, pipeline.samplers.clone()];
                closure
                    .invoke(context, registry, &args)
                    .read::<Color>()
                    .map(|color| color.to_pixel())
                    .unwrap_or_else(|| Rgba([0.0, 0.0, 0.0, 0.0]))
            })
            .flat_map(|pixel| pixel.0)
            .collect::<Vec<_>>();
        Rgba32FImage::from_vec(width, height, result)
            .map(|buffer| Reference::new(Image { buffer }, registry))
            .unwrap_or_default()
    }

    #[intuicio_method(use_context, use_registry)]
    pub fn process_multi_thread(
        context: &mut Context,
        registry: &Registry,
        pipeline: Reference,
        closure: Reference,
    ) -> Reference {
        let threads_count = available_parallelism()
            .ok()
            .map(|count| count.get() as u32)
            .unwrap_or_default();
        if threads_count <= 1 {
            return Self::process_single_thread(context, registry, pipeline, closure);
        }
        let closure = closure.read::<Closure>().unwrap();
        let captures = closure
            .captured
            .iter()
            .map(Primitive::from_value)
            .collect::<Vec<_>>();
        let host_producer = match context.custom::<HostProducer>(Jobs::HOST_PRODUCER_CUSTOM) {
            Some(host_producer) => host_producer.clone(),
            None => return Reference::null(),
        };
        let pipeline = pipeline.read::<Pipeline>().unwrap();
        let width = *pipeline.width.read::<Integer>().unwrap() as u32;
        let height = *pipeline.height.read::<Integer>().unwrap() as u32;
        let samplers = pipeline.samplers.read::<Map>().unwrap();
        let signature = closure.function.handle().unwrap().signature();
        let function_name = signature.name.to_owned();
        let function_module_name = signature.module_name.to_owned();
        let pixels_count = width * height;
        let pixels_per_thread = if pixels_count.is_multiple_of(threads_count) {
            pixels_count / threads_count
        } else {
            1 + pixels_count / threads_count
        };
        let output_width = width.saturating_sub(1) as Real;
        let output_height = height.saturating_sub(1) as Real;
        let mut offset = 0;
        let function_name_ = &function_name;
        let function_module_name_ = &function_module_name;
        let captures_ = captures.as_slice();
        let threads = (0..threads_count)
            .map(move |_| {
                let start = offset;
                offset += pixels_per_thread;
                let end = offset.min(pixels_count);
                let samplers = samplers
                    .iter()
                    .map(|(key, value)| (key.to_owned(), value.read::<Sampler>().unwrap().clone()))
                    .collect::<HashMap<_, _>>();
                let host_producer = host_producer.clone();
                let function_name = function_name_.to_owned();
                let function_module_name = function_module_name_.to_owned();
                let captures = captures_.to_owned();
                spawn(move || {
                    let mut host = host_producer.produce();
                    host.context()
                        .set_custom(Jobs::HOST_PRODUCER_CUSTOM, host_producer);
                    let (context, registry) = host.context_and_registry();
                    let samplers = samplers
                        .into_iter()
                        .map(|(key, value)| (key, Reference::new(value, registry)))
                        .collect::<HashMap<_, _>>();
                    let samplers = Reference::new_map(samplers, registry);
                    let captures = captures
                        .iter()
                        .map(|primitive| primitive.to_value(registry))
                        .collect::<Vec<_>>();
                    let size = (end - start) as usize;
                    if let Some(function) = registry.find_function(FunctionQuery {
                        name: Some(function_name.into()),
                        module_name: function_module_name.map(|name| name.into()),
                        ..Default::default()
                    }) {
                        (start..end)
                            .map(|index| {
                                let x = (index % width) as Integer;
                                let y = (index / width) as Integer;
                                let u = x as Real / output_width;
                                let v = y as Real / output_height;
                                let fragment = Reference::new(
                                    Fragment {
                                        index: Reference::new_integer(index as Integer, registry),
                                        col: Reference::new_integer(x, registry),
                                        row: Reference::new_integer(y, registry),
                                        u: Reference::new_real(u, registry),
                                        v: Reference::new_real(v, registry),
                                        width: Reference::new_integer(width as Integer, registry),
                                        height: Reference::new_integer(height as Integer, registry),
                                        start: Reference::new_integer(start as Integer, registry),
                                        end: Reference::new_integer(end as Integer, registry),
                                    },
                                    registry,
                                );
                                context.stack().push(samplers.clone());
                                context.stack().push(fragment);
                                for value in captures.iter().rev() {
                                    context.stack().push(value.clone());
                                }
                                function.invoke(context, registry);
                                context
                                    .stack()
                                    .pop::<Reference>()
                                    .unwrap()
                                    .read::<Color>()
                                    .map(|color| color.to_pixel())
                                    .unwrap_or_else(|| Rgba([0.0, 0.0, 0.0, 0.0]))
                            })
                            .collect()
                    } else {
                        vec![Rgba([0.0, 0.0, 0.0, 0.0]); size]
                    }
                })
            })
            .collect::<Vec<_>>();
        let result = threads
            .into_iter()
            .flat_map(|handle| handle.join().unwrap())
            .flat_map(|pixel| pixel.0)
            .collect::<Vec<_>>();
        Rgba32FImage::from_vec(width, height, result)
            .map(|buffer| Reference::new(Image { buffer }, registry))
            .unwrap_or_default()
    }
}

#[derive(IntuicioStruct, Default, Clone)]
#[intuicio(name = "Sampler", module_name = "image_sampler")]
pub struct Sampler {
    #[intuicio(ignore)]
    image: Arc<Image>,
}

#[intuicio_methods(module_name = "image_sampler")]
impl Sampler {
    #[allow(clippy::new_ret_no_self)]
    #[intuicio_method(use_registry)]
    pub fn new(registry: &Registry, image: Reference) -> Reference {
        let image = image.read::<Image>().unwrap();
        Reference::new(
            Sampler {
                image: Arc::new(image.clone()),
            },
            registry,
        )
    }

    #[intuicio_method(use_registry)]
    pub fn clone(registry: &Registry, sampler: Reference) -> Reference {
        Reference::new(sampler.read::<Sampler>().unwrap().clone(), registry)
    }

    #[intuicio_method(use_registry)]
    pub fn sample(
        registry: &Registry,
        sampler: Reference,
        u: Reference,
        v: Reference,
        interpolate: Reference,
        wrap: Reference,
    ) -> Reference {
        let sampler = sampler.read::<Sampler>().unwrap();
        let u = *u.read::<Real>().unwrap() as f32;
        let v = *v.read::<Real>().unwrap() as f32;
        let wrap = *wrap.read::<Boolean>().unwrap();
        let interpolate = *interpolate.read::<Boolean>().unwrap();
        let result = sampler.image.sample_inner(u, v, wrap, interpolate);
        Reference::new(Color::from_pixel(&result, registry), registry)
    }

    #[intuicio_method(use_registry)]
    pub fn fetch(
        registry: &Registry,
        sampler: Reference,
        col: Reference,
        row: Reference,
    ) -> Reference {
        let sampler = sampler.read::<Sampler>().unwrap();
        let col = *col.read::<Integer>().unwrap() as u32;
        let row = *row.read::<Integer>().unwrap() as u32;
        let result = sampler.image.get_pixel_inner(col, row);
        Reference::new(Color::from_pixel(&result, registry), registry)
    }
}

pub fn install(registry: &mut Registry) {
    registry.add_type(Fragment::define_struct(registry));
    registry.add_type(Sampler::define_struct(registry));
    registry.add_type(Pipeline::define_struct(registry));
    registry.add_function(Pipeline::process_single_thread__define_function(registry));
    registry.add_function(Pipeline::process_multi_thread__define_function(registry));
    registry.add_function(Sampler::new__define_function(registry));
    registry.add_function(Sampler::clone__define_function(registry));
    registry.add_function(Sampler::sample__define_function(registry));
    registry.add_function(Sampler::fetch__define_function(registry));
}
