//! OBJ file format parsing.
//!
//! File format documentation: <https://paulbourke.net/dataformats/obj>

use std::collections::HashMap;
use std::fs::read_to_string;
use std::path::PathBuf;

use engine_ecs::util::Either;

use crate::file_format::wavefront::common::{
    keyword,
    parse_statement_line,
    Keyword as _,
    ParsingError,
    Statement,
    Triplet,
};
use crate::mesh::vertex_buffer::{
    NamedVertexAttr,
    VertexAttrInfo,
    VertexBuffer as MeshVertexBuffer,
    VertexLabel,
};
use crate::mesh::{Mesh, VertexAttrType};
use crate::vector::{Vec2, Vec3};

/// Parses the content of a Wavefront `.obj`.
///
/// # Errors
/// Will return `Err` if the `.obj` content is formatted incorrectly.
pub fn parse(obj_content: &str) -> Result<Obj, Error>
{
    let lines = obj_content
        .lines()
        .enumerate()
        .map(|(line_index, line)| (line_index + 1, line));

    let mut item_counts = ItemCounts::default();

    for (line_index, line) in obj_content.lines().enumerate() {
        if line.is_empty() || line.starts_with('#') {
            continue;
        }

        let Some((keyword, _)) = line.split_once(" ") else {
            continue;
        };

        let Ok(keyword) = Keyword::from_str(keyword, line_index + 1) else {
            continue;
        };

        item_counts.increment_item_cnt_by_keyword(keyword);
    }

    let mut obj = Obj::with_capacities_from_item_cnts(item_counts);

    let mut parsing_state = ParsingState::default();

    for (line_no, line) in lines {
        let Some(statement) = parse_statement_line::<Keyword>(line, line_no)? else {
            continue;
        };

        obj.handle_statement(line_no, statement, &mut parsing_state)?;
    }

    Ok(obj)
}

/// The data of a Wavefront object file.
#[derive(Debug)]
#[non_exhaustive]
pub struct Obj
{
    pub vertex_positions: Vec<Vec3<f32>>,
    pub vertex_normals: Vec<Vec3<f32>>,
    pub texture_positions: Vec<Vec2<f32>>,
    pub faces: Vec<Face>,
    pub mtl_libs: Vec<PathBuf>,
    pub unique_used_material_names: Vec<Box<str>>,
}

impl Obj
{
    /// Creates a [`Mesh`] from this Wavefront object file data.
    ///
    /// # Errors
    /// Returns `Err` if:
    /// - A face's vertex position cannot be found
    /// - A face's texture position cannot be found
    /// - A face's vertex normal cannot be found
    /// - A face index does not fit in a [`u32`]
    pub fn to_mesh(&self) -> Result<Mesh, Error>
    {
        let mut vertex_buf = MeshVertexBuffer::with_capacity(
            &[
                VertexAttrInfo {
                    label: VertexLabel::Position,
                    ty: VertexAttrType::Float32Array { length: 3 },
                },
                VertexAttrInfo {
                    label: VertexLabel::Uv,
                    ty: VertexAttrType::Float32Array { length: 2 },
                },
                VertexAttrInfo {
                    label: VertexLabel::Normal,
                    ty: VertexAttrType::Float32Array { length: 3 },
                },
            ],
            self.faces.len() * 3,
        );

        let mut indices = Vec::<u32>::with_capacity(self.faces.len() * 3);

        let mut added_face_vertices =
            HashMap::<FaceVertex, u32>::with_capacity(self.faces.len() * 3);

        for face in &self.faces {
            for face_vertex in &face.vertices {
                if let Some(index) = added_face_vertices.get(&face_vertex) {
                    indices.push(*index);

                    continue;
                }

                let pos = self
                    .vertex_positions
                    .get(face_vertex.position as usize - 1)
                    .ok_or(Error::FaceVertexPositionNotFound {
                        vertex_pos_index: face_vertex.position,
                    })?
                    .clone();

                let texture_pos = face_vertex.texture.map_or_else(
                    || {
                        if !self.texture_positions.is_empty() {
                            tracing::warn!(concat!(
                                "Wavefront OBJ has texture coordinates ",
                                "but face vertex does not specify one"
                            ));
                        }

                        Ok(Vec2::default())
                    },
                    |face_vertex_texture| {
                        self.texture_positions
                            .get(face_vertex_texture as usize - 1)
                            .ok_or(Error::FaceTexturePositionNotFound {
                                texture_pos_index: face_vertex_texture,
                            })
                            .cloned()
                    },
                )?;

                let normal = face_vertex.normal.map_or_else(
                    || {
                        if !self.vertex_normals.is_empty() {
                            tracing::warn!(concat!(
                                "Wavefront OBJ has normals ",
                                "but face vertex does not specify one"
                            ));
                        }

                        Ok(Vec3::default())
                    },
                    |face_vertex_normal| {
                        self.vertex_normals
                            .get(face_vertex_normal as usize - 1)
                            .ok_or(Error::FaceVertexNormalNotFound {
                                vertex_normal_index: face_vertex_normal,
                            })
                            .cloned()
                    },
                )?;

                vertex_buf.push((
                    NamedVertexAttr {
                        label: VertexLabel::Position,
                        value: pos.into_array(),
                    },
                    NamedVertexAttr {
                        label: VertexLabel::Uv,
                        value: texture_pos.into_array(),
                    },
                    NamedVertexAttr {
                        label: VertexLabel::Normal,
                        value: normal.into_array(),
                    },
                ));

                let vertex_index = vertex_buf.len() - 1;

                let vertex_index = u32::try_from(vertex_index)
                    .map_err(|_| Error::FaceIndexTooBig(vertex_index))?;

                indices.push(vertex_index);

                added_face_vertices.insert(face_vertex.clone(), vertex_index);
            }
        }

        Ok(Mesh::builder()
            .vertices(vertex_buf)
            .indices(indices)
            .build())
    }

