use multipeek::multipeek;

use crate::*;
use super::lex::*;
use super::ast::*;
use super::ast::Binding::*;
use super::ast::Control::*;
use Token::*;
use Literal::*;
use Punctuation::*;

struct Input(multipeek::MultiPeek<std::vec::IntoIter<Token>>);

impl Input {
    /// Map input.next() to return Results for use with the propagation operator
    fn next(&mut self) -> Result<Token> {
        self.0.next().ok_or("end of input".into())
    }

    /// Check if the next character is an expected token, and if so, advance the iterator and return true
    fn peek(&mut self, expected: Token) -> bool {
        if self.0.peek() == Some(&expected) {
            self.0.next();
            true
        } else {
            false
        }
    }

    /// Expose input.peek() (we don't want EOF to trigger an error when peeking)
    fn peek_opt(&mut self) -> Option<&Token> {
        self.0.peek()
    }

    /// Expose input.peek_nth()
    fn peek_nth(&mut self, n: usize) -> Option<&Token> {
        self.0.peek_nth(n)
    }

    /// Asserts the next character to be a known token
    fn then(&mut self, expected: Token) -> Result<()> {
        match self.next()? {
            token if expected == token => Ok(()),
            token => Err(format!("expected token {} but found {}", expected, token).into())
        }
    }

    /// Allow usage of `len`
    fn len(&self) -> usize {
        self.0.len()
    }
}

/// Convert a basic TokenStream into an AbstractSyntaxTree
pub fn astify(input: TokenStream, name: &str) -> Result<Expr> {
    let mut input = Input(multipeek(input));
    let body = parse_body(&mut input)?;
    if input.len() > 0 {
        return Err(format!("additional tokens remaining after the body!").into());
    }
    Ok(Expr::Binding(Module{ id: name.to_string(), body }))
}

/// Parse a series of Exprs, for ex. the body of a function.
/// Body ::= Expr | ('{' Expr (';' Expr)* '}')
fn parse_body(input: &mut Input) -> Result<Vec<Expr>> {
    let mut res = Vec::new();
    if !input.peek(Sep(ScopeLeftBrace)) {
        res.push(parse_expr(input)?);
        return Ok(res);
    }
    while !input.peek(Sep(ScopeRightBrace)) {
        res.push(parse_expr(input)?);
        // consume semicolons. there doesn't *have* to be a semicolon though.
        // this should probably be checked to be a semicolon or a right brace.
        input.peek(Sep(Semicolon));
    }
    Ok(res)
}

/// Expr ::= Let | Var | Const | Func | Type | Mod | Import |
///   Block | Static | For | While | Loop | If | When | Try | Match
fn parse_expr(input: &mut Input) -> Result<Expr> {
    use Keyword::*;
    // Note that this match consumes, as peeking is redundant.
    // This is why subsequent functions do not check for their leading keyword, i.e. 'let'
    match input.next()? {
        Key(word) => match word {
            Pub => {
                match input.next()? {
                    Key(word) => match word {
                        Const => parse_const(input, true),
                        Func => parse_func(input, true),
                        Macro => parse_macro(input, true),
                        Type => parse_typedecl(input, true),
                        Mod => parse_mod(input, true),
                        _ => return Err("unrecognized keyword following pub".into()),
                    }
                    _ => return Err("unrecognized thing following pub".into()),
                }
            },
            Let => parse_let(input),
            Var => parse_var(input),
            Const => parse_const(input, false),
            Func => parse_func(input, false),
            Macro => parse_macro(input, false),
            Type => parse_typedecl(input, false),
            Mod => parse_mod(input, false),
            Use => parse_use(input),
            Block => parse_block(input),
            Static => parse_static(input),
            For => parse_for(input),
            While => parse_while(input),
            Loop => parse_loop(input),
            If => parse_if(input),
            When => parse_when(input),
            Try => parse_try(input),
            Match => parse_match(input),
            word => return Err(format!("attempting to parse {} out of context", word).into()),
        },
        _ => todo!(), // what can i do with this?? match line here
    }
}

/// `Let ::= 'let' Pattern Annotation? '=' Expr`
fn parse_let(input: &mut Input) -> Result<Expr> {
    let id = parse_pattern(input)?;
    let kind = parse_annotation(input)?;
    input.then(Sep(Equals))?;
    let value = Box::new(parse_expr(input)?);
    Ok(Expr::Binding(Let { id, kind, value }))
}

