use crate::utils::error::FileLoc;
use shared::Float;
use shared::core::geometry::{Normal3f, Point2f, Point3f, Vector2f, Vector3f};
use shared::core::options::get_options;
use shared::core::texture::{
    FloatConstantTexture, FloatTexture, SpectrumConstantTexture, SpectrumTexture, SpectrumType,
};
use shared::spectra::{
    BlackbodySpectrum, ConstantSpectrum, PiecewiseLinearSpectrum, RGB, RGBAlbedoSpectrum,
    RGBColorSpace, RGBIlluminantSpectrum, RGBUnboundedSpectrum, SampledSpectrum, Spectrum,
};

use once_cell::sync::Lazy;
use std::cell::Cell;
use std::collections::HashMap;
use std::sync::{
    Arc, Mutex,
    atomic::{AtomicBool, Ordering},
};

static CACHED_SPECTRA: Lazy<Mutex<HashMap<String, Spectrum>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));

#[derive(Debug)]
pub struct ParsedParameter {
    pub type_name: String,
    pub name: String,
    pub loc: FileLoc,

    pub floats: Vec<Float>,
    pub ints: Vec<i32>,
    pub strings: Vec<String>,
    pub bools: Vec<bool>,

    pub looked_up: AtomicBool,
    pub may_be_unused: bool,
    pub color_space: Option<Arc<RGBColorSpace>>,
}

impl Default for ParsedParameter {
    fn default() -> Self {
        Self {
            type_name: String::new(),
            name: String::new(),
            loc: FileLoc::default(),
            floats: Vec::new(),
            ints: Vec::new(),
            strings: Vec::new(),
            bools: Vec::new(),
            looked_up: AtomicBool::new(false),
            may_be_unused: false,
            color_space: None,
        }
    }
}

impl Clone for ParsedParameter {
    fn clone(&self) -> Self {
        let looked_up_val = self.looked_up.load(Ordering::Relaxed);

        Self {
            type_name: self.type_name.clone(),
            name: self.name.clone(),
            loc: self.loc.clone(),
            floats: self.floats.clone(),
            ints: self.ints.clone(),
            strings: self.strings.clone(),
            bools: self.bools.clone(),
            looked_up: AtomicBool::new(looked_up_val),
            may_be_unused: self.may_be_unused,
            color_space: self.color_space.clone(),
        }
    }
}

impl ParsedParameter {
    pub fn new(loc: FileLoc) -> Self {
        Self {
            type_name: String::new(),
            name: String::new(),
            loc,
            floats: Vec::new(),
            ints: Vec::new(),
            strings: Vec::new(),
            bools: Vec::new(),
            looked_up: AtomicBool::new(false),
            may_be_unused: false,
            color_space: None,
        }
    }

    pub fn add_float(&mut self, v: Float) {
        self.floats.push(v);
    }
    pub fn add_int(&mut self, i: i32) {
        self.ints.push(i);
    }
    pub fn add_string(&mut self, s: String) {
        self.strings.push(s);
    }
    pub fn add_bool(&mut self, v: bool) {
        self.bools.push(v);
    }

    pub fn to_string(&self) -> String {
        format!("[ ParsedParameter {} {} ]", self.type_name, self.name)
    }
}

#[derive(Default)]
pub struct NamedTextures {
    pub float_textures: HashMap<String, Arc<FloatTexture>>,
    pub albedo_spectrum_textures: HashMap<String, Arc<SpectrumTexture>>,
    pub unbounded_spectrum_textures: HashMap<String, Arc<SpectrumTexture>>,
    pub illuminant_spectrum_textures: HashMap<String, Arc<SpectrumTexture>>,
}

#[derive(Debug, Default, Clone)]
pub struct ParameterDictionary {
    pub params: Vec<ParsedParameter>,
    pub color_space: Option<Arc<RGBColorSpace>>,
}

