use crate::ast::Field;
use crate::ast::{
    AssertProof, MonoidElt, MonoidStmtType, ShardableType, SimplStmt, SpecialOp, SplitKind, SubIdx,
    TransitionKind, TransitionStmt, SM,
};
use crate::check_bind_stmts::get_binding_ty;
use crate::check_bind_stmts::uses_bind;
use crate::transitions::get_field;
use proc_macro2::Span;
use quote::{quote, quote_spanned};
use syn::spanned::Spanned;
use syn::{Expr, Ident, Type};

/// Simplify out `update` statements, including `add_element` etc.
///
/// Note: for 'readonly' stuff, there's less to do because we don't need to handle
/// updates. However, we still need to handle 'guard' and 'have' statements, which will
/// be translated into 'asserts'.

// Implementation:
//
// The simplification process has 3 primary passes here:
//
//  1. Turn the TransitionStmt into a SimplStmt by translating all the update, init,
//     and special ops into require, assert, and assign statements.
//     For example, `update x = 5` becomes `assign tmp_update_x = 5;`
//     or `add x += z` becomes
//        `assert (tmp_update_x + z is valid); assign tmp_update_x = tmp_update_x + z`
//
//  2. Add PostCondition statements that relate the final values of the tmp variables
//     to the post.{field} values.
//
//  3. A transformation pass to remove 'assign' statements, turning them into 'let' statements.
//     This makes it easier to turn the statement into a pure expression.
//
// So for example if we have fields {a, b, c} and a transition declared as
//
//    `update a = 5;`
//
// Then after steps 1 and 2 we would have:
//
//    assign tmp_update_a = pre.a;      // Prologue initializing each tmp
//    assign tmp_update_b = pre.b;      // variable to the initial state
//    assign tmp_update_c = pre.c;      // (This part is skipped for init transitions)
//
//    assign tmp_update_a = 5;          // Translated from the 'update' statement
//
//    postcondition post.a == tmp_update_a    // Added in phase 2
//    postcondition post.b == tmp_update_a
//    postcondition post.c == tmp_update_a
//
// Thus for the second phase, placing the PostCondition statements,
// the key is that the PostCondition statement has to go some
// place where temp_foo has taken on its final value (i.e., the PostCondition can't come
// before any statement which might update the value). Granted, one option would be to
// always put them at the end (in which case one might ask why we bother).
//
// The reason is because of conditionals. Consider,
//
//      if cond {
//          update(foo, x);
//      } else {
//          update(foo, y);
//      }
//
// One option would be to generate a relation that looks like,
//      `post.foo == (if cond { x } else { y })`
// But for better user experience, we'd ideally want one predicate per 'update' statement,
// since more fine-grained predicates make it easier to diagnose errors and each predicate
// could then be associated with the source line of a single 'update' statement.
// So we would place the PostCondition statements like this:
//
//      if cond {
//          assign tmp_update_foo = x;                 // from update(foo, x)
//          PostCondition(post.foo == tmp_update_foo); // inserted
//      } else {
//          assign tmp_update_foo = y;                 // from update(foo, y)
//          PostCondition(post.foo == tmp_update_foo); // inserted
//      }
//
// Which then generates relations like:
//
//      `cond ==> post.foo == x`
//      `!cond ==> post.foo == y`
//
// Thus the purpose of this phase is to find these ideal positions for the
// PostCondition statements.

pub fn simplify_ops(sm: &SM, ts: &TransitionStmt, kind: TransitionKind) -> Vec<SimplStmt> {
    // Phase 1: translate the update, init, and special ops into SimplStmts
    let sops = simplify_ops_with_pre(&ts, &sm, kind);

    // Phase 2: Add PostCondition statements
    let sops = add_postconditions(sm, *ts.get_span(), sops);

    // Phase 3: Simplify out all 'assign' statements
    crate::simplify_assigns::simplify_assigns(&sm, &sops)
}