/// `Var ::= 'var' Pattern Annotation? ('=' Expr)?`
fn parse_var(input: &mut Input) -> Result<Expr> {
    let id = parse_pattern(input)?;
    let kind = parse_annotation(input)?;
    let mut value = None;
    if input.peek(Sep(Equals)) {
        value = Some(Box::new(parse_expr(input)?));
    }
    Ok(Expr::Binding(Var { id, kind, value }))
}

/// `Const ::= 'pub'? 'const' Pattern Annotation? '=' Expr`
fn parse_const(input: &mut Input, public: bool) -> Result<Expr> {
    let id = parse_pattern(input)?;
    let kind = parse_annotation(input)?;
    input.then(Sep(Equals))?;
    let value = Box::new(parse_expr(input)?);
    Ok(Expr::Binding(Const { public, id, kind, value }))
}

/// Annotation ::= (':' Type)?
fn parse_annotation(input: &mut Input) -> Result<Option<Type>> {
    if input.peek(Sep(Colon)) {
        Ok(Some(parse_type(input)?))
    } else {
        Ok(None)
    }
}

/// `Func ::= 'pub'? 'func' Ident ('[' Parameters ']')? ('(' Parameters ')')? Annotation '=' Body`
fn parse_func(input: &mut Input, public: bool) -> Result<Expr> {
    let effect = None;
    let id = parse_ident(input)?;
    let mut generics = Vec::new();
    if input.peek(Sep(GenericLeftBracket)) {
        generics = parse_parameters(input)?;
        input.then(Sep(GenericRightBracket))?;
    }
    let mut parameters = Vec::new();
    if input.peek(Sep(FuncLeftParen)) { // todo: rewrite to map over an input
        // let temp_parameters = parse_parameters(input)?;
        // if temp_parameters.last().is_none() {
        //     return Err("expected a type annotation on the last function parameter".into());
        // }
        // parameters = parse_parameters(input)?.iter().map(|x| x).collect();
        let mut stack = Vec::new();
        let (id, kind) = parse_parameter(input)?;
        if kind.is_some() {
            parameters.push((id, kind.unwrap()));
        } else {
            stack.push(id);
        }
        while input.peek(Sep(Comma)) {
            let (id, kind) = parse_parameter(input)?;
            stack.push(id);
            if kind.is_some() {
                for id in &stack {
                    parameters.push((id.clone(), kind.clone().unwrap()));
                }
                stack.clear();
            }
        }
        if stack.len() != 0 {
            return Err("expected a type annotation on the last function parameter".into());
        }
        input.then(Sep(FuncRightParen))?;
    }
    let mut kind = Type::Void;
    if input.peek(Sep(Colon)) {
        kind = parse_type(input)?;
    }
    input.then(Sep(Equals))?;
    let body = parse_body(input)?;
    Ok(Expr::Binding(Func { public, effect, id, generics, parameters, kind, body }))
}

/// `Macro ::= 'pub'? 'macro' Ident ('[' Parameters ']')? ('(' Parameters ')')? (':' Type) '=' Body`
fn parse_macro(input: &mut Input, public: bool) -> Result<Expr> {
    let id = parse_ident(input)?;
    let mut generics = Vec::new();
    if input.peek(Sep(GenericLeftBracket)) {
        generics = parse_parameters(input)?;
        input.then(Sep(GenericRightBracket))?;
    }
    let mut parameters = Vec::new();
    if input.peek(Sep(FuncLeftParen)) {
        parameters = parse_parameters(input)?;
        input.then(Sep(FuncRightParen))?;
    }
    let mut kind = None;
    if input.peek(Sep(Colon)) {
        kind = Some(parse_type(input)?);
    }
    input.then(Sep(Equals))?;
    let body = parse_body(input)?;
    Ok(Expr::Binding(Macro { public, id, generics, parameters, kind, body }))
}

/// `TypeDecl ::= 'pub'? 'type' Ident ('[' Parameters ']')? '=' Type
fn parse_typedecl(input: &mut Input, public: bool) -> Result<Expr> {
    let id = parse_ident(input)?;
    let mut generics = Vec::new();
    if input.peek(Sep(GenericLeftBracket)) {
        generics = parse_parameters(input)?;
        input.then(Sep(GenericRightBracket))?;
    }
    input.then(Sep(Equals))?;
    let alias = parse_type(input)?;
    Ok(Expr::Binding(TypeDecl { id, generics, alias }))
}

/// `Parameter ::= Ident (':' Type)?`
fn parse_parameter(input: &mut Input) -> Result<(Id, Option<Type>)> {
    let id = parse_ident(input)?;
    let mut kind = None;
    if input.peek(Sep(Colon)) {
        kind = Some(parse_type(input)?);
    }
    Ok((id, kind))
}

