use crate::common::{ContractSubject, GhostLet, ghost_int_lit_suffix};
use proc_macro::TokenStream as TS1;
use quote::ToTokens as _;
use syn::{Block, parse_macro_input, visit_mut::VisitMut};

// #[proc_macro]

pub fn pearlite(_: TS1) -> TS1 {
    TS1::new()
}

pub fn proof_assert(_: TS1) -> TS1 {
    quote::quote! { {} }.into()
}

pub fn snapshot(_: TS1) -> TS1 {
    quote::quote! { ::creusot_std::snapshot::Snapshot::from_fn(|| panic!() ) }.into()
}

pub fn declare_namespace(_: TS1) -> TS1 {
    TS1::new()
}

pub fn ghost(body: TS1) -> TS1 {
    let group = proc_macro2::Group::new(proc_macro2::Delimiter::Brace, body.into());
    let body = ghost_int_lit_suffix(group.into_token_stream()).into();
    let mut body = parse_macro_input!(body as Block);
    DeleteInvariants.visit_block_mut(&mut body);
    quote::quote! { if false {
        ::creusot_std::ghost::Ghost::new({ #body })
    } else {
        ::creusot_std::ghost::Ghost::conjure()
    } }
    .into()
}

pub fn ghost_let(body: TS1) -> TS1 {
    let body = ghost_int_lit_suffix(proc_macro2::TokenStream::from(body)).into();
    let GhostLet { mutability, var, mut body } = parse_macro_input!(body);
    DeleteInvariants.visit_expr_mut(&mut body);
    quote::quote! { let #mutability #var = if false {
        ::creusot_std::ghost::Ghost::new(#body)
    } else {
        ::creusot_std::ghost::Ghost::conjure()
    } }
    .into()
}

pub fn requires(_: TS1, tokens: TS1) -> TS1 {
    let mut item = syn::parse_macro_input!(tokens as ContractSubject);
    delete_invariants(&mut item);
    TS1::from(item.into_token_stream())
}

pub fn ensures(_: TS1, tokens: TS1) -> TS1 {
    let mut item = syn::parse_macro_input!(tokens as ContractSubject);
    delete_invariants(&mut item);
    TS1::from(item.into_token_stream())
}

pub fn variant(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn invariant(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn check(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn logic(_: TS1, _: TS1) -> TS1 {
    TS1::new()
}

pub fn trusted(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn opaque(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn intrinsic(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn builtin(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn extern_spec(_: TS1) -> TS1 {
    TS1::new()
}

pub fn maintains(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn erasure(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn open_inv_result(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn bitwise_proof(_: TS1, tokens: TS1) -> TS1 {
    tokens
}

pub fn derive_deep_model(_: TS1) -> TS1 {
    TS1::new()
}

/// Visitor to delete all `#[invariant]` and `#[creusot_std::invariant]`
/// attributes on loops.
struct DeleteInvariants;

impl VisitMut for DeleteInvariants {
    fn visit_expr_for_loop_mut(&mut self, i: &mut syn::ExprForLoop) {
        delete_invariants_attrs(&mut i.attrs);
        syn::visit_mut::visit_expr_for_loop_mut(self, i)
    }

    fn visit_expr_while_mut(&mut self, i: &mut syn::ExprWhile) {
        delete_invariants_attrs(&mut i.attrs);
        syn::visit_mut::visit_expr_while_mut(self, i)
    }

    fn visit_expr_loop_mut(&mut self, i: &mut syn::ExprLoop) {
        delete_invariants_attrs(&mut i.attrs);
        syn::visit_mut::visit_expr_loop_mut(self, i)
    }

    fn visit_expr_macro_mut(&mut self, i: &mut syn::ExprMacro) {
        delete_trusted_attrs(&mut i.attrs);
        syn::visit_mut::visit_expr_macro_mut(self, i)
    }

    fn visit_stmt_macro_mut(&mut self, i: &mut syn::StmtMacro) {
        delete_trusted_attrs(&mut i.attrs);
        syn::visit_mut::visit_stmt_macro_mut(self, i)
    }
}

// `invariant` or `creusot_std::invariant` or `::creusot_std::invariant`
fn is(name: &str, path: &syn::Path) -> bool {
    path.is_ident(name) || is_qualified(name, path)
}

fn is_qualified(name: &str, path: &syn::Path) -> bool {
    let mut segments = path.segments.iter();
    if let Some(first) = segments.next()
        && let Some(second) = segments.next()
    {
        return first.ident == "creusot_std" && second.ident == name;
    }
    false
}

fn delete_invariants_attrs(attrs: &mut Vec<syn::Attribute>) {
    attrs.retain(|attr| {
        if let syn::Meta::List(meta) = &attr.meta {
            !(is("invariant", &meta.path) || is("variant", &meta.path))
        } else {
            true
        }
    });
}

fn delete_contracts(attrs: &mut Vec<syn::Attribute>) {
    attrs.retain(|attr| {
        if let syn::Meta::List(meta) = &attr.meta {
            !is("requires", &meta.path) && !is("ensures", &meta.path)
        } else {
            true
        }
    });
}

fn delete_trusted_attrs(attrs: &mut Vec<syn::Attribute>) {
    attrs.retain(
        |attr| {
            if let syn::Meta::Path(path) = &attr.meta { !is("trusted", &path) } else { true }
        },
    )
}

// Also delete other contracts to avoid redundant passes
fn delete_invariants(item: &mut ContractSubject) {
    match item {
        ContractSubject::FnOrMethod(fn_or_method) => {
            delete_contracts(&mut fn_or_method.attrs);
            fn_or_method.body.iter_mut().for_each(|body| DeleteInvariants.visit_block_mut(body));
        }
        ContractSubject::Closure(expr) => {
            delete_contracts(&mut expr.attrs);
            DeleteInvariants.visit_expr_closure_mut(expr);
        }
    }
}
