use alloy_primitives::U256;
use eyre::Result;
use solar::ast::{ExprKind, Span};

use crate::mutation::{
    mutant::{Mutant, MutationType, OwnedLiteral},
    mutators::{MutationContext, Mutator},
    visitor::AssignVarTypes,
};

pub struct AssignmentMutator;

impl Mutator for AssignmentMutator {
    fn generate_mutants(&self, context: &MutationContext<'_>) -> Result<Vec<Mutant>> {
        let (assign_var_type, replacement_span) = match extract_rhs_info(context) {
            Some(info) => info,
            None => return Ok(vec![]), // is_applicable should filter this
        };

        let original = context.original_text();
        let source_line = context.source_line();
        let line_number = context.line_number();
        let column_number = context.column_number();

        match assign_var_type {
            AssignVarTypes::Literal(ref lit) => match lit {
                OwnedLiteral::Bool(val) => Ok(vec![Mutant {
                    span: replacement_span,
                    mutation: MutationType::Assignment(AssignVarTypes::Literal(
                        OwnedLiteral::Bool(!val),
                    )),
                    path: context.path.clone(),
                    original,
                    source_line,
                    line_number,
                    column_number,
                }]),
                OwnedLiteral::Number(val) if *val == U256::ZERO => Ok(vec![]),
                OwnedLiteral::Number(val) => Ok(vec![
                    Mutant {
                        span: replacement_span,
                        mutation: MutationType::Assignment(AssignVarTypes::Literal(
                            OwnedLiteral::Number(U256::ZERO),
                        )),
                        path: context.path.clone(),
                        original: original.clone(),
                        source_line: source_line.clone(),
                        line_number,
                        column_number,
                    },
                    // Negation of a numeric literal must be carried textually:
                    // applying `-*val` on a `U256` wraps via two's complement
                    // and would render as a huge unsigned literal (e.g. `1`
                    // becomes `2^256 - 1`), producing wrong-source mutants.
                    Mutant {
                        span: replacement_span,
                        mutation: MutationType::Assignment(AssignVarTypes::Literal(
                            OwnedLiteral::NegatedNumber(*val),
                        )),
                        path: context.path.clone(),
                        original,
                        source_line,
                        line_number,
                        column_number,
                    },
                ]),
                OwnedLiteral::Str { .. } => Ok(vec![]),
                OwnedLiteral::Rational(_) => Ok(vec![]),
                OwnedLiteral::Address(_) => Ok(vec![]),
                OwnedLiteral::Err(_) => Ok(vec![]),
                // `NegatedNumber` is only ever constructed *as* a mutant; it
                // does not appear as an original literal in the source AST,
                // so there is nothing to mutate here.
                OwnedLiteral::NegatedNumber(_) => Ok(vec![]),
            },
            AssignVarTypes::Identifier(ref ident) => Ok(vec![
                Mutant {
                    span: replacement_span,
                    mutation: MutationType::Assignment(AssignVarTypes::Literal(
                        OwnedLiteral::Number(U256::ZERO),
                    )),
                    path: context.path.clone(),
                    original: original.clone(),
                    source_line: source_line.clone(),
                    line_number,
                    column_number,
                },
                Mutant {
                    span: replacement_span,
                    mutation: MutationType::Assignment(AssignVarTypes::NegatedIdentifier(
                        ident.clone(),
                    )),
                    path: context.path.clone(),
                    original,
                    source_line,
                    line_number,
                    column_number,
                },
            ]),
            AssignVarTypes::NegatedIdentifier(_) => Ok(vec![]),
        }
    }

    /// Match is the expr is an assign with a var definition having a literal or identifier as
    /// initializer
    fn is_applicable(&self, context: &MutationContext<'_>) -> bool {
        if let Some(expr) = context.expr {
            if let ExprKind::Assign(_lhs, _op_opt, rhs_actual_expr) = &expr.kind {
                matches!(rhs_actual_expr.kind, ExprKind::Lit(..) | ExprKind::Ident(..))
            } else {
                false // Not an assign
            }
        } else if let Some(var_definition) = context.var_definition {
            if let Some(init) = &var_definition.initializer {
                matches!(&init.kind, ExprKind::Lit(..) | ExprKind::Ident(..))
            } else {
                false // No initializer
            }
        } else {
            false // Not an expression or var_definition
        }
    }
}

fn extract_rhs_info<'ast>(context: &MutationContext<'ast>) -> Option<(AssignVarTypes, Span)> {
    let relevant_expr_for_rhs = if let Some(var_definition) = context.var_definition {
        var_definition.initializer.as_ref()?
    } else {
        let expr = context.expr?;
        match &expr.kind {
            ExprKind::Assign(_lhs, _op_opt, rhs_actual_expr) => &**rhs_actual_expr,
            // If the context.expr is already what we want to get the type from
            // (e.g. a simple Lit or Ident being passed directly, though is_applicable filters this)
            ExprKind::Lit(..) | ExprKind::Ident(..) => expr,
            _ => return None,
        }
    };

    match &relevant_expr_for_rhs.kind {
        ExprKind::Lit(kind, _) => {
            let owned = OwnedLiteral::from(&kind.kind);
            Some((AssignVarTypes::Literal(owned), relevant_expr_for_rhs.span))
        }
        ExprKind::Ident(val) => {
            Some((AssignVarTypes::Identifier(val.to_string()), relevant_expr_for_rhs.span))
        }
        _ => None,
    }
}
