use std::fs;
use std::path::Path;

use pest::error::Error;
use pest::iterators::Pair;
use pest::pratt_parser::PrattParser;
use pest::Parser;

use crate::ast::untyped::module::Module;
use crate::ast::untyped::*;
use crate::ast::{Import, ModulePath};
use crate::typing::Type;

#[derive(pest_derive::Parser)]
#[grammar = "parsing/backend/pest/grammar.pest"]
struct LilaParser;

use lazy_static;
lazy_static::lazy_static! {
    static ref PRATT_PARSER: PrattParser<Rule> = {
        use pest::pratt_parser::{Assoc::*, Op};
        use Rule::*;

        // Precedence is defined lowest to highest
        PrattParser::new()
            .op(Op::infix(equal, Left) | Op::infix(not_equal, Left))
            .op(Op::infix(add, Left) | Op::infix(subtract, Left))
            .op(Op::infix(modulo, Left))
            .op(Op::infix(multiply, Left) | Op::infix(divide, Left))
    };
}

pub fn parse_file(path: &Path) -> Result<Module, Error<Rule>> {
    let source = fs::read_to_string(&path).expect("could not read source file");
    let module_path = ModulePath::from(path);
    let mut module = parse_as_module(&source, module_path)?;
    module.file = Some(path.to_owned());
    Ok(module)
}

pub fn parse_as_module(source: &str, path: ModulePath) -> Result<Module, Error<Rule>> {
    let mut pairs = LilaParser::parse(Rule::source_file, &source)?;

    assert!(pairs.len() == 1);
    let module = parse_module(pairs.next().unwrap().into_inner().next().unwrap(), path);

    Ok(module)
}

pub fn parse_module(pair: Pair<Rule>, path: ModulePath) -> Module {
    assert!(pair.as_rule() == Rule::module_items);

    let mut module = Module::new(path);

    let pairs = pair.into_inner();
    for pair in pairs {
        match pair.as_rule() {
            Rule::definition => {
                let def = parse_definition(pair.into_inner().next().unwrap());
                module.definitions.push(def);
            }
            Rule::use_statement => {
                let path = parse_import(pair.into_inner().next().unwrap());
                module.imports.push(path);
            }
            _ => panic!("unexpected rule in source_file: {:?}", pair.as_rule()),
        }
    }

    module
}

fn parse_block(pair: Pair<Rule>) -> Block {
    let mut statements = vec![];
    let mut value = None;

    for pair in pair.into_inner() {
        match pair.as_rule() {
            Rule::statement => statements.push(parse_statement(pair)),
            Rule::expr => value = Some(parse_expression(pair)),
            _ => panic!("unexpected rule {:?} in block", pair.as_rule()),
        }
    }

    Block { statements, value }
}

fn parse_statement(pair: Pair<Rule>) -> Statement {
    let pair = pair.into_inner().next().unwrap();
    match pair.as_rule() {
        Rule::assign_statement => {
            let mut pairs = pair.into_inner();
            let identifier = pairs.next().unwrap().as_str().to_string();
            let expr = parse_expression(pairs.next().unwrap());
            Statement::AssignStatement(identifier, expr)
        }
        Rule::declare_statement => {
            let mut pairs = pair.into_inner();
            let identifier = pairs.next().unwrap().as_str().to_string();
            let expr = parse_expression(pairs.next().unwrap());
            Statement::DeclareStatement(identifier, expr)
        }
        Rule::return_statement => {
            let expr = if let Some(pair) = pair.into_inner().next() {
                Some(parse_expression(pair))
            } else {
                None
            };
            Statement::ReturnStatement(expr)
        }
        Rule::call_statement => {
            let call = parse_call(pair.into_inner().next().unwrap());
            Statement::CallStatement(call)
        }
        Rule::use_statement => {
            let import = parse_import(pair.into_inner().next().unwrap());
            Statement::UseStatement(import)
        }
        Rule::if_statement => {
            let mut pairs = pair.into_inner();
            let condition = parse_expression(pairs.next().unwrap());
            let block = parse_block(pairs.next().unwrap());
            Statement::IfStatement(condition, block)
        }
        Rule::while_statement => {
            let mut pairs = pair.into_inner();
            let condition = parse_expression(pairs.next().unwrap());
            let block = parse_block(pairs.next().unwrap());
            Statement::WhileStatement(Box::new(condition), Box::new(block))
        }
        _ => unreachable!("unexpected rule '{:?}' in parse_statement", pair.as_rule()),
    }
}