    /// Reads and parses the material libraries of this `Obj`.
    ///
    /// # Errors
    /// Returns `Err` if:
    /// - Reading the contents of a material library fails
    /// - Parsing a material library fails
    pub fn read_and_parse_material_libs<ParsedMaterial, ParseError>(
        &self,
        parse_material_lib: impl Fn(&str) -> Result<Vec<ParsedMaterial>, ParseError>,
    ) -> Result<Vec<ParsedMaterial>, MaterialLibsError<ParseError>>
    {
        self.mtl_libs
            .iter()
            .map(|mtl_lib| {
                Ok(parse_material_lib(&read_to_string(mtl_lib).map_err(
                    |err| MaterialLibsError::ReadingMaterialLibFailed {
                        source: err,
                        material_lib: mtl_lib.clone(),
                    },
                )?)?)
            })
            .flat_map(|res| match res {
                Ok(inner) => Either::A(inner.into_iter().map(Ok)),
                Err(err) => Either::B(vec![Err(err)].into_iter()),
            })
            .collect::<Result<Vec<_>, _>>()
    }

    fn with_capacities_from_item_cnts(item_counts: ItemCounts) -> Self
    {
        Self {
            vertex_positions: Vec::with_capacity(item_counts.pos_cnt),
            vertex_normals: Vec::with_capacity(item_counts.normal_cnt),
            texture_positions: Vec::with_capacity(item_counts.uv_cnt),
            faces: Vec::with_capacity(item_counts.face_cnt),
            mtl_libs: Vec::with_capacity(item_counts.mtl_lib_cnt),
            unique_used_material_names: Vec::with_capacity(
                item_counts.material_usage_cnt,
            ),
        }
    }

    fn handle_statement(
        &mut self,
        line_no: usize,
        statement: Statement<Keyword>,
        parsing_state: &mut ParsingState,
    ) -> Result<(), Error>
    {
        match statement.keyword {
            Keyword::V => self.handle_position_statement(line_no, statement),
            Keyword::Vn => self.handle_normal_statement(line_no, statement),
            Keyword::Vt => self.handle_uv_statement(line_no, statement),
            Keyword::F => self.handle_face_statement(line_no, statement, parsing_state),
            Keyword::Mtllib => self.handle_material_lib_statement(line_no, statement),
            Keyword::Usemtl => {
                self.handle_material_usage_statement(line_no, statement, parsing_state)
            }
            Keyword::O | Keyword::S => Ok(()),
        }
    }

