use std::cmp::Ordering;
use std::hash::{Hash, Hasher};
use std::iter::Peekable;
use std::str::{Chars, FromStr};

#[derive(Debug, Clone, Copy)]
pub struct Number(pub f64);

impl PartialEq for Number {
    fn eq(&self, other: &Self) -> bool {
        self.0.to_bits() == other.0.to_bits()
    }
}

impl Eq for Number {}

impl PartialOrd for Number {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Number {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.total_cmp(&other.0)
    }
}

impl Hash for Number {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.to_bits().hash(state);
    }
}

impl std::fmt::Display for Number {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.0.fract() == 0.0 {
            write!(f, "{:.0}", self.0)
        } else {
            write!(f, "{}", self.0)
        }
    }
}

#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct Variable {
    pub r#type: VariableType,
    pub index: String
}

impl FromStr for Variable {
    type Err = ();
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if let Some(index) = s.strip_prefix("..") {
            if index.is_empty() {
                return Err(());
            }

            return Ok(Variable {
                r#type: VariableType::Ellipsis,
                index: index.to_string(),
            });
        }

        if let Some(index) = s.strip_prefix('!') {
            if index.is_empty() {
                return Err(());
            }

            return Ok(Variable {
                r#type: VariableType::NonNumberExpr,
                index: index.to_string(),
            });
        }

        match s.chars().nth(0) {
            Some('$') => Ok(Variable {
                r#type: VariableType::Number,
                index: s[1..].to_string()
            }),
            Some('@') => Ok(Variable {
                r#type: VariableType::Expr,
                index: s[1..].to_string()
            }),
            _ => Err(()),
        }
    }
}

#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum VariableType {
    Number,
    Expr,
    NonNumberExpr,
    Ellipsis,
}

#[derive(Debug, PartialEq, Clone, Eq, Hash)]
pub enum Atom {
    Number(Number),
    Builtin(String),
    Variable(Variable),
}

impl PartialOrd for Atom {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Atom {
    fn cmp(&self, other: &Self) -> Ordering {
        match (self, other) {
            (Atom::Number(i1), Atom::Number(i2)) => i1.cmp(i2),
            (Atom::Number(_), _) => Ordering::Less,
            (_, Atom::Number(_)) => Ordering::Greater,

            (Atom::Builtin(b1), Atom::Builtin(b2)) => b1.cmp(b2),
            (Atom::Builtin(_), Atom::Variable(_)) => Ordering::Less,

            (Atom::Variable(v1), Atom::Variable(v2)) => {
                v1.index.cmp(&v2.index)
            }
            (Atom::Variable(_), Atom::Builtin(_)) => Ordering::Greater,
        }
    }
}

impl std::fmt::Display for Atom {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Atom::Number(i) => write!(f, "{}", i),
            Atom::Builtin(b) => write!(f, "{}", b),
            Atom::Variable(v) => write!(f, "{}{}", 
                match v.r#type {
                    VariableType::Number => "$",
                    VariableType::Expr => "@",
                    VariableType::NonNumberExpr => "!",
                    VariableType::Ellipsis => "..",
                },
                v.index
            ),
        }
    }
}

#[derive(Debug, PartialEq, Clone, Eq, Hash)]
pub enum Expr {
    Atom(Atom),
    Application(Vec<Expr>),
    OrderedList(Vec<Expr>),
}

impl PartialOrd for Expr {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Expr {
    fn cmp(&self, other: &Self) -> Ordering {
        match (self, other) {
            (Expr::Atom(_), Expr::Application(_) | Expr::OrderedList(_)) => Ordering::Less,
            (Expr::Application(_) | Expr::OrderedList(_), Expr::Atom(_)) => Ordering::Greater,
            (Expr::Application(_), Expr::OrderedList(_)) => Ordering::Less,
            (Expr::OrderedList(_), Expr::Application(_)) => Ordering::Greater,
            (Expr::Atom(a1), Expr::Atom(a2)) => a1.cmp(a2),
            (Expr::Application(v1), Expr::Application(v2)) => {
                v1.len().cmp(&v2.len())
                    .then_with(|| v1.cmp(v2))
            }
            (Expr::OrderedList(v1), Expr::OrderedList(v2)) => {
                v1.len().cmp(&v2.len())
                    .then_with(|| v1.cmp(v2))
            }
        }
    }
}

impl std::fmt::Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Expr::Atom(a) => write!(f, "{}", a),
            Expr::Application(args) => {
                let pieces: Vec<String> = args.into_iter().map(ToString::to_string).collect();
                write!(f, "({})", pieces.join(" "))
            }
            Expr::OrderedList(args) => {
                let pieces: Vec<String> = args.into_iter().map(ToString::to_string).collect();
                write!(f, "[{}]", pieces.join(" "))
            }
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
enum ParseResult {
    Expr(Expr),
    Eof,
}

pub struct Parser<'a> {
    input: Peekable<Chars<'a>>,
}

#[cfg(test)]
mod tests {
    use super::{Atom, Expr, Number, Parser, Variable, VariableType};