/// `Parameters ::= Parameter (',' Parameter)*
fn parse_parameters(input: &mut Input) -> Result<Vec<(Id, Option<Type>)>> {
    let mut res = Vec::new();
    res.push(parse_parameter(input)?);
    while input.peek(Sep(Comma)) {
        res.push(parse_parameter(input)?);
    }
    Ok(res)
}

/// `Mod ::= 'pub'? 'mod' Ident ':' Body`
fn parse_mod(input: &mut Input, public: bool) -> Result<Expr> {
    let id = parse_ident(input)?;
    input.then(Sep(Colon))?;
    let body = parse_body(input)?;
    Ok(Expr::Binding(Module { id, body }))
}

/// `Use ::= 'use' Ident ('/' Ident)* ('/' (('[' Ident (',' Ident)* ']') | '*'))?`
fn parse_use(input: &mut Input) -> Result<Expr> { todo!() }

/// `Block ::= 'block' Ident? ':' Body`
fn parse_block(input: &mut Input) -> Result<Expr> {
    match input.next()? {
        Sep(Colon) => {
            let id = None;
            let body = parse_body(input)?;
            Ok(Expr::Control(Block { id, body }))
        },
        Word(label) => {
            input.then(Sep(Colon))?;
            let id = Some(label);
            let body = parse_body(input)?;
            Ok(Expr::Control(Block { id, body }))
        },
        _ => return Err("unexpected thing following block keyword".into()),
    }
}

/// `Static ::= 'static' ':' Body`
fn parse_static(input: &mut Input) -> Result<Expr> {
    input.then(Sep(Colon))?;
    let body = parse_body(input)?;
    Ok(Expr::Control(Static { body }))
}

/// `For ::= 'for' Pattern 'in' Expr ':' Body`
fn parse_for(input: &mut Input) -> Result<Expr> {
    let binding = parse_pattern(input)?;
    input.then(Key(Keyword::In))?;
    let range = Box::new(parse_expr(input)?);
    input.then(Sep(Colon))?;
    let body = parse_body(input)?;
    Ok(Expr::Control(For { binding, range, body }))
}

/// `While ::= 'while' Expr ':' Body`
fn parse_while(input: &mut Input) -> Result<Expr> {
    let cond = Box::new(parse_expr(input)?);
    input.then(Sep(Colon))?;
    let body = parse_body(input)?;
    Ok(Expr::Control(While { cond, body }))
}

/// `Loop ::= 'loop' ':' Body`
fn parse_loop(input: &mut Input) -> Result<Expr> {
    input.then(Sep(Colon))?;
    let body = parse_body(input)?;
    Ok(Expr::Control(Loop { body }))
}

/// `If ::= 'if' CondBranch ('elif' CondBranch)* ('else' ':' Body)?`
fn parse_if(input: &mut Input) -> Result<Expr> {
    let mut branches = Vec::new();
    branches.push(parse_cond_branch(input)?);
    while input.peek(Key(Keyword::Elif)) {
        branches.push(parse_cond_branch(input)?);
    }
    let mut else_body = None;
    if input.peek(Key(Keyword::Else)) {
        else_body = Some(parse_body(input)?);
    }
    Ok(Expr::Control(If { branches, else_body }))
}

/// `When ::= 'when' CondBranch ('elif' CondBranch)* ('else' ':' Body)?`
fn parse_when(input: &mut Input) -> Result<Expr> {
    let mut branches = Vec::new();
    branches.push(parse_cond_branch(input)?);
    while input.peek(Key(Keyword::Elif)) {
        branches.push(parse_cond_branch(input)?);
    }
    let mut else_body = None;
    if input.peek(Key(Keyword::Else)) {
        input.then(Sep(Colon))?;
        else_body = Some(parse_body(input)?);
    }
    let mut body = Vec::new();
    body.push(Expr::Control(If { branches, else_body }));
    Ok(Expr::Control(Static { body }))
}

/// `CondBranch ::= Expr ':' Body`
fn parse_cond_branch(input: &mut Input) -> Result<CondBranch> {
    let cond = parse_expr(input)?;
    input.then(Sep(Colon))?;
    let body = parse_body(input)?;
    Ok(CondBranch { cond, body })
}

