use crate::ast::{self, kw};
use crate::parser::{Cursor, Parse, Parser, Peek, Result};

/// A parsed representation of a `*.wast` file.
///
/// WAST files are not officially specified but are used in the official test
/// suite to write official spec tests for wasm. This type represents a parsed
/// `*.wast` file which parses a list of directives in a file.
pub struct Wast<'a> {
    #[allow(missing_docs)]
    pub directives: Vec<WastDirective<'a>>,
}

impl<'a> Parse<'a> for Wast<'a> {
    fn parse(parser: Parser<'a>) -> Result<Self> {
        let mut directives = Vec::new();

        // If it looks like a directive token is in the stream then we parse a
        // bunch of directives, otherwise assume this is an inline module.
        if parser.peek2::<WastDirectiveToken>() {
            while !parser.is_empty() {
                directives.push(parser.parens(|p| p.parse())?);
            }
        } else {
            let module = parser.parse::<ast::Wat>()?.module;
            directives.push(WastDirective::Module(module));
        }
        Ok(Wast { directives })
    }
}

struct WastDirectiveToken;

impl Peek for WastDirectiveToken {
    fn peek(cursor: Cursor<'_>) -> bool {
        let kw = match cursor.keyword() {
            Some((kw, _)) => kw,
            None => return false,
        };
        kw.starts_with("assert_") || kw == "module" || kw == "register" || kw == "invoke"
    }

    fn display() -> &'static str {
        unimplemented!()
    }
}

/// The different kinds of directives found in a `*.wast` file.
///
/// It's not entirely clear to me what all of these are per se, but they're only
/// really interesting to test harnesses mostly.
#[allow(missing_docs)]
pub enum WastDirective<'a> {
    Module(ast::Module<'a>),
    AssertMalformed {
        span: ast::Span,
        module: QuoteModule<'a>,
        message: &'a str,
    },
    AssertInvalid {
        span: ast::Span,
        module: ast::Module<'a>,
        message: &'a str,
    },
    Register {
        span: ast::Span,
        name: &'a str,
        module: Option<ast::Id<'a>>,
    },
    Invoke(WastInvoke<'a>),
    AssertTrap {
        span: ast::Span,
        exec: WastExecute<'a>,
        message: &'a str,
    },
    AssertReturn {
        span: ast::Span,
        exec: WastExecute<'a>,
        results: Vec<ast::Expression<'a>>,
    },
    AssertReturnCanonicalNan {
        span: ast::Span,
        invoke: WastInvoke<'a>,
    },
    AssertReturnArithmeticNan {
        span: ast::Span,
        invoke: WastInvoke<'a>,
    },
    AssertReturnFunc {
        span: ast::Span,
        invoke: WastInvoke<'a>,
    },
    AssertExhaustion {
        span: ast::Span,
        call: WastInvoke<'a>,
        message: &'a str,
    },
    AssertUnlinkable {
        span: ast::Span,
        module: ast::Module<'a>,
        message: &'a str,
    },
}

impl WastDirective<'_> {
    /// Returns the location in the source that this directive was defined at
    pub fn span(&self) -> ast::Span {
        match self {
            WastDirective::Module(m) => m.span,
            WastDirective::AssertMalformed { span, .. } |
            WastDirective::Register { span, .. } |
            WastDirective::AssertTrap { span, .. } |
            WastDirective::AssertReturn { span, .. } |
            WastDirective::AssertReturnCanonicalNan { span, .. } |
            WastDirective::AssertReturnArithmeticNan { span, .. } |
            WastDirective::AssertReturnFunc { span, .. } |
            WastDirective::AssertExhaustion { span, .. } |
            WastDirective::AssertUnlinkable { span, .. } |
            WastDirective::AssertInvalid { span, .. } => *span,
            WastDirective::Invoke(i) => i.span,
        }
    }
}

