use miette::{bail, miette, Context, IntoDiagnostic, LabeledSpan, Result, Severity};
use nom::{
    branch::alt,
    bytes::complete::{tag, take_till, take_till1},
    character,
    error::Error,
    multi::{fill, many1},
    number::complete::double,
    Finish, IResult,
};
use ray_tracing_core::prelude::*;
use std::{collections::HashMap, path::Path};

#[derive(Debug)]
pub struct ObjData {
    pub vertices: Vec<Pos3>,
    pub normals: Vec<Dir3>,
    pub texcoord: Vec<[Float; 2]>,
    pub faces: Vec<Face>,
    pub materials: Vec<ObjMaterial>,
}

#[derive(Debug)]
pub struct ObjMaterial {
    pub name: String,
    pub map: HashMap<String, String>,
}

#[derive(Debug)]
pub struct Face {
    pub v: Vec<u32>,
    pub t: Option<Vec<u32>>,
    pub n: Option<Vec<u32>>,
    pub mat: Option<u32>,
}

pub struct TriangleFace {
    pub v: [u32; 3],
    pub t: Option<[u32; 3]>,
    pub n: Option<[u32; 3]>,
    pub mat: Option<u32>,
}

#[derive(Debug)]
pub struct RelativeFace {
    pub v: Vec<i32>,
    pub t: Option<Vec<i32>>,
    pub n: Option<Vec<i32>>,
}