    fn handle_position_statement(
        &mut self,
        line_no: usize,
        statement: Statement<Keyword>,
    ) -> Result<(), Error>
    {
        debug_assert_eq!(statement.keyword, Keyword::V);

        if statement.arguments.len() == 4 {
            return Err(Error::UnsupportedArgumentCount {
                keyword: statement.keyword.to_string(),
                arg_count: statement.arguments.len(),
                line_no: line_no,
            });
        }

        if statement.arguments.len() > 4 {
            return Err(Error::InvalidArgumentCount {
                keyword: statement.keyword.to_string(),
                arg_count: statement.arguments.len(),
                line_no: line_no,
            });
        }

        let x = statement.get_float_arg(0, line_no)?;
        let y = statement.get_float_arg(1, line_no)?;
        let z = statement.get_float_arg(2, line_no)?;

        self.vertex_positions.push(Vec3 { x, y, z });

        Ok(())
    }

    fn handle_uv_statement(
        &mut self,
        line_no: usize,
        statement: Statement<Keyword>,
    ) -> Result<(), Error>
    {
        debug_assert_eq!(statement.keyword, Keyword::Vt);

        if statement.arguments.len() == 3 {
            return Err(Error::UnsupportedArgumentCount {
                keyword: statement.keyword.to_string(),
                arg_count: statement.arguments.len(),
                line_no: line_no,
            });
        }

        if statement.arguments.len() > 3 {
            return Err(Error::InvalidArgumentCount {
                keyword: statement.keyword.to_string(),
                arg_count: statement.arguments.len(),
                line_no: line_no,
            });
        }

        let u = statement.get_float_arg(0, line_no)?;
        let v = statement.get_float_arg(1, line_no)?;

        self.texture_positions.push(Vec2 { x: u, y: v });

        Ok(())
    }

    fn handle_normal_statement(
        &mut self,
        line_no: usize,
        statement: Statement<Keyword>,
    ) -> Result<(), Error>
    {
        debug_assert_eq!(statement.keyword, Keyword::Vn);

        if statement.arguments.len() > 3 {
            return Err(Error::InvalidArgumentCount {
                keyword: statement.keyword.to_string(),
                arg_count: statement.arguments.len(),
                line_no: line_no,
            });
        }

        let i = statement.get_float_arg(0, line_no)?;
        let j = statement.get_float_arg(1, line_no)?;
        let k = statement.get_float_arg(2, line_no)?;

        self.vertex_normals.push(Vec3 { x: i, y: j, z: k });

        Ok(())
    }

    fn handle_material_usage_statement(
        &mut self,
        line_no: usize,
        statement: Statement<Keyword>,
        parsing_state: &mut ParsingState,
    ) -> Result<(), Error>
    {
        debug_assert_eq!(statement.keyword, Keyword::Usemtl);

        if statement.arguments.len() > 1 {
            return Err(Error::InvalidArgumentCount {
                keyword: statement.keyword.to_string(),
                arg_count: statement.arguments.len(),
                line_no: line_no,
            });
        }

        let material_name = statement
            .get_text_arg(0, line_no)?
            .to_owned()
            .into_boxed_str();

        if !self.unique_used_material_names.contains(&material_name) {
            self.unique_used_material_names.push(material_name.clone());
        }

        parsing_state.current_material_name = Some(material_name);

        Ok(())
    }

    fn handle_face_statement(
        &mut self,
        line_no: usize,
        statement: Statement<Keyword>,
        parsing_state: &mut ParsingState,
    ) -> Result<(), Error>
    {
        debug_assert_eq!(statement.keyword, Keyword::F);

        if statement.arguments.len() > 3 {
            return Err(Error::UnsupportedArgumentCount {
                keyword: statement.keyword.to_string(),
                arg_count: statement.arguments.len(),
                line_no: line_no,
            });
        }

        let vertex_a = statement.get_triplet_arg(0, line_no)?;
        let vertex_b = statement.get_triplet_arg(1, line_no)?;
        let vertex_c = statement.get_triplet_arg(2, line_no)?;

        self.faces.push(Face {
            vertices: [vertex_a.into(), vertex_b.into(), vertex_c.into()],
            material_name: parsing_state.current_material_name.clone(),
        });

        Ok(())
    }

    fn handle_material_lib_statement(
        &mut self,
        line_no: usize,
        statement: Statement<Keyword>,
    ) -> Result<(), Error>
    {
        debug_assert_eq!(statement.keyword, Keyword::Mtllib);

        for (index, value) in statement.arguments.iter().enumerate() {
            let mtl_lib = PathBuf::from(value.to_text(index, line_no)?);

            self.mtl_libs.push(mtl_lib);
        }

        Ok(())
    }
}