/// `Try ::= 'try' ':' Body ('except' Exception (',' Exception)* ':' Body)* ('finally' ':' Body)?`
fn parse_try(input: &mut Input) -> Result<Expr> {
    input.then(Sep(Colon))?;
    let body = parse_body(input)?;
    let mut catches = Vec::new();
    while input.peek(Key(Keyword::Catch)) {
        let mut exceptions = Vec::new();
        exceptions.push(parse_catch_exception(input)?);
        while input.peek(Sep(Comma)) {
            exceptions.push(parse_catch_exception(input)?);
        }
        input.then(Sep(Colon))?;
        let body = parse_body(input)?;
        catches.push(CatchBranch { exceptions, body });
    }
    let mut finally = None;
    if input.peek(Key(Keyword::Finally)) {
        input.then(Sep(Colon))?;
        finally = Some(parse_body(input)?);
    }
    Ok(Expr::Control(Try { body, catches, finally }))
}

/// `Exception ::= Ident ('as' Ident)?`
fn parse_catch_exception(input: &mut Input) -> Result<(Id, Option<Id>)> {
    let id = parse_ident(input)?;
    let mut alias = None;
    if input.peek(Key(Keyword::As)) {
        alias = Some(parse_ident(input)?);
    }
    Ok((id, alias))
}

/// `Match ::= 'match' Expr ('of' Pattern (',' Pattern)* ('where' Expr)? ':' Body)+`
fn parse_match(input: &mut Input) -> Result<Expr> {
    let item = parse_pattern(input)?; // fixme
    let mut branches = Vec::new();
    while input.peek(Key(Keyword::Of)) {
        let mut patterns = Vec::new();
        patterns.push(parse_pattern(input)?);
        while input.peek(Sep(Comma)) {
            patterns.push(parse_pattern(input)?);
        }
        let mut guard = None;
        if input.peek(Key(Keyword::Where)) {
            guard = Some(parse_expr(input)?)
        }
        input.then(Sep(Colon))?;
        let body = parse_body(input)?;
        branches.push(MatchBranch { patterns, guard, body })
    }
    Ok(Expr::Control(Match { item, branches }))
}

/// `Type ::= (('distinct' | 'ref' | 'ptr' | 'mut' | 'static') BracketType?) |
///   StructType | TupleType | EnumType | UnionType | InterfaceType`
///
/// The input stream must be normalized (i.e. all type declarations must be on one line)
/// before attempting to parse types, because otherwise it's just a little bit hellish.
fn parse_type(input: &mut Input) -> Result<Type> {
    use Type::*;
    match input.next()? {
        Key(word) => match word {
            Keyword::Distinct => Ok(Distinct(Box::new(parse_wrapped_type(input)?))),
            Keyword::Ref => Ok(Reference(Box::new(parse_wrapped_type(input)?))),
            Keyword::Ptr => Ok(Pointer(Box::new(parse_wrapped_type(input)?))),
            Keyword::Mut => Ok(Mutable(Box::new(parse_wrapped_type(input)?))),
            Keyword::Static => Ok(Static(Box::new(parse_wrapped_type(input)?))),
            Keyword::Struct => parse_struct_type(input),
            Keyword::Tuple => parse_tuple_type(input),
            Keyword::Enum => parse_enum_type(input),
            Keyword::Union => parse_union_type(input),
            Keyword::Interface => parse_interface(input),
            _ => return Err("invalid keyword present in type!".into())
        },
        Word(id) => {
            let mut generics = Vec::new();
            if input.peek(Sep(GenericLeftBracket)) {
                generics.push(parse_type(input)?);
                while input.peek(Sep(Comma)) {
                    generics.push(parse_type(input)?);
                }
                input.then(Sep(GenericRightBracket))?;
            }
            Ok(Alias { id, generics })
        },
        _ => return Err("error".into())
    }
}

/// `StructType ::= 'struct' ('[' Ident ':' Type (',' Ident ':' Type)* ']')?`
fn parse_struct_type(input: &mut Input) -> Result<Type> {
    let mut res = Vec::new();
    if input.peek(Sep(GenericLeftBracket)) {
        res.push(parse_struct_field(input)?);
        while input.peek(Sep(Comma)) {
            res.push(parse_struct_field(input)?);
        }
        input.then(Sep(GenericRightBracket))?;
    }
    Ok(Type::Struct(res))
}

fn parse_struct_field(input: &mut Input) -> Result<(Id, Box<Type>)> {
    let id = parse_ident(input)?;
    input.then(Sep(Colon))?;
    let kind = Box::new(parse_type(input)?);
    Ok((id, kind))
}

/// `TupleType ::= 'tuple' ('[' (Ident ':')? Type (',' (Ident ':')? Type)* ']')?`
fn parse_tuple_type(input: &mut Input) -> Result<Type> {
    let mut res = Vec::new();
    if input.peek(Sep(GenericLeftBracket)) {
        res.push(parse_tuple_field(input)?);
        while input.peek(Sep(Comma)) {
            res.push(parse_tuple_field(input)?);
        }
        input.then(Sep(GenericRightBracket))?;
    }
    Ok(Type::Tuple(res))
}

