use pest::{iterators::Pair, Parser};
use pest_derive::Parser;

use crate::{expression::Expression, statement::Statement, value::Value};

#[derive(Parser)]
#[grammar = "grammar.pest"]
pub struct LangParser;

pub fn parse<S: ToString>(source: &S) -> (Statement, FuncHolder) {
    let source = source.to_string();
    let pairs = LangParser::parse(Rule::root_statement, &source).unwrap();
    handle_rules(pairs)
}
type Func = (String, Vec<String>, Vec<Statement>);

#[derive(Default)]
pub struct FuncHolder(Vec<Func>);

impl FuncHolder {
    fn insert(mut self, name: String, params: Vec<String>, body: Vec<Statement>) -> Self {
        self.0.push((name, params, body));
        self
    }
}

impl Extend<Func> for FuncHolder {
    fn extend<T>(&mut self, iter: T)
    where
        T: IntoIterator<Item = Func>,
    {
        self.0.extend(iter);
    }
}

impl Extend<FuncHolder> for FuncHolder {
    fn extend<T>(&mut self, iter: T)
    where
        T: IntoIterator<Item = FuncHolder>,
    {
        self.0
            .extend(iter.into_iter().flat_map(|h| h.0.into_iter()));
    }
}

impl IntoIterator for FuncHolder {
    type Item = Func;
    type IntoIter = std::vec::IntoIter<Func>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

fn handle_rules(mut p: pest::iterators::Pairs<Rule>) -> (Statement, FuncHolder) {
    let z = p.next().unwrap();
    parse_statement(z)
}

fn parse_statement(pair: Pair<Rule>) -> (Statement, FuncHolder) {
    match pair.as_rule() {
        Rule::root_statement => {
            let (z, f) = pair.into_inner().map(parse_statement).unzip();
            (Statement::Body(z), f)
        }
        Rule::while_block => {
            let mut inner = pair.into_inner();
            let expr = inner.next().unwrap();
            let body = inner.next().unwrap();
            (
                Statement::While(parse_expression(expr), Box::new(parse_statement(body).0)),
                FuncHolder::default(),
            )
        }
        Rule::if_block => {
            let mut inner = pair.into_inner();
            let mut the_list = vec![];
            let mut the_else = None;
            while let Some(expr) = inner.next() {
                if expr.as_rule() == Rule::expression {
                    let body = inner.next().unwrap();
                    the_list.push((parse_expression(expr), parse_statement(body).0));
                } else {
                    the_else = Some(Box::new(parse_statement(expr).0));
                }
            }
            (Statement::If(the_list, the_else), FuncHolder::default())
        }
        Rule::assignment => {
            let mut inner = pair.into_inner();
            let name = inner.next().unwrap();
            let expr = inner.next().unwrap();
            (
                Statement::Assignment(name.as_str().to_string(), parse_expression(expr), false),
                FuncHolder::default(),
            )
        }
        Rule::assignment_let => {
            let mut inner = pair.into_inner();
            let name = inner.next().unwrap();
            let expr = inner.next().unwrap();
            (
                Statement::Assignment(name.as_str().to_string(), parse_expression(expr), true),
                FuncHolder::default(),
            )
        }
        Rule::expression => (
            Statement::Expression(parse_expression(pair)),
            FuncHolder::default(),
        ),
        Rule::statement => {
            let inner = pair.into_inner().next().unwrap();
            parse_statement(inner)
        }
        Rule::block => {
            let (z, _): (_, FuncHolder) = pair.into_inner().map(parse_statement).unzip();
            (Statement::Body(z), FuncHolder::default())
        }
        Rule::func_block => {
            let (statements, expression): (Vec<_>, Vec<_>) = pair
                .into_inner()
                .partition(|x| x.as_rule() != Rule::expression);
            // Only functions defined at the root are supported for now, so we drop the ones
            // defined in a block.
            let (mut v, _): (Vec<Statement>, FuncHolder) =
                statements.into_iter().map(parse_statement).unzip();

            if let Some(expr) = expression.into_iter().next() {
                let ret = parse_expression(expr);
                v.push(Statement::Return(Some(ret)));
            }
            (Statement::Body(v), FuncHolder::default())
        }
        Rule::def => {
            let mut inner = pair.into_inner();
            let name = inner.next().unwrap();
            let args = inner
                .next()
                .unwrap()
                .into_inner()
                .map(|a| a.as_str().to_string())
                .collect();

            let body = inner.next().unwrap();
            let body = parse_statement(body);
            (
                Statement::Body(vec![]),
                FuncHolder::default().insert(name.as_str().to_string(), args, vec![body.0]),
            )
        }
        Rule::ret => {
            let mut inner = pair.into_inner();
            if let Some(expr) = inner.next() {
                (
                    Statement::Return(Some(parse_expression(expr))),
                    FuncHolder::default(),
                )
            } else {
                (Statement::Return(None), FuncHolder::default())
            }
        }
        Rule::EOI => (Statement::Body(vec![]), FuncHolder::default()),
        x => unreachable!("How did we get to {:?}", x),
    }
}

fn parse_expression(pair: Pair<Rule>) -> Expression {
    match pair.as_rule() {
        Rule::expression => {
            let mut inner = pair.into_inner();
            parse_expression(inner.next().unwrap())
        }
        Rule::equality | Rule::inequality | Rule::sum | Rule::product => {
            let mut inner = pair.into_inner();
            let lhs = parse_expression(inner.next().unwrap());
            if let Some(op) = inner.next() {
                let rhs = parse_expression(inner.next().unwrap());
                Expression::Binary(Box::new(lhs), op.as_str().parse().unwrap(), Box::new(rhs))
            } else {
                lhs
            }
        }
        Rule::value => {
            let mut inner = pair.into_inner();
            let v = parse_value(inner.next().unwrap());
            if let Some((name, args)) = method_call(inner.next().unwrap()) {
                Expression::CallMethod(Box::new(v), name, args)
            } else {
                v
            }
        }
        Rule::if_block_expr => {
            let mut inner = pair.into_inner();
            let mut the_list = vec![];
            let mut the_else = None;
            while let Some(expr) = inner.next() {
                if let Some(if_block) = inner.next() {
                    let expr = parse_expression(expr);
                    let if_block_body = if_block.into_inner();
                    let (a, b): (Vec<_>, Vec<_>) =
                        if_block_body.partition(|a| a.as_rule() == Rule::statement);
                    let (statements, _): (_, FuncHolder) =
                        a.into_iter().map(parse_statement).unzip();
                    let z = b.into_iter().map(parse_expression).collect::<Vec<_>>();
                    the_list.push((expr, Statement::Body(statements), z[0].clone()));
                } else {
                    let if_block_body = expr.into_inner();
                    let (a, b): (Vec<_>, Vec<_>) =
                        if_block_body.partition(|a| a.as_rule() == Rule::statement);
                    let (statements, _): (_, FuncHolder) =
                        a.into_iter().map(parse_statement).unzip();
                    let z = b.into_iter().map(parse_expression).collect::<Vec<_>>();
                    the_else = Some((
                        Box::new(Statement::Body(statements)),
                        Box::new(z[0].clone()),
                    ));
                }
            }

            Expression::If(the_list, the_else.unwrap())
        }
        _ => unreachable!(),
    }
}

fn method_call(pair: Pair<Rule>) -> Option<(String, Vec<Expression>)> {
    match pair.as_rule() {
        Rule::method_call => {
            let mut inner = pair.into_inner();
            if let Some(name) = inner.next() {
                let args = inner
                    .next()
                    .unwrap()
                    .into_inner()
                    .map(parse_expression)
                    .collect();
                Some((name.as_str().to_string(), args))
            } else {
                None
            }
        }
        _ => None,
    }
}

fn parse_value(pair: Pair<Rule>) -> Expression {
    match pair.as_rule() {
        Rule::const_value => {
            Expression::Constant(parse_const_value(&pair.into_inner().next().unwrap()))
        }
        Rule::identifier => Expression::Variable(pair.as_str().to_string()),
        Rule::call => {
            let mut inner = pair.into_inner();
            let name = inner.next().unwrap();
            let args = inner
                .next()
                .unwrap()
                .into_inner()
                .map(parse_expression)
                .collect();
            Expression::Call(name.as_str().to_string(), args)
        }
        Rule::parens => {
            let mut inner = pair.into_inner();
            parse_expression(inner.next().unwrap())
        }
        _ => unreachable!(),
    }
}

fn parse_const_value(pair: &Pair<Rule>) -> Value {
    match pair.as_rule() {
        Rule::number => {
            if let Ok(f) = pair.as_str().parse::<i32>() {
                Value::Int(f)
            } else if let Ok(f) = pair.as_str().parse::<f32>() {
                Value::Float(f)
            } else {
                unreachable!()
            }
        }
        Rule::string => Value::String(pair.as_str().to_string()),
        Rule::boolean => Value::Bool(pair.as_str() == "true"),
        Rule::char => Value::Char(pair.as_str().chars().next().unwrap()),
        _ => unreachable!(),
    }
}