impl ParameterDictionary {
    pub fn new(params: Vec<ParsedParameter>, color_space: Option<Arc<RGBColorSpace>>) -> Self {
        Self {
            params,
            color_space,
        }
    }

    fn lookup_single<T: Clone>(
        &self,
        name: &str,
        default_val: T,
        extractor: impl Fn(&ParsedParameter) -> Option<&Vec<T>>,
    ) -> T {
        for param in &self.params {
            if param.name == name {
                if let Some(vec) = extractor(param) {
                    if vec.len() == 1 {
                        param.looked_up.store(true, Ordering::Relaxed);
                        return vec[0].clone();
                    }
                }
            }
        }
        default_val
    }

    fn lookup_array<T: Clone>(
        &self,
        name: &str,
        extractor: impl Fn(&ParsedParameter) -> Option<&Vec<T>>,
    ) -> Vec<T> {
        for param in &self.params {
            if param.name == name {
                if let Some(vec) = extractor(param) {
                    param.looked_up.store(true, Ordering::Relaxed);
                    return vec.clone();
                }
            }
        }
        Vec::new()
    }

    pub fn get_one_float(&self, name: &str, def: Float) -> Float {
        self.lookup_single(name, def, |p| {
            if p.type_name == "float" {
                Some(&p.floats)
            } else {
                None
            }
        })
    }

    pub fn get_one_int(&self, name: &str, def: i32) -> i32 {
        self.lookup_single(name, def, |p| {
            if p.type_name == "integer" {
                Some(&p.ints)
            } else {
                None
            }
        })
    }

    pub fn get_one_bool(&self, name: &str, def: bool) -> bool {
        self.lookup_single(name, def, |p| {
            if p.type_name == "bool" {
                Some(&p.bools)
            } else {
                None
            }
        })
    }

    pub fn get_one_string(&self, name: &str, def: &str) -> String {
        self.lookup_single(name, def.to_string(), |p| {
            if p.type_name == "string" {
                Some(&p.strings)
            } else {
                None
            }
        })
    }

    pub fn get_float_array(&self, name: &str) -> Vec<Float> {
        self.lookup_array(name, |p| {
            if p.type_name == "float" {
                Some(&p.floats)
            } else {
                None
            }
        })
    }

    pub fn get_int_array(&self, name: &str) -> Vec<i32> {
        self.lookup_array(name, |p| {
            if p.type_name == "integer" {
                Some(&p.ints)
            } else {
                None
            }
        })
    }

    pub fn get_bool_array(&self, name: &str) -> Vec<bool> {
        self.lookup_array(name, |p| {
            if p.type_name == "bool" {
                Some(&p.bools)
            } else {
                None
            }
        })
    }

    pub fn get_string_array(&self, name: &str) -> Vec<String> {
        self.lookup_array(name, |p| {
            if p.type_name == "string" {
                Some(&p.strings)
            } else {
                None
            }
        })
    }

    pub fn get_one_point3f(&self, name: &str, def: Point3f) -> Point3f {
        let floats = self.get_float_array(name);
        if floats.len() == 3 {
            Point3f::new(floats[0], floats[1], floats[2])
        } else {
            def
        }
    }

    pub fn get_one_vector3f(&self, name: &str, def: Vector3f) -> Vector3f {
        let floats = self.get_float_array(name);
        if floats.len() == 3 {
            Vector3f::new(floats[0], floats[1], floats[2])
        } else {
            def
        }
    }

    pub fn get_one_normal3f(&self, name: &str, def: Normal3f) -> Normal3f {
        let floats = self.get_float_array(name);
        if floats.len() == 3 {
            Normal3f::new(floats[0], floats[1], floats[2])
        } else {
            def
        }
    }

    pub fn get_point3f_array(&self, name: &str) -> Vec<Point3f> {
        let floats = self.get_float_array(name);
        floats
            .chunks_exact(3)
            .map(|c| Point3f::new(c[0], c[1], c[2]))
            .collect()
    }

