use crate::context::BodyCtxt;
use crate::verus_items::RustItem;
use rustc_span::Span;
use std::sync::Arc;
use vir::ast::{
    BinaryOp, Expr, ExprX, FunctionX, Mode, Place, PlaceX, SpannedTyped, VirErr, VirErrAs,
};
use vir::messages::WarningAllow;

/// Traits with special handling
#[derive(Clone, Copy, Debug)]
pub enum SpecialTrait {
    Clone,
    //PartialEq,
}

/// What to do for a given automatically-derived trait impl
#[derive(Debug, Clone, Copy)]
pub enum AutomaticDeriveAction {
    Special(SpecialTrait),
    VerifyAsIs,
    /// Ignore, optionally providing a warning
    Ignore,
}

pub fn get_action(rust_item: Option<RustItem>) -> AutomaticDeriveAction {
    match rust_item {
        Some(RustItem::PartialEq | RustItem::Eq) => AutomaticDeriveAction::Ignore,
        Some(RustItem::Clone) => AutomaticDeriveAction::Special(SpecialTrait::Clone),

        Some(RustItem::Copy) => AutomaticDeriveAction::VerifyAsIs,

        Some(RustItem::Hash)
        | Some(RustItem::Default)
        | Some(RustItem::Debug)
        | Some(RustItem::Ord)
        | Some(RustItem::PartialOrd) => AutomaticDeriveAction::Ignore,

        Some(_) | None => AutomaticDeriveAction::VerifyAsIs,
    }
}

pub fn is_automatically_derived(attrs: &[rustc_hir::Attribute]) -> bool {
    for attr in attrs.iter() {
        match attr {
            rustc_hir::Attribute::Unparsed(item) => match &item.path.segments[..] {
                [segment] => {
                    if segment.as_str() == "automatically_derived" {
                        return true;
                    }
                }
                _ => {}
            },
            rustc_hir::Attribute::Parsed(rustc_hir::attrs::AttributeKind::AutomaticallyDerived) => {
                return true;
            }
            _ => {}
        }
    }
    false
}

pub fn modify_derived_item<'tcx>(
    bctx: &BodyCtxt<'tcx>,
    id: rustc_span::def_id::DefId,
    inputs: &Vec<rustc_middle::ty::Ty>,
    span: Span,
    action: AutomaticDeriveAction,
    function: &mut FunctionX,
) -> Result<(), VirErr> {
    let AutomaticDeriveAction::Special(special) = action else {
        return Ok(());
    };
    match special {
        SpecialTrait::Clone => {
            if &*function.name.path.last_segment() == "clone" {
                return clone_add_post_condition(bctx, id, inputs, span, function);
            }
        }
    }
    Ok(())
}

fn clone_add_post_condition<'tcx>(
    bctx: &BodyCtxt<'tcx>,
    mut id: rustc_span::def_id::DefId,
    inputs: &Vec<rustc_middle::ty::Ty>,
    span: Span,
    functionx: &mut FunctionX,
) -> Result<(), VirErr> {
    if inputs.len() >= 1 {
        use rustc_middle::ty::{AdtDef, TyKind};
        if let TyKind::Ref(_, t, _) = inputs[0].kind() {
            if let TyKind::Adt(AdtDef(adt_def_data), _) = t.kind() {
                // It's more convenient to put verifier::allow on the datatype than on the function
                id = adt_def_data.did;
            }
        }
    }
    let warn = |msg: &str| {
        crate::attributes::warning_maybe(
            bctx.ctxt.tcx,
            id,
            span,
            &WarningAllow::AutoderiveCloneWithoutSpec,
            || msg,
            |msg| bctx.ctxt.diagnostics.borrow_mut().push(VirErrAs::Warning(msg)),
        );
    };
    let warn_unexpected = || {
        warn(
            "autoderive Clone impl does not take the form Verus expects; continuing, but without adding a specification for the derived Clone impl",
        )
    };
    let warn_unsupported = || {
        warn(
            "Verus does not currently support autoderive Clone impl when the clone is not a copy; continuing, but without adding a specification for the derived Clone impl",
        )
    };

    let Some(body) = &functionx.body else {
        return Ok(());
    };

    let uses_copy;
    let self_var;

    match &body.x {
        ExprX::Block(_stmts, Some(last_expr)) => match &last_expr.x {
            ExprX::ReadPlace(pl, _) => match &pl.x {
                PlaceX::Local(id) if &*id.0 == "self" => {
                    uses_copy = true;
                    self_var = Some(last_expr.clone());
                }
                _ => {
                    warn_unexpected();
                    return Ok(());
                }
            },
            ExprX::Ctor { .. } => {
                uses_copy = false;
                self_var = None;
            }
            _ => {
                warn_unexpected();
                return Ok(());
            }
        },
        _ => {
            warn_unexpected();
            return Ok(());
        }
    }

    if functionx.ensure.0.len() != 0 {
        warn_unexpected();
        return Ok(());
    }

    if uses_copy {
        // Add `ensures ret == self`
        let self_var = self_var.unwrap();
        let ret_var = SpannedTyped::new(
            &self_var.span,
            &self_var.typ,
            ExprX::Var(functionx.inner_ret.x.name.clone()),
        );
        let eq_expr = SpannedTyped::new(
            &self_var.span,
            &vir::ast_util::bool_typ(),
            ExprX::Binary(BinaryOp::Eq(Mode::Spec), ret_var.clone(), self_var.clone()),
        );

        let eq_expr = cleanup_span_ids(bctx, span, &eq_expr);
        functionx.ensure.0 = Arc::new(vec![eq_expr]);
    } else {
        warn_unsupported();
    }

    Ok(())
}

/// Give each node a unique ID. This inserts 'None' for every var into the local_hir_vir_ids
/// list, so it's only safe for pure expressions (e.g., in asserts or signatures).
fn cleanup_span_ids<'tcx>(bctx: &BodyCtxt<'tcx>, span: Span, expr: &Expr) -> Expr {
    vir::ast_visitor::map_expr_place_visitor(
        expr,
        &|e: &Expr| {
            let e = bctx.ctxt.spans.spanned_typed_new(span, &e.typ, e.x.clone());
            if matches!(&e.x, ExprX::Var(..) | ExprX::VarAt(..)) {
                let mut erasure_info = bctx.erasure_info.borrow_mut();
                erasure_info.local_hir_vir_ids.push((None, e.span.id));
            }
            Ok(e)
        },
        &|p: &Place| {
            let p = bctx.ctxt.spans.spanned_typed_new(span, &p.typ, p.x.clone());
            if matches!(&p.x, PlaceX::Local(..)) {
                let mut erasure_info = bctx.erasure_info.borrow_mut();
                erasure_info.local_hir_vir_ids.push((None, p.span.id));
            }
            Ok(p)
        },
    )
    .unwrap()
}