// Phase 2. Adding the PostCondition operations.
//
// The key correctness criteria are:
//
//    1. A postcondition for a field `foo` cannot come before a statement that updates `foo`
//    2. Every control-flow path must encounter exactly one PostCondition statement
//
// Other than that, we want the PostCondition to come as soon as possible.
// So we basically just walk the tree backwards, keeping track of which fields we have
// made placements for. When we encounter the first update statement (the first from the end,
// that is) we add the PostCondition.
//
// For each conditional, we have to check if we added a statement on one branch but not
// the other, and if so, resolve.
// Finally at the very end, we add postconditions for any fields that were never updated.

fn add_postconditions(sm: &SM, span: Span, sops: Vec<SimplStmt>) -> Vec<SimplStmt> {
    let mut found = Vec::new();
    let mut sops = sops;
    add_postconditions_vec(&mut sops, &mut found);

    for field in &sm.fields {
        if !contains_ident(&found, &field.name) {
            let fs = postcondition_stmt(span, field.name.clone());
            sops.push(fs);
        }
    }

    sops
}

fn add_postconditions_vec(sops: &mut Vec<SimplStmt>, found: &mut Vec<Ident>) {
    // To be appended at the end
    let mut new_sops = Vec::new();

    for sop in sops.iter_mut().rev() {
        // Add a postcondition statement if necessary.

        match &sop {
            SimplStmt::Assign(span, tmp_ident, _, _) => {
                let f_ident_opt = field_name_from_tmp(tmp_ident);
                if let Some(f_ident) = f_ident_opt {
                    if !contains_ident(found, &f_ident) {
                        // If it's an assign statement, AND we haven't added
                        // a postcondition for this field yet, then add a postcondition
                        // at the end of this block of statements.
                        // (And leave the current statement unchanged).

                        found.push(f_ident.clone());
                        new_sops.push(postcondition_stmt(*span, f_ident));
                    }
                }
            }
            _ => {}
        }

        // Recurse.

        match sop {
            SimplStmt::Let(_, _, _, _, v) => {
                add_postconditions_vec(v, found);
            }
            SimplStmt::Split(span, _, es) => {
                let idx = found.len();
                let mut found_inner = vec![found.clone(); es.len()];

                for i in 0..es.len() {
                    add_postconditions_vec(&mut es[i], &mut found_inner[i]);
                }

                // For each side of the conditional (or arm of the match),
                // look at any newly-found
                // fields from that conditional (those after `idx`, the original
                // length of the array). For such a field, if it wasn't ALSO found
                // in some other branch, then we go ahead and add it to the other
                // branch now. Thus we maintain that, for each field and for each
                // conditional, we will either get a postcondition on both branches,
                // or on neither.

                // Step 1: Collect any field which is newly updated in any branch.

                let mut all_new = Vec::new();

                for j in 0..es.len() {
                    for i in idx..found_inner[j].len() {
                        if !contains_ident(&all_new, &found_inner[j][i]) {
                            all_new.push(found_inner[j][i].clone());
                        }
                    }
                }

                // Step 2: For each field and each branch, check if it needs to be
                // added, and if so, do so.

                for f in &all_new {
                    for j in 0..es.len() {
                        if !contains_ident(&found_inner[j], f) {
                            found_inner[j].push(f.clone());
                            es[j].push(postcondition_stmt(*span, f.clone()));
                        }
                    }
                }

                // Make sure we end with `found` (the &mut argument) containing the
                // union of all the fields that were found on either branch.

                *found = found_inner[0].clone();
            }

            SimplStmt::Require(..)
            | SimplStmt::PostCondition(..)
            | SimplStmt::Assert(..)
            | SimplStmt::Assign(..) => {}
        }
    }

    sops.append(&mut new_sops);
}