// annoyingly complex to parse. `TupleField ::= (Ident ':')? Type`
fn parse_tuple_field(input: &mut Input) -> Result<(Option<Id>, Box<Type>)> {
    match input.peek_opt().clone() {
        Some(Word(id)) if input.peek_nth(1) == Some(&Sep(Colon)) => {
            input.next()?;
            input.then(Sep(Colon))?;
            Ok((Some(id.to_string()), Box::new(parse_type(input)?)))
        },
        _ => Ok((None, Box::new(parse_type(input)?)))
    }
}

/// `EnumType ::= 'enum' ('[' Ident ('=' Pattern)? (Ident ('=' Pattern)?)* ']')?`
fn parse_enum_type(input: &mut Input) -> Result<Type> {
    let mut res = Vec::new();
    if input.peek(Sep(GenericLeftBracket)) {
        res.push(parse_enum_variant(input)?);
        while input.peek(Sep(Comma)) {
            res.push(parse_enum_variant(input)?);
        }
        input.then(Sep(GenericRightBracket))?;
    }
    todo!()
}

fn parse_enum_variant(input: &mut Input) -> Result<(Id, Option<Pattern>)> {
    let id = parse_ident(input)?;
    let mut kind = None;
    if input.peek(Sep(Equals)) {
        kind = Some(parse_pattern(input)?);
    }
    Ok((id, kind))
}

/// `UnionType ::= 'union' ('[' Ident (':' Type)? (',' Ident (':' Type)?)* ']')?`
fn parse_union_type(input: &mut Input) -> Result<Type> {
    let mut res = Vec::new();
    if input.peek(Sep(GenericLeftBracket)) {
        res.push(parse_union_variant(input)?);
        while input.peek(Sep(Comma)) {
            res.push(parse_union_variant(input)?);
        }
        input.then(Sep(GenericRightBracket))?;
    }
    Ok(Type::Union(res))
}

fn parse_union_variant(input: &mut Input) -> Result<(Id, Box<Type>)> {
    let id = parse_ident(input)?;
    let mut kind = Box::new(Type::Alias { id: "unit".to_string(), generics: Vec::new() });
    if input.peek(Sep(Colon)) {
        kind = Box::new(parse_type(input)?);
    }
    Ok((id, kind))
}

/// `Interface ::= 'interface' ('[' Signature (',' Signature)* ']')?`
fn parse_interface(input: &mut Input) -> Result<Type> {
    let mut res = Vec::new();
    if input.peek(Sep(GenericLeftBracket)) {
        res.push(parse_signature(input)?);
        while input.peek(Sep(Comma)) {
            res.push(parse_signature(input)?);
        }
        input.then(Sep(GenericRightBracket))?;
    }
    Ok(Type::Interface(res))
}

/// `Signature ::= Ident ('[' Parameters ']')? ('(' Type (',' Type)* ')')? (':' Type)?`
fn parse_signature(input: &mut Input) -> Result<Sig> {
    let effect = None;
    let id = parse_ident(input)?;
    let mut generics = Vec::new();
    if input.peek(Sep(GenericLeftBracket)) {
        generics = parse_parameters(input)?;
        input.then(Sep(GenericRightBracket))?;
    }
    let mut parameters = Vec::new();
    if input.peek(Sep(FuncLeftParen)) {
        parameters.push(parse_type(input)?);
        if input.peek(Sep(Comma)) {
            parameters.push(parse_type(input)?);
        }
        input.then(Sep(FuncRightParen))?;
    }
    let mut kind = None;
    if input.peek(Sep(Colon)) {
        kind = Some(parse_type(input)?);
    }
    Ok(Sig { effect, id, generics, parameters, kind })
}

/// `WrappedType ::= Type | ('[' Type ']')`
fn parse_wrapped_type(input: &mut Input) -> Result<Type> {
    if input.peek(Sep(GenericLeftBracket)) {
        let result = parse_type(input)?;
        input.then(Sep(GenericRightBracket))?;
        Ok(result)
    } else {
        parse_type(input)
    }
}

/// Pattern ::= Literal | Ident | '(' Pattern (',' Pattern)* ')' | Ident '(' Pattern (',' Pattern)* ')'
fn parse_pattern(input: &mut Input) -> Result<Pattern> { todo!() }

fn parse_ident(input: &mut Input) -> Result<Id> {
    match input.next()? {
        Word(id) => Ok(id),
        token => Err(format!("expected identifier but found token {}", token).into())
    }
}