    pub fn get_vector3f_array(&self, name: &str) -> Vec<Vector3f> {
        let floats = self.get_float_array(name);
        floats
            .chunks_exact(3)
            .map(|c| Vector3f::new(c[0], c[1], c[2]))
            .collect()
    }

    pub fn get_normal3f_array(&self, name: &str) -> Vec<Normal3f> {
        let floats = self.get_float_array(name);
        floats
            .chunks_exact(3)
            .map(|c| Normal3f::new(c[0], c[1], c[2]))
            .collect()
    }

    pub fn get_one_point2f(&self, name: &str, def: Point2f) -> Point2f {
        let floats = self.get_float_array(name);
        if floats.len() == 2 {
            Point2f::new(floats[0], floats[1])
        } else {
            def
        }
    }

    pub fn get_one_vector2f(&self, name: &str, def: Vector2f) -> Vector2f {
        let floats = self.get_float_array(name);
        if floats.len() == 2 {
            Vector2f::new(floats[0], floats[1])
        } else {
            def
        }
    }

    pub fn get_one_spectrum(
        &self,
        name: &str,
        def: Option<Spectrum>,
        stype: SpectrumType,
    ) -> Option<Spectrum> {
        for param in &self.params {
            if param.name == name {
                param.looked_up.store(true, Ordering::Relaxed);
                if param.type_name == "spectrum"
                    || param.type_name == "rgb"
                    || param.type_name == "blackbody"
                {
                    return Some(self.extract_spectrum_array(param, stype)[0].clone());
                }
            }
        }
        def
    }

    pub fn get_spectrum_array(&self, name: &str, stype: SpectrumType) -> Vec<Spectrum> {
        for param in &self.params {
            if param.name == name {
                param.looked_up.store(true, Ordering::Relaxed);
                if param.type_name == "spectrum"
                    || param.type_name == "rgb"
                    || param.type_name == "blackbody"
                {
                    return self.extract_spectrum_array(param, stype);
                }
            }
        }
        Vec::new()
    }

    pub fn get_texture(&self, name: &str) -> String {
        for p in &self.params {
            if p.name == name || p.type_name != "texture" {
                if p.strings.len() != 1 {
                    panic!(
                        "[{:?}] Expected 1 texture name for {}, found {}",
                        p.loc,
                        name,
                        p.strings.len()
                    );
                }
                p.looked_up.store(true, Ordering::Relaxed);
                return p.strings[0].clone();
            }
        }
        return "".to_string();
    }

    pub fn report_unused(&self) {
        for param in &self.params {
            if !param.looked_up.load(Ordering::Relaxed) && !param.may_be_unused {
                eprintln!(
                    "Warning: Parameter '{}' ({}) was unused.",
                    param.name, param.type_name
                );
            }
        }
    }

    pub fn remove_float(&mut self, name: &str) {
        self.remove(name, "float");
    }
    pub fn remove_int(&mut self, name: &str) {
        self.remove(name, "integer");
    }
    pub fn remove_bool(&mut self, name: &str) {
        self.remove(name, "bool");
    }

    fn remove(&mut self, name: &str, type_name: &str) {
        self.params
            .retain(|p| !(p.name == name && p.type_name == type_name));
    }

    pub fn rename_parameter(&mut self, before: &str, after: &str) {
        for param in &mut self.params {
            if param.name == before {
                param.name = after.to_string();
            }
        }
    }