#[derive(Debug)]
#[non_exhaustive]
pub struct Face
{
    pub vertices: [FaceVertex; 3],
    pub material_name: Option<Box<str>>,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FaceVertex
{
    pub position: u32,
    pub texture: Option<u32>,
    pub normal: Option<u32>,
}

impl From<Triplet> for FaceVertex
{
    fn from(triplet: Triplet) -> Self
    {
        Self {
            position: triplet.0,
            texture: triplet.1,
            normal: triplet.2,
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum Error
{
    #[error(transparent)]
    ParsingError(#[from] ParsingError),

    #[error(
        "Face vertex position with index {vertex_pos_index} (1-based) was not found"
    )]
    FaceVertexPositionNotFound
    {
        vertex_pos_index: u32
    },

    #[error(
        "Face texture position with index {texture_pos_index} (1-based) was not found"
    )]
    FaceTexturePositionNotFound
    {
        texture_pos_index: u32
    },

    #[error(
        "Face vertex normal with index {vertex_normal_index} (1-based) was not found"
    )]
    FaceVertexNormalNotFound
    {
        vertex_normal_index: u32
    },

    #[error("Face index {0} is too big to fit into a 32-bit integer")]
    FaceIndexTooBig(usize),

    #[error(
        "Unsupported number of arguments ({arg_count}) to {keyword} at line {line_no}"
    )]
    UnsupportedArgumentCount
    {
        keyword: String,
        arg_count: usize,
        line_no: usize,
    },

    #[error("Invalid number of arguments ({arg_count}) to {keyword} at line {line_no}")]
    InvalidArgumentCount
    {
        keyword: String,
        arg_count: usize,
        line_no: usize,
    },
}

#[derive(Debug, thiserror::Error)]
pub enum MaterialLibsError<MaterialLibParseError>
{
    #[error("Parsing material library failed")]
    ParsingFailed(#[from] MaterialLibParseError),

    #[error("Failed to read material library {}", material_lib.display())]
    ReadingMaterialLibFailed
    {
        #[source]
        source: std::io::Error,
        material_lib: PathBuf,
    },
}

#[derive(Debug, Default)]
struct ParsingState
{
    current_material_name: Option<Box<str>>,
}

#[derive(Debug, Default)]
struct ItemCounts
{
    pos_cnt: usize,
    uv_cnt: usize,
    normal_cnt: usize,
    face_cnt: usize,
    mtl_lib_cnt: usize,
    material_usage_cnt: usize,
}

impl ItemCounts
{
    fn increment_item_cnt_by_keyword(&mut self, keyword: Keyword)
    {
        match keyword {
            Keyword::V => self.pos_cnt += 1,
            Keyword::Vn => self.normal_cnt += 1,
            Keyword::Vt => self.uv_cnt += 1,
            Keyword::F => self.face_cnt += 1,
            Keyword::Mtllib => self.mtl_lib_cnt += 1,
            Keyword::Usemtl => self.material_usage_cnt += 1,
            Keyword::O | Keyword::S => {}
        }
    }
}

keyword! {
    #[derive(Debug, PartialEq, Eq, Clone, Copy)]
    enum Keyword {
        #[keyword(rename = "v")]
        V,

        #[keyword(rename = "vn")]
        Vn,

        #[keyword(rename = "vt")]
        Vt,

        #[keyword(rename = "o")]
        O,

        #[keyword(rename = "s")]
        S,

        #[keyword(rename = "f")]
        F,

        #[keyword(rename = "mtllib")]
        Mtllib,

        #[keyword(rename = "usemtl")]
        Usemtl,
    }
}

#[cfg(test)]
mod tests
{
    use super::parse;

    #[test]
    fn parse_containing_usemtl_works()
    {
        let obj = parse(concat!(
            "usemtl dark-green\n",
            "f 6/8/1 7/5/3 1/3/2\n",
            "f 1/2/4 4/1/8 3/4/2\n",
            "usemtl light-pink\n",
            "f 9/1/7 5/1/2 3/2/7"
        ))
        .expect("Expected Ok");

        assert_eq!(obj.faces.len(), 3);

        assert_eq!(
            obj.faces[0]
                .material_name
                .as_deref()
                .expect("Expected Some"),
            "dark-green"
        );

        assert_eq!(
            obj.faces[1]
                .material_name
                .as_deref()
                .expect("Expected Some"),
            "dark-green"
        );

        assert_eq!(
            obj.faces[2]
                .material_name
                .as_deref()
                .expect("Expected Some"),
            "light-pink"
        );
    }
}