    fn builtin(name: &str) -> Expr {
        Expr::Atom(Atom::Builtin(name.to_string()))
    }

    fn number(i: f64) -> Expr {
        Expr::Atom(Atom::Number(Number(i)))
    }

    #[test]
    fn parses_ordered_list() {
        let mut parser = Parser::new("[a b 3]");
        let parsed = parser.parse_one();

        assert_eq!(
            parsed,
            Expr::OrderedList(vec![builtin("a"), builtin("b"), number(3.0)])
        );
    }

    #[test]
    fn parses_nested_ordered_list_and_application() {
        let mut parser = Parser::new("(+ [a b] (* [1 2] 3))");
        let parsed = parser.parse_one();

        assert_eq!(
            parsed,
            Expr::Application(vec![
                builtin("+"),
                Expr::OrderedList(vec![builtin("a"), builtin("b")]),
                Expr::Application(vec![
                    builtin("*"),
                    Expr::OrderedList(vec![number(1.0), number(2.0)]),
                    number(3.0),
                ]),
            ])
        );
    }

    #[test]
    fn parses_decimal_number() {
        let mut parser = Parser::new("1.25");
        let parsed = parser.parse_one();

        assert_eq!(parsed, number(1.25));
    }

    #[test]
    fn parses_non_number_variable() {
        let mut parser = Parser::new("!term");
        let parsed = parser.parse_one();

        assert_eq!(
            parsed,
            Expr::Atom(Atom::Variable(Variable {
                r#type: VariableType::NonNumberExpr,
                index: "term".to_string(),
            }))
        );
    }
}

impl<'a> Parser<'a> {
    pub fn new(input: &'a str) -> Self {
        Parser {
            input: input.chars().peekable(),
        }
    }

    pub fn parse_all(&mut self) -> Vec<Expr> {
        let mut expressions = Vec::new();
        loop {
            match self.parse_one_optional() {
                ParseResult::Expr(expr) => expressions.push(expr),
                ParseResult::Eof => break,
            }
        }
        expressions
    }

    fn parse_one_optional(&mut self) -> ParseResult {
        self.skip_whitespace();

        match self.input.peek() {
            Some('(') => ParseResult::Expr(self.parse_list()),
            Some('[') => ParseResult::Expr(self.parse_ordered_list()),
            Some(_) => ParseResult::Expr(self.parse_atom()),
            None => ParseResult::Eof,
        }
    }

    pub fn parse_one(&mut self) -> Expr {
        match self.parse_one_optional() {
            ParseResult::Expr(expr) => expr,
            ParseResult::Eof => panic!("Unexpected end of input during parse!"),
        }
    }

    fn skip_whitespace(&mut self) {
        while let Some(&c) = self.input.peek() {
            if c.is_whitespace() {
                self.input.next();
            } else {
                break;
            }
        }
    }

    fn parse_list(&mut self) -> Expr {
        if self.input.next() != Some('(') {
            panic!("Parser error: Expected '('");
        }

        let mut expressions = Vec::new();

        loop {
            self.skip_whitespace();

            match self.input.peek() {
                Some(')') => {
                    self.input.next();
                    break;
                }
                None => panic!("Parser error: Unexpected EOF, unclosed list"),
                Some(_) => {
                    match self.parse_one_optional() {
                        ParseResult::Expr(expr) => expressions.push(expr),
                        ParseResult::Eof => panic!("Parser error: Unexpected EOF after space inside list"),
                    }
                }
            }
        }

        Expr::Application(expressions)
    }

    fn parse_ordered_list(&mut self) -> Expr {
        if self.input.next() != Some('[') {
            panic!("Parser error: Expected '['");
        }

        let mut expressions = Vec::new();

        loop {
            self.skip_whitespace();

            match self.input.peek() {
                Some(']') => {
                    self.input.next();
                    break;
                }
                None => panic!("Parser error: Unexpected EOF, unclosed ordered list"),
                Some(_) => {
                    match self.parse_one_optional() {
                        ParseResult::Expr(expr) => expressions.push(expr),
                        ParseResult::Eof => panic!(
                            "Parser error: Unexpected EOF after space inside ordered list"
                        ),
                    }
                }
            }
        }

        Expr::OrderedList(expressions)
    }

    fn parse_atom(&mut self) -> Expr {
        let mut buffer = String::new();

        while let Some(&c) = self.input.peek() {
            if c.is_whitespace() || c == '(' || c == ')' || c == '[' || c == ']' {
                break;
            }
            buffer.push(c);
            self.input.next();
        }

        if buffer.is_empty() {
            panic!("Parser error: Tried to parse an atom, but it was empty.");
        }

        // Check for variable
        if let Ok(v) = buffer.parse::<Variable>() {
            return Expr::Atom(Atom::Variable(v));
        }

        // Check for number
        if let Ok(i) = buffer.parse::<f64>() {
            return Expr::Atom(Atom::Number(Number(i)));
        }

        // Default to builtin
        Expr::Atom(Atom::Builtin(buffer))
    }
}