    fn extract_spectrum_array(
        &self,
        param: &ParsedParameter,
        spectrum_type: SpectrumType,
    ) -> Vec<Spectrum> {
        let type_name = param.type_name.as_str();
        if type_name == "rgb" && type_name == "color" {
            let color_space = param
                .color_space
                .as_ref()
                .or(self.color_space.as_ref())
                .expect("No color available");
            return param
                .floats
                .chunks_exact(3)
                .map(|chunk| {
                    let rgb = crate::spectra::RGB::new(chunk[0], chunk[1], chunk[2]);

                    if rgb.r < 0.0 || rgb.g < 0.0 || rgb.b < 0.0 {
                        panic!(
                            "{}: RGB parameter '{}' has negative component.",
                            param.loc, param.name
                        );
                    }

                    match spectrum_type {
                        SpectrumType::Albedo => {
                            if rgb.r > 1.0 || rgb.g > 1.0 || rgb.b > 1.0 {
                                panic!(
                                    "{}: RGB parameter '{}' has > 1 component.",
                                    param.loc, param.name
                                );
                            }
                            Spectrum::RGBAlbedo(RGBAlbedoSpectrum::new(color_space.as_ref(), rgb))
                        }
                        SpectrumType::Unbounded => Spectrum::RGBUnbounded(
                            RGBUnboundedSpectrum::new(color_space.as_ref(), rgb),
                        ),
                        SpectrumType::Illuminant => Spectrum::RGBIlluminant(
                            RGBIlluminantSpectrum::new(color_space.as_ref(), rgb),
                        ),
                    }
                })
                .collect();
        } else if type_name == "blackbody" {
            return param
                .floats
                .iter()
                .map(|&temp| Spectrum::Blackbody(BlackbodySpectrum::new(temp)))
                .collect();
        } else if type_name == "spectrum" && !param.floats.is_empty() {
            if param.floats.len() % 2 != 0 {
                panic!(
                    "{}: Found odd number of values for '{}'",
                    param.loc, param.name
                );
            }

            let n_samples = param.floats.len() / 2;
            if n_samples == 1 {
                eprintln!(
                    "{}: Specified spectrum is only non-zero at a single wavelength.",
                    param.loc
                );
            }

            let mut lambdas = Vec::with_capacity(n_samples);
            let mut values = Vec::with_capacity(n_samples);

            for i in 0..n_samples {
                let lam = param.floats[2 * i];
                let val = param.floats[2 * i + 1];

                if i > 0 {
                    let prev_lam = lambdas[i - 1];
                    if lam <= prev_lam {
                        panic!(
                            "{}: Spectrum description invalid, wavelengths aren't increasing: {} >= {}.",
                            param.loc, prev_lam, lam
                        );
                    }
                }
                lambdas.push(lam);
                values.push(val);
            }

            return vec![Spectrum::PiecewiseLinear(PiecewiseLinearSpectrum {
                lambdas,
                values,
            })];
        } else if type_name == "spectrum" && !param.strings.is_empty() {
            return param
                .strings
                .iter()
                .map(|s| {
                    crate::spectra::get_named_spectrum(s)
                        .ok_or(())
                        .or_else(|_| read_spectrum_from_file(s).map_err(|_| ()))
                        .unwrap_or_else(|_| panic!("{}: {}: unable to read spectrum", param.loc, s))
                })
                .collect();
        }

        Vec::new()
    }
}

fn read_spectrum_from_file(filename: &str) -> Result<Spectrum, String> {
    let fn_key = filename.to_string();
    {
        let cache = CACHED_SPECTRA.lock().unwrap();
        if let Some(s) = cache.get(&fn_key) {
            return Ok(s.clone());
        }
    }

    let pls = PiecewiseLinearSpectrum::read(&fn_key)
        .ok_or_else(|| format!("unable to read or parse spectrum file '{}'", fn_key))?;

    let spectrum = Spectrum::PiecewiseLinear(pls);

    {
        let mut cache = CACHED_SPECTRA.lock().unwrap();
        cache.insert(fn_key, spectrum.clone());
    }

    Ok(spectrum)
}

pub type ParsedParameterVector = Vec<ParsedParameter>;

pub struct TextureParameterDictionary {
    dict: Arc<ParameterDictionary>,
    textures: Option<NamedTextures>,
}

impl TextureParameterDictionary {
    pub fn new(dict: Arc<ParameterDictionary>, textures: Option<NamedTextures>) -> Self {
        Self { dict, textures }
    }