fn postcondition_stmt(span: Span, f: Ident) -> SimplStmt {
    let cur = get_cur(&f);
    SimplStmt::PostCondition(
        span,
        Expr::Verbatim(quote! {
            ::builtin::equal(post.#f, #cur)
        }),
    )
}

fn contains_ident(v: &Vec<Ident>, id: &Ident) -> bool {
    for id0 in v {
        if id0.to_string() == id.to_string() {
            return true;
        }
    }
    return false;
}

fn get_cur(field_name: &Ident) -> Expr {
    let id = get_cur_ident(field_name);
    Expr::Verbatim(quote! { #id })
}

pub const UPDATE_TMP_PREFIX: &str = "update_tmp_";

fn get_cur_ident(field_name: &Ident) -> Ident {
    let name = UPDATE_TMP_PREFIX.to_string() + &field_name.to_string();
    Ident::new(&name, field_name.span())
}

fn field_name_from_tmp(tmp_name: &Ident) -> Option<Ident> {
    let s = tmp_name.to_string();
    if s.starts_with(UPDATE_TMP_PREFIX) {
        Some(Ident::new(&s[UPDATE_TMP_PREFIX.len()..], tmp_name.span()))
    } else {
        None
    }
}

// Phase 1. Give meaning to all the update, init, and special op statements.
// User-provided require and assert statements are just translated as they are.

fn simplify_ops_with_pre(ts: &TransitionStmt, sm: &SM, kind: TransitionKind) -> Vec<SimplStmt> {
    let mut ops = simplify_ops_rec(ts, sm);
    if kind == TransitionKind::Init {
        ops
    } else {
        let mut ops1: Vec<SimplStmt> = sm
            .fields
            .iter()
            .map(|f| {
                let f_ident = &f.name;
                SimplStmt::Assign(
                    f_ident.span(),
                    get_cur_ident(f_ident),
                    f.get_type(),
                    Expr::Verbatim(quote! {pre.#f_ident}),
                )
            })
            .collect();
        ops1.append(&mut ops);
        ops1
    }
}

fn simplify_ops_rec(ts: &TransitionStmt, sm: &SM) -> Vec<SimplStmt> {
    match ts {
        TransitionStmt::Block(_span, v) => {
            let mut res = Vec::new();
            for t in v {
                let mut t = simplify_ops_rec(t, sm);
                res.append(&mut t);
            }
            res
        }
        TransitionStmt::Split(span, split_kind, es) => match split_kind {
            SplitKind::If(..) | SplitKind::Match(..) => {
                let mut new_es: Vec<Vec<SimplStmt>> = Vec::new();
                for e in es {
                    let new_e1 = simplify_ops_rec(&e, sm);
                    new_es.push(new_e1);
                }

                vec![SimplStmt::Split(*span, split_kind.clone(), new_es)]
            }
            SplitKind::Let(pat, ty, _lk, e) => {
                assert!(es.len() == 1);
                let child = es.into_iter().next().expect("child");
                let new_child = simplify_ops_rec(child, sm);
                vec![SimplStmt::Let(*span, pat.clone(), ty.clone(), e.clone(), new_child)]
            }
            SplitKind::Special(f, op, proof, pat_opt) => {
                let field = get_field(&sm.fields, f);
                let (cond_ops, mut update_ops) = simplify_special_op(*span, &field, op, proof);

                match pat_opt {
                    None => {
                        assert!(es.len() == 1 && es[0].is_trivial());
                        assert!(!uses_bind(&op.elt));
                        let mut res = cond_ops;
                        res.append(&mut update_ops);
                        res
                    }
                    Some(pat) => {
                        assert!(es.len() == 1);
                        assert!(uses_bind(&op.elt));
                        let child = es.into_iter().next().expect("child");
                        let mut new_child = simplify_ops_rec(child, sm);

                        let mut all_children = update_ops;
                        all_children.append(&mut new_child);

                        let ty = get_binding_ty(&field.stype);
                        let initializer = get_initializer_expr(f, op);

                        let mut res = cond_ops;
                        res.push(SimplStmt::Let(*span, pat.clone(), ty, initializer, all_children));
                        res
                    }
                }
            }
        },

        TransitionStmt::PostCondition(span, e) => {
            vec![SimplStmt::PostCondition(*span, e.clone())]
        }
        TransitionStmt::Require(span, e) => {
            vec![SimplStmt::Require(*span, e.clone())]
        }
        TransitionStmt::Assert(span, e, pf) => {
            vec![SimplStmt::Assert(*span, e.clone(), pf.clone())]
        }

        TransitionStmt::Initialize(span, f, e) | TransitionStmt::Update(span, f, e) => {
            let field = get_field(&sm.fields, f);
            vec![SimplStmt::Assign(*span, get_cur_ident(f), field.get_type(), e.clone())]
        }
        TransitionStmt::SubUpdate(span, f, subs, e) => {
            let field = get_field(&sm.fields, f);
            vec![SimplStmt::Assign(
                *span,
                get_cur_ident(f),
                field.get_type(),
                update_sub_expr(&get_cur(f), subs, 0, e),
            )]
        }
    }
}

fn simplify_special_op(
    span: Span,
    field: &Field,
    op: &SpecialOp,
    proof: &AssertProof,
) -> (Vec<SimplStmt>, Vec<SimplStmt>) {
    match op {
        SpecialOp { stmt: MonoidStmtType::Have, elt } => {
            let cur = get_cur(&field.name);
            let prec = expr_ge(&field.stype, &cur, elt);
            (vec![SimplStmt::Require(span, prec)], vec![])
        }

        SpecialOp { stmt: MonoidStmtType::Add, elt } => {
            let cur = get_cur(&field.name);
            let new_val = expr_add(&field.stype, &cur, elt);

            (
                match expr_can_add(&field.stype, &cur, elt) {
                    Some(safety) => vec![SimplStmt::Assert(span, safety, proof.clone())],
                    None => vec![],
                },
                vec![SimplStmt::Assign(
                    span,
                    get_cur_ident(&field.name),
                    field.get_type(),
                    new_val,
                )],
            )
        }

        SpecialOp { stmt: MonoidStmtType::Remove, elt } => {
            let cur = get_cur(&field.name);
            let new_val = expr_remove(&field.stype, &cur, elt);
            let prec = expr_ge(&field.stype, &cur, elt);
            (
                vec![SimplStmt::Require(span, prec)],
                vec![SimplStmt::Assign(
                    span,
                    get_cur_ident(&field.name),
                    field.get_type(),
                    new_val,
                )],
            )
        }

        SpecialOp { stmt: MonoidStmtType::Guard, elt } => {
            let cur = get_cur(&field.name);
            let prec = expr_ge(&field.stype, &cur, elt);
            (vec![SimplStmt::Assert(span, prec, proof.clone())], vec![])
        }

        SpecialOp { stmt: MonoidStmtType::Deposit, elt } => {
            let cur = get_cur(&field.name);
            let new_val = expr_add(&field.stype, &cur, elt);
            (
                match expr_can_add(&field.stype, &cur, elt) {
                    Some(safety) => vec![SimplStmt::Assert(span, safety, proof.clone())],
                    None => vec![],
                },
                vec![SimplStmt::Assign(
                    span,
                    get_cur_ident(&field.name),
                    field.get_type(),
                    new_val,
                )],
            )
        }

        SpecialOp { stmt: MonoidStmtType::Withdraw, elt } => {
            let cur = get_cur(&field.name);
            let new_val = expr_remove(&field.stype, &cur, elt);
            let prec = expr_ge(&field.stype, &cur, elt);
            (
                vec![SimplStmt::Assert(span, prec, proof.clone())],
                vec![SimplStmt::Assign(
                    span,
                    get_cur_ident(&field.name),
                    field.get_type(),
                    new_val,
                )],
            )
        }
    }
}

fn get_initializer_expr(f: &Ident, op: &SpecialOp) -> Expr {
    let cur = get_cur(f);
    match &op.elt {
        MonoidElt::OptionSome(None) => Expr::Verbatim(quote! {
            #cur.get_Some_0()
        }),
        MonoidElt::SingletonKV(key, None) => Expr::Verbatim(quote! {
            #cur.index(#key)
        }),
        _ => {
            panic!("simplify_let_value got unexpected elt");
        }
    }
}

fn update_sub_expr(root: &Expr, subs: &Vec<SubIdx>, i: usize, val: &Expr) -> Expr {
    if i == subs.len() {
        val.clone()
    } else {
        match &subs[i] {
            SubIdx::Field(_field) => {
                panic!("dot-fields in sub-updates not yet supported");
            }
            SubIdx::Idx(idx_e) => {
                let child = Expr::Verbatim(quote_spanned! { idx_e.span() => #root.index(#idx_e) });
                let r = update_sub_expr(&child, subs, i + 1, val);
                Expr::Verbatim(quote_spanned! { idx_e.span() =>
                    #root.update(#idx_e, #r)
                })
            }
        }
    }
}

fn get_opt_type(stype: &ShardableType) -> Type {
    match stype {
        ShardableType::Option(ty) => ty.clone(),
        ShardableType::StorageOption(ty) => ty.clone(),
        _ => {
            panic!("get_opt_type expected option");
        }
    }
}

fn expr_can_add(stype: &ShardableType, cur: &Expr, elt: &MonoidElt) -> Option<Expr> {
    if stype.is_persistent() {
        match elt {
            MonoidElt::SingletonKV(key, val) => Some(Expr::Verbatim(quote! {
                ::builtin::imply(
                    (#cur).dom().contains(#key),
                    ::builtin::equal((#cur).index(#key), #val),
                )
            })),
            MonoidElt::General(e) => match stype {
                ShardableType::PersistentMap(_, _) => Some(Expr::Verbatim(quote! {
                    (#cur).agrees(#e).dom()
                })),
                _ => {
                    panic!("expr_can_add invalid case");
                }
            },
            _ => {
                panic!("expr_can_add invalid case");
            }
        }
    } else {
        match elt {
            MonoidElt::OptionSome(_e) => Some(Expr::Verbatim(quote! { (#cur).is_None() })),
            MonoidElt::SingletonKV(key, _val) => {
                Some(Expr::Verbatim(quote! { !(#cur).dom().contains(#key) }))
            }
            MonoidElt::SingletonMultiset(_e) => None,
            MonoidElt::General(e) => match stype {
                ShardableType::Option(_) | ShardableType::StorageOption(_) => {
                    Some(Expr::Verbatim(quote! {
                        ((#e).is_None() || (#cur).is_None())
                    }))
                }

                ShardableType::Map(_, _) | ShardableType::StorageMap(_, _) => {
                    Some(Expr::Verbatim(quote! {
                        (#cur).dom().disjoint((#e).dom())
                    }))
                }

                ShardableType::Multiset(_) => None,
                ShardableType::Count => None,

                _ => {
                    panic!("expected option/map/multiset/count");
                }
            },
        }
    }
}

/// Add the elements, assuming they are addable. Build the expression assuming
/// to "prefer" the content from the right-hand side, even though it shouldn't really matter.
/// That's often easier, anyway.
fn expr_add(stype: &ShardableType, cur: &Expr, elt: &MonoidElt) -> Expr {
    if stype.is_persistent() {
        match elt {
            MonoidElt::SingletonKV(key, val) => {
                Expr::Verbatim(quote! { (#cur).insert(#key, #val) })
            }
            MonoidElt::General(e) => match stype {
                ShardableType::PersistentMap(_, _) => {
                    Expr::Verbatim(quote! { (#cur).union_prefer_right(#e) })
                }
                _ => {
                    panic!("expr_can_add invalid case");
                }
            },
            _ => {
                panic!("expr_can_add invalid case");
            }
        }
    } else {
        match elt {
            MonoidElt::OptionSome(Some(e)) => {
                let ty = get_opt_type(stype);
                Expr::Verbatim(quote! {
                    crate::pervasive::option::Option::<#ty>::Some(#e)
                })
            }
            MonoidElt::OptionSome(None) => {
                panic!("expr_add: no value found");
            }
            MonoidElt::SingletonKV(key, Some(val)) => Expr::Verbatim(quote! {
                (#cur).insert(#key, #val)
            }),
            MonoidElt::SingletonKV(_key, None) => {
                panic!("expr_add: no value found");
            }
            MonoidElt::SingletonMultiset(e) => Expr::Verbatim(quote! {
                (#cur).insert(#e)
            }),
            MonoidElt::General(e) => match stype {
                ShardableType::Option(_) | ShardableType::StorageOption(_) => {
                    Expr::Verbatim(quote! {
                        (#e).or(#cur)
                    })
                }

                ShardableType::Map(_, _) | ShardableType::StorageMap(_, _) => {
                    Expr::Verbatim(quote! {
                        (#cur).union_prefer_right(#e)
                    })
                }

                ShardableType::Multiset(_) => Expr::Verbatim(quote! {
                    (#cur).add(#e)
                }),

                ShardableType::Count => Expr::Verbatim(quote! {
                    (#cur) + (#e)
                }),

                _ => {
                    panic!("expected option/map/multiset");
                }
            },
        }
    }
}

fn expr_ge(stype: &ShardableType, cur: &Expr, elt: &MonoidElt) -> Expr {
    // note: persistent case should always be the same as the normal case

    match elt {
        MonoidElt::OptionSome(None) => Expr::Verbatim(quote! { (#cur).is_Some() }),
        MonoidElt::OptionSome(Some(e)) => {
            let ty = get_opt_type(stype);
            Expr::Verbatim(quote! {
                ::builtin::equal(
                    #cur,
                    crate::pervasive::option::Option::<#ty>::Some(#e)
                )
            })
        }
        MonoidElt::SingletonKV(key, None) => Expr::Verbatim(quote! {
            (#cur).dom().contains(#key)
        }),
        MonoidElt::SingletonKV(key, Some(val)) => Expr::Verbatim(quote! {
            (#cur).contains_pair(#key, #val)
        }),
        MonoidElt::SingletonMultiset(e) => Expr::Verbatim(quote! {
            (#cur).count(#e) >= 1
        }),
        MonoidElt::General(e) => match stype {
            ShardableType::Option(_) | ShardableType::StorageOption(_) => Expr::Verbatim(quote! {
                ((#e).is_Some() >>= ::builtin::equal(#cur, #e))
            }),

            ShardableType::Map(_, _)
            | ShardableType::PersistentMap(_, _)
            | ShardableType::StorageMap(_, _) => Expr::Verbatim(quote! {
                (#e).le(#cur)
            }),

            ShardableType::Multiset(_) => Expr::Verbatim(quote! {
                (#e).le(#cur)
            }),

            ShardableType::Count => Expr::Verbatim(quote! {
                (#cur) >= (#e)
            }),

            _ => {
                panic!("expected option/map/multiset");
            }
        },
    }
}

fn expr_remove(stype: &ShardableType, cur: &Expr, elt: &MonoidElt) -> Expr {
    assert!(!stype.is_persistent());
    match elt {
        MonoidElt::OptionSome(_e) => {
            let ty = get_opt_type(stype);
            Expr::Verbatim(quote! {
                crate::pervasive::option::Option::<#ty>::None
            })
        }
        MonoidElt::SingletonKV(key, _val) => Expr::Verbatim(quote! {
            (#cur).remove(#key)
        }),
        MonoidElt::SingletonMultiset(e) => Expr::Verbatim(quote! {
            (#cur).remove(#e)
        }),
        MonoidElt::General(e) => match stype {
            ShardableType::Option(_) | ShardableType::StorageOption(_) => Expr::Verbatim(quote! {
                (if (#e).is_Some() { crate::pervasive::option::Option::None } else { #cur })
            }),

            ShardableType::Map(_, _) | ShardableType::StorageMap(_, _) => Expr::Verbatim(quote! {
                (#cur).remove_keys(#e.dom())
            }),

            ShardableType::Multiset(_) => Expr::Verbatim(quote! {
                (#cur).sub(#e)
            }),

            ShardableType::Count => Expr::Verbatim(quote! {
                (#cur) - (#e)
            }),

            _ => {
                panic!("expected option/map/multiset");
            }
        },
    }
}
