// 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.

//! A state machine generator for tokio codecs.

#![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<usize>,
    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::<LitInt>()?.value() as usize)
        } 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_span_impl(struct_: &ExtractStruct) -> TokenStream {
    let generics = &struct_.generics;
    let ident = &struct_.ident;
    let size = struct_.size.unwrap();
    quote!(
        impl #generics ::tarrasque::Span for #ident #generics {
            const SPAN: usize = #size;
        }
    )
}

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;
    match &field.expr {
        ExtractExpr::Expr(expr) => quote!(let #ident = #expr;),
        ExtractExpr::Extract(expr) => if let Some(expr) = expr {
            quote!(let #ident = stream.extract(#expr)?;)
        } else {
            quote!(let #ident = stream.extract(())?;)
        }
    }
}

fn generate_extract_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);
    quote!(
        impl #prefix ::tarrasque::Extract<'s, #ty> for #ident #suffix {
            fn extract(
                mut stream: &mut ::tarrasque::Stream<'s>, #pat: #ty
            ) -> ::tarrasque::ExtractResult<'s, Self> {
                #(#extracts)*
                Ok(#ident { #(#fields),* })
            }
        }
    )
}

/// Provides the implementation for `extract! { ... }`.
#[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_));
    if struct_.size.is_some() { tokens.extend(generate_span_impl(&struct_)); }
    tokens.extend(generate_extract_impl(&struct_));

    tokens.into()
}