// Copyright 2018 Kyle Mayes
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Macro implementations for `tarrasque`.

#![recursion_limit="128"]

#[macro_use]
extern crate quote;

extern crate proc_macro;
extern crate proc_macro2;
extern crate syn;

use proc_macro2::{TokenStream};
use syn::*;
use syn::parse::{Parse, ParseStream, Result};
use syn::punctuated::{Punctuated};

/// An extract struct parameter (e.g., `size: i32`).
struct ExtractParam {
    ident: Ident,
    ty: Type,
}

impl Parse for ExtractParam {
    fn parse(input: ParseStream) -> Result<Self> {
        let ident = input.parse()?;
        input.parse::<Token![:]>()?;
        let ty = input.parse()?;
        Ok(ExtractParam { ident, ty })
    }
}

/// An extract struct field expression (e.g., `[]`, `[size]`, or `foo + bar`).
enum ExtractExpr {
    Expr(Expr),
    Extract(Option<Expr>),
}

impl Parse for ExtractExpr {
    fn parse(input: ParseStream) -> Result<Self> {
        let lookahead = input.lookahead1();
        if lookahead.peek(token::Bracket) {
            let content;
            bracketed!(content in input);
            if !content.is_empty() {
                Ok(ExtractExpr::Extract(Some(content.parse()?)))
            } else {
                Ok(ExtractExpr::Extract(None))
            }
        } else {
            Ok(ExtractExpr::Expr(input.parse()?))
        }
    }
}

/// An extract struct field (e.g., `pub value: u64 = [size]`).
struct ExtractField {
    attrs: Vec<Attribute>,
    vis: Visibility,
    ident: Ident,
    ty: Type,
    expr: ExtractExpr,
}

impl Parse for ExtractField {
    fn parse(input: ParseStream) -> Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        let vis = input.parse()?;
        let ident = input.parse()?;
        input.parse::<Token![:]>()?;
        let ty = input.parse()?;
        input.parse::<Token![=]>()?;
        let expr = input.parse()?;
        Ok(ExtractField { attrs, vis, ident, ty, expr })
    }
}

/// An extract struct.
struct ExtractStruct {
    attrs: Vec<Attribute>,
    vis: Visibility,
    ident: Ident,
    generics: Generics,
    size: Option<Expr>,
    params: Option<Punctuated<ExtractParam, Token![,]>>,
    where_: Option<WhereClause>,
    fields: Punctuated<ExtractField, Token![,]>,
}

impl Parse for ExtractStruct {
    fn parse(input: ParseStream) -> Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        let vis = input.parse()?;
        input.parse::<Token![struct]>()?;
        let ident = input.parse()?;
        let generics = input.parse()?;

        // Parse size if provided.
        let lookahead = input.lookahead1();
        let size = if lookahead.peek(token::Bracket) {
            let content;
            bracketed!(content in input);
            Some(content.parse()?)
        } else {
            None
        };

        // Parse parameters if provided.
        let lookahead = input.lookahead1();
        let params = if lookahead.peek(token::Paren) {
            let content;
            parenthesized!(content in input);
            Some(content.parse_terminated(ExtractParam::parse)?)
        } else {
            None
        };

        // Parse where clause if provided.
        let lookahead = input.lookahead1();
        let where_ = if lookahead.peek(Token![where]) {
            Some(input.parse()?)
        } else {
            None
        };

        let content;
        braced!(content in input);
        let fields = content.parse_terminated(ExtractField::parse)?;

        Ok(ExtractStruct {
            attrs, vis, ident, generics, where_, size, params, fields
        })
    }
}