impl<'a> Parse<'a> for WastDirective<'a> {
    fn parse(parser: Parser<'a>) -> Result<Self> {
        let mut l = parser.lookahead1();
        if l.peek::<kw::module>() {
            Ok(WastDirective::Module(parser.parse()?))
        } else if l.peek::<kw::assert_malformed>() {
            let span = parser.parse::<kw::assert_malformed>()?.0;
            Ok(WastDirective::AssertMalformed {
                span,
                module: parser.parens(|p| p.parse())?,
                message: parser.parse()?,
            })
        } else if l.peek::<kw::assert_invalid>() {
            let span = parser.parse::<kw::assert_invalid>()?.0;
            Ok(WastDirective::AssertInvalid {
                span,
                module: parser.parens(|p| p.parse())?,
                message: parser.parse()?,
            })
        } else if l.peek::<kw::register>() {
            let span = parser.parse::<kw::register>()?.0;
            Ok(WastDirective::Register {
                span,
                name: parser.parse()?,
                module: parser.parse()?,
            })
        } else if l.peek::<kw::invoke>() {
            Ok(WastDirective::Invoke(parser.parse()?))
        } else if l.peek::<kw::assert_trap>() {
            let span = parser.parse::<kw::assert_trap>()?.0;
            Ok(WastDirective::AssertTrap {
                span,
                exec: parser.parens(|p| p.parse())?,
                message: parser.parse()?,
            })
        } else if l.peek::<kw::assert_return>() {
            let span = parser.parse::<kw::assert_return>()?.0;
            let exec = parser.parens(|p| p.parse())?;
            let mut results = Vec::new();
            while !parser.is_empty() {
                results.push(parser.parens(|p| p.parse())?);
            }
            Ok(WastDirective::AssertReturn {
                span,
                exec,
                results,
            })
        } else if l.peek::<kw::assert_return_canonical_nan>() {
            let span = parser.parse::<kw::assert_return_canonical_nan>()?.0;
            Ok(WastDirective::AssertReturnCanonicalNan {
                span,
                invoke: parser.parens(|p| p.parse())?,
            })
        } else if l.peek::<kw::assert_return_arithmetic_nan>() {
            let span = parser.parse::<kw::assert_return_arithmetic_nan>()?.0;
            Ok(WastDirective::AssertReturnArithmeticNan {
                span,
                invoke: parser.parens(|p| p.parse())?,
            })
        } else if l.peek::<kw::assert_return_func>() {
            let span = parser.parse::<kw::assert_return_func>()?.0;
            Ok(WastDirective::AssertReturnFunc {
                span,
                invoke: parser.parens(|p| p.parse())?,
            })
        } else if l.peek::<kw::assert_exhaustion>() {
            let span = parser.parse::<kw::assert_exhaustion>()?.0;
            Ok(WastDirective::AssertExhaustion {
                span,
                call: parser.parens(|p| p.parse())?,
                message: parser.parse()?,
            })
        } else if l.peek::<kw::assert_unlinkable>() {
            let span = parser.parse::<kw::assert_unlinkable>()?.0;
            Ok(WastDirective::AssertUnlinkable {
                span,
                module: parser.parens(|p| p.parse())?,
                message: parser.parse()?,
            })
        } else {
            Err(l.error())
        }
    }
}

#[allow(missing_docs)]
pub enum WastExecute<'a> {
    Invoke(WastInvoke<'a>),
    Module(ast::Module<'a>),
    Get {
        module: Option<ast::Id<'a>>,
        global: &'a str,
    },
}

impl<'a> Parse<'a> for WastExecute<'a> {
    fn parse(parser: Parser<'a>) -> Result<Self> {
        let mut l = parser.lookahead1();
        if l.peek::<kw::invoke>() {
            Ok(WastExecute::Invoke(parser.parse()?))
        } else if l.peek::<kw::module>() {
            Ok(WastExecute::Module(parser.parse()?))
        } else if l.peek::<kw::get>() {
            parser.parse::<kw::get>()?;
            Ok(WastExecute::Get {
                module: parser.parse()?,
                global: parser.parse()?,
            })
        } else {
            Err(l.error())
        }
    }
}

#[allow(missing_docs)]
pub struct WastInvoke<'a> {
    pub span: ast::Span,
    pub module: Option<ast::Id<'a>>,
    pub name: &'a str,
    pub args: Vec<ast::Expression<'a>>,
}

impl<'a> Parse<'a> for WastInvoke<'a> {
    fn parse(parser: Parser<'a>) -> Result<Self> {
        let span = parser.parse::<kw::invoke>()?.0;
        let module = parser.parse()?;
        let name = parser.parse()?;
        let mut args = Vec::new();
        while !parser.is_empty() {
            args.push(parser.parens(|p| p.parse())?);
        }
        Ok(WastInvoke {
            span,
            module,
            name,
            args,
        })
    }
}

#[allow(missing_docs)]
pub enum QuoteModule<'a> {
    Module(ast::Module<'a>),
    Quote(Vec<&'a str>),
}

impl<'a> Parse<'a> for QuoteModule<'a> {
    fn parse(parser: Parser<'a>) -> Result<Self> {
        if parser.peek2::<kw::quote>() {
            parser.parse::<kw::module>()?;
            parser.parse::<kw::quote>()?;
            let mut src = Vec::new();
            while !parser.is_empty() {
                src.push(parser.parse()?);
            }
            Ok(QuoteModule::Quote(src))
        } else {
            Ok(QuoteModule::Module(parser.parse()?))
        }
    }
}