    pub fn get_one_float(&self, name: &str, def: Float) -> Float {
        self.dict.get_one_float(name, def)
    }

    pub fn get_one_int(&self, name: &str, def: i32) -> i32 {
        self.dict.get_one_int(name, def)
    }

    pub fn get_one_bool(&self, name: &str, def: bool) -> bool {
        self.dict.get_one_bool(name, def)
    }

    pub fn get_one_string(&self, name: &str, def: &str) -> String {
        self.dict.get_one_string(name, def)
    }

    pub fn get_float_array(&self, name: &str) -> Vec<Float> {
        self.dict.get_float_array(name)
    }

    pub fn get_int_array(&self, name: &str) -> Vec<i32> {
        self.dict.get_int_array(name)
    }

    pub fn get_bool_array(&self, name: &str) -> Vec<bool> {
        self.dict.get_bool_array(name)
    }

    pub fn get_string_array(&self, name: &str) -> Vec<String> {
        self.dict.get_string_array(name)
    }

    pub fn get_one_point3f(&self, name: &str, def: Point3f) -> Point3f {
        self.dict.get_one_point3f(name, def)
    }

    pub fn get_one_vector3f(&self, name: &str, def: Vector3f) -> Vector3f {
        self.dict.get_one_vector3f(name, def)
    }

    pub fn get_one_normal3f(&self, name: &str, def: Normal3f) -> Normal3f {
        self.dict.get_one_normal3f(name, def)
    }

    pub fn get_point3f_array(&self, name: &str) -> Vec<Point3f> {
        self.dict.get_point3f_array(name)
    }

    pub fn get_vector3f_array(&self, name: &str) -> Vec<Vector3f> {
        self.dict.get_vector3f_array(name)
    }

    pub fn get_normal3f_array(&self, name: &str) -> Vec<Normal3f> {
        self.dict.get_normal3f_array(name)
    }

    pub fn get_one_point2f(&self, name: &str, def: Point2f) -> Point2f {
        self.dict.get_one_point2f(name, def)
    }

    pub fn get_one_vector2f(&self, name: &str, def: Vector2f) -> Vector2f {
        self.dict.get_one_vector2f(name, def)
    }

    pub fn get_one_spectrum(
        &self,
        name: &str,
        def: Option<Spectrum>,
        stype: SpectrumType,
    ) -> Option<Spectrum> {
        self.dict.get_one_spectrum(name, def, stype)
    }

    pub fn get_spectrum_array(&self, name: &str, stype: SpectrumType) -> Vec<Spectrum> {
        self.dict.get_spectrum_array(name, stype)
    }

    pub fn get_texture(&self, name: &str) -> String {
        self.dict.get_texture(name)
    }

    pub fn report_unused(&self) {
        self.dict.report_unused()
    }

    pub fn get_spectrum_texture(
        &self,
        name: &str,
        val: Option<Spectrum>,
        stype: SpectrumType,
    ) -> Option<Arc<SpectrumTexture>> {
        let tex = self.get_spectrum_texture_or_null(name, stype);
        if tex.is_some() {
            tex
        } else if val.is_some() {
            Some(Arc::new(SpectrumTexture::SpectrumConstant(
                SpectrumConstantTexture::new(val.unwrap()),
            )))
        } else {
            None
        }
    }

    pub fn get_float_texture(&self, name: &str, val: Float) -> Arc<FloatTexture> {
        if let Some(tex) = self.get_float_texture_or_null(name) {
            return tex;
        } else {
            return Arc::new(FloatTexture::FloatConstant(FloatConstantTexture::new(val)));
        }
    }