fn generate_field(field: &ExtractField) -> TokenStream {
    let attrs = &field.attrs;
    let vis = &field.vis;
    let ident = &field.ident;
    let ty = &field.ty;
    quote!(#(#attrs)* #vis #ident: #ty)
}

fn generate_struct(struct_: &ExtractStruct) -> TokenStream {
    let attrs = &struct_.attrs;
    let vis = &struct_.vis;
    let ident = &struct_.ident;
    let generics = &struct_.generics;
    let where_ = &struct_.where_;
    let fields = struct_.fields.iter().map(generate_field);
    quote!(#(#attrs)* #vis struct #ident #generics #where_ { #(#fields),* })
}

fn generate_params(struct_: &ExtractStruct) -> (TokenStream, TokenStream) {
    if let Some(params) = &struct_.params {
        let idents = params.iter().map(|a| &a.ident);
        let tys = params.iter().map(|a| &a.ty);
        (quote!((#(#idents),*)), quote!((#(#tys),*)))
    } else {
        (quote!(_), quote!(()))
    }
}

fn generate_expr(field: &ExtractField) -> TokenStream {
    let ident = &field.ident;
    let ty = &field.ty;
    match &field.expr {
        ExtractExpr::Expr(expr) => quote!(let #ident: #ty = #expr;),
        ExtractExpr::Extract(expr) => if let Some(expr) = expr {
            quote!(let #ident: #ty = stream.extract(#expr)?;)
        } else {
            quote!(let #ident: #ty = stream.extract(())?;)
        }
    }
}

fn generate_impl(struct_: &ExtractStruct) -> TokenStream {
    let params = &struct_.generics.params;
    let clause = &struct_.generics.where_clause;
    let prefix = quote!(<'s, #params> #clause);

    let generics = &struct_.generics;
    let where_ = &struct_.where_;
    let suffix = quote!(#generics #where_);

    let (pat, ty) = generate_params(&struct_);
    let ident = &struct_.ident;
    let extracts = struct_.fields.iter().map(generate_expr);
    let fields = struct_.fields.iter().map(|f| &f.ident);
    let mut tokens = quote!(
        impl #prefix ::tarrasque::Extract<'s, #ty> for #ident #suffix {
            #[inline]
            fn extract(
                mut stream: &mut ::tarrasque::Stream<'s>, #pat: #ty
            ) -> ::tarrasque::ExtractResult<'s, Self> {
                #(#extracts)*
                Ok(#ident { #(#fields),* })
            }
        }
    );

    if let Some(size) = &struct_.size {
        tokens.extend(quote!(
            impl #prefix ::tarrasque::Span for #ident #suffix {
                const SPAN: usize = #size;
            }
        ));
    }

    tokens
}

/// Generates a struct and an implementation of `Extract` (and optionally
/// `Span`) for that struct.
///
/// # Example
///
/// The following code is an example of the usage of the `extract!` macro:
///
/// ```skip
/// use tarrasque::{Endianness, extract};
///
/// extract! {
///     /// A 2D point.
///     #[derive(Copy, Clone, Debug, PartialEq, Eq)]
///     pub struct Point[4](endianness: Endianness) {
///         /// The x-coordinate of this point.
///         pub x: u16 = [endianness],
///         /// The y-coordinate of this point.
///         pub y: u16 = [endianness],
///     }
/// }
/// # fn main() { }
/// ```
///
/// The `[4]` that follows `Point` is an optional component that statically
/// defines the number of bytes consumed by the generated `Extract`
/// implementation. When this component is present, a `Span` implementation is
/// generated using the specified span.
///
/// The `(endianness: Endianness)` that follows `[4]` is another optional
/// component that defines the parameter used by the `Extract` implementation.
/// In this case, a specified endianness is used to extract the `u16`s instead
/// of using the default (big-endian). When this component is omitted, a
/// placeholder parameter of type `()` is used instead.
///
/// Though not shown here, structs may also make use of generics and lifetimes.
/// See below for an example.
///
/// Each field has a corresponding expression which is used to populate it in
/// the `Extract` implementation. Square brackets indicate that the field is to
/// be populated by extracting a value of that type from the stream. The
/// expression in the square brackets is used as the argument for the `extract`
/// method. If the square brackets do not contain an expression, `()` is used as
/// the argument.
///
/// Expressions not wrapped in square brackets are treated as normal
/// expressions. The `[endianness]` expression could be replaced with the
/// equivalent `stream.extract(endianness)?`. Note the usage of `stream` which
/// is implicitly available in these expressions.
///
/// ## Expansion
///
/// The usage of `extract!` shown above generates the following code:
///
/// ```skip
/// use tarrasque::{Endianness, extract};
///
/// /// A 2D point.
/// #[derive(Copy, Clone, Debug, PartialEq, Eq)]
/// pub struct Point {
///     /// The x-coordinate of this point.
///     pub x: u16,
///     /// The y-coordinate of this point.
///     pub y: u16,
/// }
///
/// impl<'s> ::tarrasque::Extract<'s, Endianness> for Point {
///     #[inline]
///     fn extract(
///         mut stream: &mut ::tarrasque::Stream<'s>, endianness: Endianness
///     ) -> ::tarrasque::ExtractResult<'s, Self> {
///         let x = stream.extract(endianness)?;
///         let y = stream.extract(endianness)?;
///         Ok(Point { x, y })
///     }
/// }
///
/// impl ::tarrasque::Span for Point {
///     const SPAN: usize = 4;
/// }
/// ```
///
/// ## Generics and Lifetimes
///
/// The following code is an example of the usage of the `extract!` macro with
/// generics and lifetimes.
///
/// ```skip
/// use tarrasque::{Extract, Span, extract};
///
/// extract! {
///     /// A pair of values.
///     pub struct Pair<'s, T>[2 * T::SPAN] where T: Extract<'s, ()> + Span {
///         pub bytes: &'s [u8] = &stream[..],
///         pub left: T = [],
///         pub right: T = [],
///     }
/// }
/// ```
///
/// **Note:** As shown above, the lifetime `'s` must be used to refer to the
/// lifetime of the stream.
#[proc_macro]
pub fn extract(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let struct_ = parse_macro_input!(input as ExtractStruct);

    let mut tokens = TokenStream::new();
    tokens.extend(generate_struct(&struct_));
    tokens.extend(generate_impl(&struct_));

    tokens.into()
}