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

#[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> {
        match s.chars().nth(0) {
            Some('$') => Ok(Variable {
                r#type: VariableType::Integer,
                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 {
    Integer,
    Expr
}

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

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::Int(i1), Atom::Int(i2)) => i1.cmp(i2),
            (Atom::Int(_), _) => Ordering::Less,
            (_, Atom::Int(_)) => Ordering::Greater,

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

            (Atom::Ellipsis, Atom::Ellipsis) => Ordering::Equal,
            (Atom::Ellipsis, _) => Ordering::Greater,
        }
    }
}

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

#[derive(Debug, PartialEq, Clone, Eq, Hash)]
pub enum Expr {
    Atom(Atom),
    Application(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(_)) => Ordering::Less,
            (Expr::Application(_), Expr::Atom(_)) => 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))
            }
        }
    }
}

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(" "))
            }
        }
    }
}

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

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

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_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_atom(&mut self) -> Expr {
        let mut buffer = String::new();

        while let Some(&c) = self.input.peek() {
            if c.is_whitespace() || 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 ellipsis
        if buffer == "." {
            return Expr::Atom(Atom::Ellipsis);
        }

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

        // Check for integer
        if let Ok(i) = buffer.parse::<i64>() {
            return Expr::Atom(Atom::Int(i));
        }

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