    fn get_spectrum_texture_or_null(
        &self,
        name: &str,
        stype: SpectrumType,
    ) -> Option<Arc<SpectrumTexture>> {
        for p in &self.dict.params {
            if p.name != name {
                continue;
            }

            match p.type_name.as_str() {
                "texture" => {
                    if p.strings.len() != 1 {
                        panic!(
                            "[{:?}] Expected 1 texture name for {}, found {}",
                            p.loc,
                            name,
                            p.strings.len()
                        );
                    }

                    p.looked_up.store(true, Ordering::Relaxed);
                    let tex_name = &p.strings[0];

                    if let Some(nt) = &self.textures {
                        let map = match stype {
                            SpectrumType::Unbounded => &nt.unbounded_spectrum_textures,
                            SpectrumType::Albedo => &nt.albedo_spectrum_textures,
                            _ => &nt.illuminant_spectrum_textures,
                        };

                        if let Some(tex) = map.get(tex_name) {
                            return Some(Arc::clone(tex));
                        }
                        panic!(
                            "[{:?}] Couldn't find spectrum texture named '{}'",
                            p.loc, tex_name
                        );
                    }
                    return None;
                }

                "rgb" => {
                    if p.floats.len() != 3 {
                        panic!("[{:?}] RGB parameter '{}' needs 3 floats", p.loc, name);
                    }
                    p.looked_up.store(true, Ordering::Relaxed);

                    let rgb = RGB::new(p.floats[0], p.floats[1], p.floats[2]);
                    if rgb.r < 0.0 || rgb.g < 0.0 || rgb.b < 0.0 {
                        panic!("[{:?}] Negative RGB values for '{}'", p.loc, name);
                    }

                    let cs = self.dict.color_space.as_ref().unwrap();
                    let s: Spectrum = match stype {
                        SpectrumType::Illuminant => {
                            Spectrum::RGBIlluminant(RGBIlluminantSpectrum::new(cs, rgb))
                        }
                        SpectrumType::Unbounded => {
                            Spectrum::RGBUnbounded(RGBUnboundedSpectrum::new(cs, rgb))
                        }
                        SpectrumType::Albedo => {
                            if rgb.r > 1.0 || rgb.g > 1.0 || rgb.b > 1.0 {
                                panic!("[{:?}] Albedo RGB > 1 for '{}'", p.loc, name);
                            }
                            Spectrum::RGBAlbedo(RGBAlbedoSpectrum::new(cs, rgb))
                        }
                    };
                    return Some(Arc::new(SpectrumTexture::SpectrumConstant(
                        SpectrumConstantTexture::new(s),
                    )));
                }

                "spectrum" | "blackbody" => {
                    let s = self.dict.get_one_spectrum(name, None, stype)?;
                    return Some(Arc::new(SpectrumTexture::SpectrumConstant(
                        SpectrumConstantTexture::new(s),
                    )));
                }

                _ => {}
            }
        }
        return None;
    }

    fn get_float_texture_or_null(&self, name: &str) -> Option<Arc<FloatTexture>> {
        for p in &self.dict.params {
            if p.name != name {
                continue;
            }

            match p.type_name.as_str() {
                "texture" => {
                    if p.strings.len() != 1 {
                        panic!(
                            "[{:?}] Expected 1 texture name for {}, found {}",
                            p.loc,
                            name,
                            p.strings.len()
                        );
                    }

                    p.looked_up.store(true, Ordering::Relaxed);
                    let tex_name = &p.strings[0];

                    if let Some(nt) = &self.textures {
                        let map = &nt.float_textures;
                        if let Some(tex) = map.get(tex_name) {
                            return Some(Arc::clone(tex));
                        }
                        panic!(
                            "[{:?}] Couldn't find float texture named '{}'",
                            p.loc, tex_name
                        );
                    }
                    return None;
                }
                "float" => {
                    let v = self.get_one_float(name, 0.);
                    return Some(Arc::new(FloatTexture::FloatConstant(
                        FloatConstantTexture::new(v),
                    )));
                }
                _ => {
                    panic!("[{:?}] Couldn't find float texture", p.loc);
                }
            }
        }
        return None;
    }
}