#[derive(Debug)]
enum Line<'a> {
    Vertex(Pos3),
    Normal(Dir3),
    Texcoord([Float; 2]),
    Face(RelativeFace),
    Material(&'a str),
    Mtllib(&'a str),
    Empty,
}

fn remove_whitespace(i: &str) -> IResult<&str, ()> {
    let (i, _) = take_till(|c| c != ' ')(i)?;
    Ok((i, ()))
}

fn parse_float(i: &str) -> IResult<&str, Float> {
    let (i, _) = remove_whitespace(i)?;
    double(i).map(|(i, f)| (i, f as Float))
}

fn parse_vertex(i: &'_ str) -> IResult<&'_ str, Line<'_>> {
    let mut v = [0.0; 3];
    let (i, _) = fill(parse_float, &mut v)(i)?;
    Ok((i, Line::Vertex(Pos3::new(v[0], v[1], v[2]))))
}

fn parse_normal(i: &'_ str) -> IResult<&'_ str, Line<'_>> {
    let mut v = [0.0; 3];
    let (i, _) = fill(parse_float, &mut v)(i)?;
    Ok((i, Line::Normal(Dir3::new(v[0], v[1], v[2]))))
}

fn parse_texcoord(i: &'_ str) -> IResult<&'_ str, Line<'_>> {
    let mut v = [0.0; 2];
    let (i, _) = fill(parse_float, &mut v)(i)?;
    Ok((i, Line::Texcoord(v)))
}

fn parse_identifier(i: &str) -> IResult<&str, &str> {
    let (i, ()) = remove_whitespace(i)?;
    alt((tag("#"), take_till(|c| c == ' ')))(i)
}

fn face_vertex_empty(i: &str) -> IResult<&str, i32> {
    let (i, _) = remove_whitespace(i)?;
    let (i, v) = character::complete::i32(i)?;

    Ok((i, v))
}
fn face_vertex_normal(i: &str) -> IResult<&str, (i32, i32)> {
    let (i, _) = remove_whitespace(i)?;
    let (i, v) = character::complete::i32(i)?;
    let (i, _) = tag("//")(i)?;
    let (i, n) = character::complete::i32(i)?;

    Ok((i, (v, n)))
}
fn face_vertex_tex(i: &str) -> IResult<&str, (i32, i32)> {
    let (i, _) = remove_whitespace(i)?;
    let (i, v) = character::complete::i32(i)?;
    let (i, _) = tag("/")(i)?;
    let (i, t) = character::complete::i32(i)?;

    Ok((i, (v, t)))
}
fn face_vertex_normal_tex(i: &str) -> IResult<&str, (i32, i32, i32)> {
    let (i, _) = remove_whitespace(i)?;
    let (i, v) = character::complete::i32(i)?;
    let (i, _) = tag("/")(i)?;
    let (i, t) = character::complete::i32(i)?;
    let (i, _) = tag("/")(i)?;
    let (i, n) = character::complete::i32(i)?;

    Ok((i, (v, t, n)))
}

fn parse_face(i: &'_ str) -> IResult<&'_ str, Line<'_>> {
    if let Ok((i, s)) = many1(face_vertex_normal_tex)(i) {
        if s.len() >= 3 {
            Ok((
                i,
                Line::Face(RelativeFace {
                    v: s.iter().map(|t| t.0).collect(),
                    t: Some(s.iter().map(|t| t.1).collect()),
                    n: Some(s.iter().map(|t| t.2).collect()),
                }),
            ))
        } else {
            Err(nom::Err::Error(Error {
                input: i,
                code: nom::error::ErrorKind::Tag,
            }))
        }
    } else if let Ok((i, s)) = many1(face_vertex_normal)(i) {
        if s.len() >= 3 {
            Ok((
                i,
                Line::Face(RelativeFace {
                    v: s.iter().map(|t| t.0).collect(),
                    t: None,
                    n: Some(s.iter().map(|t| t.1).collect()),
                }),
            ))
        } else {
            Err(nom::Err::Error(Error {
                input: i,
                code: nom::error::ErrorKind::Tag,
            }))
        }
    } else if let Ok((i, s)) = many1(face_vertex_tex)(i) {
        if s.len() >= 3 {
            Ok((
                i,
                Line::Face(RelativeFace {
                    v: s.iter().map(|t| t.0).collect(),
                    t: Some(s.iter().map(|t| t.1).collect()),
                    n: None,
                }),
            ))
        } else {
            Err(nom::Err::Error(Error {
                input: i,
                code: nom::error::ErrorKind::Tag,
            }))
        }
    } else if let Ok((i, s)) = many1(face_vertex_empty)(i) {
        if s.len() >= 3 {
            Ok((
                i,
                Line::Face(RelativeFace {
                    v: s,
                    t: None,
                    n: None,
                }),
            ))
        } else {
            Err(nom::Err::Error(Error {
                input: i,
                code: nom::error::ErrorKind::Tag,
            }))
        }
    } else {
        Err(nom::Err::Error(Error {
            input: i,
            code: nom::error::ErrorKind::Tag,
        }))
    }
}

fn parse_mtllib(i: &'_ str) -> IResult<&'_ str, Line<'_>> {
    let (i, s) = take_till(|c| c == ' ')(i)?;
    Ok((i, Line::Mtllib(s)))
}

fn parse_mtl(i: &'_ str) -> IResult<&'_ str, Line<'_>> {
    let (i, s) = take_till(|c| c == ' ')(i)?;
    Ok((i, Line::Material(s)))
}

fn parse_line(i: &'_ str) -> IResult<&'_ str, Line<'_>> {
    let (i, ident) = parse_identifier(i)?;
    let (i, _) = remove_whitespace(i)?;
    match ident {
        "v" => parse_vertex(i),
        "#" | "" => Ok((i, Line::Empty)),
        "f" => parse_face(i),
        "o" => Ok((i, Line::Empty)),
        "g" => Ok((i, Line::Empty)),
        "vn" => parse_normal(i),
        "vt" => parse_texcoord(i),
        "usemtl" => parse_mtl(i),
        "mtllib" => parse_mtllib(i),
        _ => Err(nom::Err::Failure(nom::error::Error {
            input: i,
            code: nom::error::ErrorKind::Fail,
        })),
    }
}

fn handle_relative(x: i32, len: i32) -> Result<u32> {
    if x == 0 {
        Err(miette!("Index cannot be 0"))
    } else if x > 0 {
        if x <= len {
            Ok(x as u32 - 1)
        } else {
            Err(miette!("Index {x} out of bounds"))
        }
    } else if len + x >= 0 {
        Ok((len + x) as u32)
    } else {
        Err(miette!("Index {x} out of bounds"))
    }
}

enum MaterialsLine<'a> {
    Material(&'a str),
    KeyValue(&'a str, &'a str),
    Empty,
}

fn parse_new_material(i: &'_ str) -> IResult<&'_ str, MaterialsLine<'_>> {
    let (i, ident) = take_till1(|c| c == ' ')(i)?;
    Ok((i, MaterialsLine::Material(ident)))
}

fn parse_line_material(i: &'_ str) -> IResult<&'_ str, MaterialsLine<'_>> {
    let (i, ident) = parse_identifier(i)?;
    let (i, _) = remove_whitespace(i)?;
    match ident {
        "newmtl" => parse_new_material(i),
        "#" | "" => Ok((i, MaterialsLine::Empty)),
        _ => Ok((i, MaterialsLine::KeyValue(ident, i))),
    }
}

fn read_mtllib(path: impl AsRef<Path>) -> Result<Vec<ObjMaterial>> {
    let string = std::fs::read_to_string(path.as_ref())
        .into_diagnostic()
        .with_context(|| format!("opening file {:?}", path.as_ref()))?;

    let mut materials = Vec::new();

    let mut current_material = None;

    for (i, l) in string.lines().enumerate() {
        let t = match parse_line_material(l).finish() {
            Ok(t) => t.1,
            Err(e) => {
                return Err(miette!(
                    severity = Severity::Error,
                    code = e.code.description(),
                    labels = vec![LabeledSpan::at_offset(l.len() - e.input.len(), "here")],
                    "unable to parse line {}",
                    i + 1
                )
                .with_source_code(l.to_string()));
            }
        };

        match t {
            MaterialsLine::Material(name) => {
                if let Some(m) = current_material.take() {
                    materials.push(m);
                }

                current_material = Some(ObjMaterial {
                    name: name.to_string(),
                    map: HashMap::new(),
                });
            }
            MaterialsLine::KeyValue(k, v) => {
                if let Some(m) = current_material.as_mut() {
                    m.map.insert(k.to_string(), v.to_string());
                } else {
                    bail!("Material property appears multiple times in line {}", i + 1);
                }
            }
            MaterialsLine::Empty => (),
        }
    }

    if let Some(m) = current_material {
        materials.push(m);
    }

    Ok(materials)
}

impl ObjData {
    pub fn new(path: impl AsRef<Path>) -> Result<Self> {
        let string = std::fs::read_to_string(path.as_ref())
            .into_diagnostic()
            .with_context(|| format!("opening file {:?}", path.as_ref()))?;

        let mut vertices = Vec::new();
        let mut normals = Vec::new();
        let mut texcoord = Vec::new();
        let mut faces = Vec::new();
        let mut materials = Vec::new();
        let mut materials_map = HashMap::new();

        let mut current_material = None;

        for (i, l) in string.lines().enumerate() {
            let t = match parse_line(l).finish() {
                Ok(t) => t.1,
                Err(e) => {
                    return Err(miette!(
                        severity = Severity::Error,
                        code = e.code.description(),
                        labels = vec![LabeledSpan::at_offset(l.len() - e.input.len(), "here")],
                        "unable to parse line {}",
                        i + 1
                    )
                    .with_source_code(l.to_string()));
                }
            };

            match t {
                Line::Vertex(v) => vertices.push(v),
                Line::Normal(n) => normals.push(n),
                Line::Texcoord(t) => texcoord.push(t),
                Line::Face(f) => faces.push(Face {
                    v: f.v
                        .iter()
                        .map(|&c| {
                            handle_relative(c, vertices.len() as i32)
                                .with_context(|| format!("In line {}", i + 1))
                        })
                        .collect::<Result<Vec<_>>>()?,
                    t: f.t
                        .map(|t| {
                            t.iter()
                                .map(|&c| {
                                    handle_relative(c, texcoord.len() as i32)
                                        .with_context(|| format!("In line {}", i + 1))
                                })
                                .collect::<Result<Vec<_>>>()
                        })
                        .transpose()?,
                    n: f.n
                        .map(|t| {
                            t.iter()
                                .map(|&c| {
                                    handle_relative(c, normals.len() as i32)
                                        .with_context(|| format!("In line {}", i + 1))
                                })
                                .collect::<Result<Vec<_>>>()
                        })
                        .transpose()?,
                    mat: current_material,
                }),
                Line::Material(s) => {
                    current_material = Some(
                        *materials_map
                            .get(s)
                            .ok_or(miette!("Unknown material \"{s}\" in line {}", i + 1))?,
                    );
                }
                Line::Mtllib(s) => {
                    for m in read_mtllib(path.as_ref().parent().ok_or(miette!("Unable "))?.join(s))?
                    {
                        materials_map.insert(m.name.clone(), materials.len() as u32);
                        materials.push(m);
                    }
                }
                Line::Empty => (),
            }
        }

        Ok(Self {
            vertices,
            normals,
            texcoord,
            faces,
            materials,
        })
    }

    pub fn vertices(&self) -> &Vec<Pos3> {
        &self.vertices
    }

    pub fn faces(&self) -> &Vec<Face> {
        &self.faces
    }

    pub fn triangle_faces(&self) -> impl Iterator<Item = TriangleFace> + '_ {
        self.faces.iter().flat_map(|f| {
            (1..f.v.len() - 1).map(|i| TriangleFace {
                v: [f.v[0], f.v[i], f.v[i + 1]],
                t: f.t.as_ref().map(|t| [t[0], t[i], t[i + 1]]),
                n: f.n.as_ref().map(|n| [n[0], n[i], n[i + 1]]),
                mat: f.mat,
            })
        })
    }

    pub fn materials(&self) -> &Vec<ObjMaterial> {
        &self.materials
    }
}