type ImportPath = ModulePath;

fn parse_import(pair: Pair<Rule>) -> Import {
    Import(pair.as_str().to_string())
}

fn parse_call(pair: Pair<Rule>) -> Call {
    let mut pairs = pair.into_inner();
    // TODO: support calls on more than identifiers (needs grammar change)
    let callee = Expr::Identifier(pairs.next().unwrap().as_str().to_string());
    let args: Vec<Expr> = pairs
        .next()
        .unwrap()
        .into_inner()
        .map(parse_expression)
        .collect();
    Call {
        callee: Box::new(callee),
        args,
    }
}

fn parse_expression(pair: Pair<Rule>) -> Expr {
    let pairs = pair.into_inner();
    PRATT_PARSER
        .map_primary(|primary| match primary.as_rule() {
            Rule::integer_literal => Expr::IntegerLiteral(primary.as_str().parse().unwrap()),
            Rule::float_literal => Expr::FloatLiteral(primary.as_str().parse().unwrap()),
            Rule::string_literal => Expr::StringLiteral(
                primary
                    .into_inner()
                    .next()
                    .unwrap()
                    .as_str()
                    .parse()
                    .unwrap(),
            ),
            Rule::expr => parse_expression(primary),
            Rule::ident => Expr::Identifier(primary.as_str().to_string()),
            Rule::call => Expr::Call(Box::new(parse_call(primary))),
            Rule::block => Expr::Block(Box::new(parse_block(primary))),
            Rule::if_expr => {
                let mut pairs = primary.into_inner();
                let condition = parse_expression(pairs.next().unwrap());
                let true_block = parse_block(pairs.next().unwrap());
                let else_value = parse_expression(pairs.next().unwrap());
                Expr::IfExpr(
                    Box::new(condition),
                    Box::new(true_block),
                    Box::new(else_value),
                )
            }
            Rule::boolean_literal => Expr::BooleanLiteral(match primary.as_str() {
                "true" => true,
                "false" => false,
                _ => unreachable!(),
            }),
            _ => unreachable!(
                "Unexpected rule '{:?}' in primary expression",
                primary.as_rule()
            ),
        })
        .map_infix(|lhs, op, rhs| {
            let operator = match op.as_rule() {
                Rule::add => BinaryOperator::Add,
                Rule::subtract => BinaryOperator::Sub,
                Rule::multiply => BinaryOperator::Mul,
                Rule::divide => BinaryOperator::Div,
                Rule::modulo => BinaryOperator::Modulo,
                Rule::equal => BinaryOperator::Equal,
                Rule::not_equal => BinaryOperator::NotEqual,
                _ => unreachable!(),
            };
            Expr::BinaryExpression(Box::new(lhs), operator, Box::new(rhs))
        })
        .parse(pairs)
}

fn parse_parameter(pair: Pair<Rule>) -> Parameter {
    assert!(pair.as_rule() == Rule::parameter);
    let mut pair = pair.into_inner();
    let name = pair.next().unwrap().as_str().to_string();
    let typ = Type::from(pair.next().unwrap().as_str());
    Parameter { name, typ }
}

fn parse_definition(pair: Pair<Rule>) -> Definition {
    match pair.as_rule() {
        Rule::func_def => {
            let line_col = pair.line_col();
            let mut pairs = pair.into_inner();
            let name = pairs.next().unwrap().as_str().to_string();
            let parameters: Vec<Parameter> = pairs
                .next()
                .unwrap()
                .into_inner()
                .map(parse_parameter)
                .collect();
            let pair = pairs.next().unwrap();
            // Before the block there is an optional return type
            let (return_type, pair) = match pair.as_rule() {
                Rule::ident => (Some(Type::from(pair.as_str())), pairs.next().unwrap()),
                Rule::block => (None, pair),
                _ => unreachable!(
                    "Unexpected rule '{:?}' in function definition, expected return type or block",
                    pair.as_rule()
                ),
            };
            let body = parse_block(pair);
            let body = Box::new(body);
            Definition::FunctionDefinition(FunctionDefinition {
                name,
                parameters,
                return_type,
                body,
                line_col,
            })
        }
        _ => panic!("unexpected node for definition: {:?}", pair.as_rule()),
    }
}