use crate::ast::Span;
use crate::mir::optimizer::MirOptimizer;
use crate::mir::optimizer_stats::OptimizationStats;
use crate::mir::{
    BarrierOp, EffectMask, MirModule, SpannedInstruction, TypeOpKind, ValueId, WeakRefOp,
};

fn idemp_enabled() -> bool {
    std::env::var("NYASH_MIR_DEV_IDEMP").ok().as_deref() == Some("1")
}

fn idemp_key(pass: &str, func_name: &str) -> String {
    format!("{}:{}", pass, func_name)
}

fn idemp_already_done(module: &crate::mir::MirModule, key: &str) -> bool {
    module.metadata.dev_processed_markers.contains(key)
}

fn idemp_mark(module: &mut crate::mir::MirModule, key: String) {
    module.metadata.dev_processed_markers.insert(key);
}

pub fn force_plugin_invoke(_opt: &mut MirOptimizer, module: &mut MirModule) -> OptimizationStats {
    use crate::mir::MirInstruction as I;
    let mut stats = OptimizationStats::new();
    let pass_name = "normalize.force_plugin_invoke";
    let func_names: Vec<String> = module.functions.keys().cloned().collect();
    for fname in func_names {
        if idemp_enabled() {
            let key = idemp_key(pass_name, &fname);
            if idemp_already_done(module, &key) {
                continue;
            }
        }
        let function = match module.functions.get_mut(&fname) {
            Some(f) => f,
            None => continue,
        };
        for (_bb, block) in &mut function.blocks {
            let old_spanned = block.drain_spanned_instructions();
            let mut new_spanned = Vec::with_capacity(old_spanned.len());
            for SpannedInstruction { inst, span } in old_spanned {
                let rewritten = match inst {
                    I::BoxCall {
                        dst,
                        box_val,
                        method,
                        args,
                        effects,
                        ..
                    } => {
                        stats.intrinsic_optimizations += 1;
                        I::PluginInvoke {
                            dst,
                            box_val,
                            method,
                            args,
                            effects,
                        }
                    }
                    other => other,
                };
                new_spanned.push(SpannedInstruction {
                    inst: rewritten,
                    span,
                });
            }

            let (insts, spans): (Vec<_>, Vec<_>) =
                new_spanned.into_iter().map(|sp| (sp.inst, sp.span)).unzip();
            block.instructions = insts;
            block.instruction_spans = spans;
            block.effects = block
                .instructions
                .iter()
                .chain(block.terminator.iter())
                .fold(EffectMask::PURE, |mask, inst| mask | inst.effects());
        }
        if idemp_enabled() {
            let key = idemp_key(pass_name, &fname);
            idemp_mark(module, key);
        }
    }
    stats
}

pub fn normalize_python_helper_calls(
    _opt: &mut MirOptimizer,
    module: &mut MirModule,
) -> OptimizationStats {
    use crate::mir::MirInstruction as I;
    let mut stats = OptimizationStats::new();
    let pass_name = "normalize.python_helper_calls";
    let func_names: Vec<String> = module.functions.keys().cloned().collect();
    for fname in func_names {
        if idemp_enabled() {
            let key = idemp_key(pass_name, &fname);
            if idemp_already_done(module, &key) {
                continue;
            }
        }
        let function = match module.functions.get_mut(&fname) {
            Some(f) => f,
            None => continue,
        };
        for (_bb, block) in &mut function.blocks {
            let old_spanned = block.drain_spanned_instructions();
            let mut new_spanned = Vec::with_capacity(old_spanned.len());
            for SpannedInstruction { mut inst, span } in old_spanned {
                if let I::PluginInvoke {
                    ref mut box_val,
                    ref method,
                    ref mut args,
                    ..
                } = inst
                {
                    if method == "getattr" && args.len() >= 2 {
                        let new_recv = args[0];
                        args.remove(0);
                        *box_val = new_recv;
                        stats.intrinsic_optimizations += 1;
                    } else if method == "call" && !args.is_empty() {
                        let new_recv = args[0];
                        args.remove(0);
                        *box_val = new_recv;
                        stats.intrinsic_optimizations += 1;
                    }
                }
                new_spanned.push(SpannedInstruction { inst, span });
            }

            let (insts, spans): (Vec<_>, Vec<_>) =
                new_spanned.into_iter().map(|sp| (sp.inst, sp.span)).unzip();
            block.instructions = insts;
            block.instruction_spans = spans;
            block.effects = block
                .instructions
                .iter()
                .chain(block.terminator.iter())
                .fold(EffectMask::PURE, |mask, inst| mask | inst.effects());
        }
        if idemp_enabled() {
            let key = idemp_key(pass_name, &fname);
            idemp_mark(module, key);
        }
    }
    stats
}

pub fn normalize_legacy_instructions(
    _opt: &mut MirOptimizer,
    module: &mut MirModule,
) -> OptimizationStats {
    use crate::mir::MirInstruction as I;
    let mut stats = OptimizationStats::new();
    let rw_dbg = crate::config::env::rewrite_debug();
    let rw_sp = crate::config::env::rewrite_safepoint();
    let rw_future = crate::config::env::rewrite_future();
    let core13 = crate::config::env::mir_core13();
    let mut array_to_boxcall = crate::config::env::mir_array_boxcall();
    if core13 {
        array_to_boxcall = true;
    }
    let pass_name = "normalize.legacy_instructions";
    let func_names: Vec<String> = module.functions.keys().cloned().collect();
    for fname in func_names {
        if idemp_enabled() {
            let key = idemp_key(pass_name, &fname);
            if idemp_already_done(module, &key) {
                continue;
            }
        }
        let function = match module.functions.get_mut(&fname) {
            Some(f) => f,
            None => continue,
        };
        for (_bb, block) in &mut function.blocks {
            let old_spanned = block.drain_spanned_instructions();
            let mut new_spanned = Vec::with_capacity(old_spanned.len());
            for SpannedInstruction { inst, span } in old_spanned {
                let rewritten = match inst {
                    I::WeakNew { dst, box_val } => {
                        stats.intrinsic_optimizations += 1;
                        I::WeakRef {
                            dst,
                            op: WeakRefOp::New,
                            value: box_val,
                        }
                    }
                    I::WeakLoad { dst, weak_ref } => {
                        stats.intrinsic_optimizations += 1;
                        I::WeakRef {
                            dst,
                            op: WeakRefOp::Load,
                            value: weak_ref,
                        }
                    }
                    I::BarrierRead { ptr } => {
                        stats.intrinsic_optimizations += 1;
                        I::Barrier {
                            op: BarrierOp::Read,
                            ptr,
                        }
                    }
                    I::BarrierWrite { ptr } => {
                        stats.intrinsic_optimizations += 1;
                        I::Barrier {
                            op: BarrierOp::Write,
                            ptr,
                        }
                    }
                    I::Print { value, .. } => {
                        stats.intrinsic_optimizations += 1;
                        I::ExternCall {
                            dst: None,
                            iface_name: "env.console".to_string(),
                            method_name: "log".to_string(),
                            args: vec![value],
                            effects: crate::mir::EffectMask::PURE.add(crate::mir::Effect::Io),
                        }
                    }
                    I::ArrayGet { dst, array, index } if array_to_boxcall => {
                        stats.intrinsic_optimizations += 1;
                        let mid =
                            crate::mir::slot_registry::resolve_slot_by_type_name("ArrayBox", "get");
                        I::BoxCall {
                            dst: Some(dst),
                            box_val: array,
                            method: "get".to_string(),
                            method_id: mid,
                            args: vec![index],
                            effects: crate::mir::EffectMask::READ,
                        }
                    }
                    I::ArraySet {
                        array,
                        index,
                        value,
                    } if array_to_boxcall => {
                        stats.intrinsic_optimizations += 1;
                        let mid =
                            crate::mir::slot_registry::resolve_slot_by_type_name("ArrayBox", "set");
                        I::BoxCall {
                            dst: None,
                            box_val: array,
                            method: "set".to_string(),
                            method_id: mid,
                            args: vec![index, value],
                            effects: crate::mir::EffectMask::WRITE,
                        }
                    }
                    I::PluginInvoke {
                        dst,
                        box_val,
                        method,
                        args,
                        effects,
                    } => {
                        stats.intrinsic_optimizations += 1;
                        I::BoxCall {
                            dst,
                            box_val,
                            method,
                            method_id: None,
                            args,
                            effects,
                        }
                    }
                    I::Debug { .. } if !rw_dbg => I::Nop,
                    I::Safepoint if !rw_sp => I::Nop,
                    I::FutureNew { dst, value } if rw_future => I::ExternCall {
                        dst: Some(dst),
                        iface_name: "env.future".to_string(),
                        method_name: "new".to_string(),
                        args: vec![value],
                        effects: crate::mir::EffectMask::PURE.add(crate::mir::Effect::Io),
                    },
                    I::FutureSet { future, value } if rw_future => I::ExternCall {
                        dst: None,
                        iface_name: "env.future".to_string(),
                        method_name: "set".to_string(),
                        args: vec![future, value],
                        effects: crate::mir::EffectMask::PURE.add(crate::mir::Effect::Io),
                    },
                    I::Await { dst, future } if rw_future => I::ExternCall {
                        dst: Some(dst),
                        iface_name: "env.future".to_string(),
                        method_name: "await".to_string(),
                        args: vec![future],
                        effects: crate::mir::EffectMask::PURE.add(crate::mir::Effect::Io),
                    },
                    other => other,
                };
                new_spanned.push(SpannedInstruction {
                    inst: rewritten,
                    span,
                });
            }

            let terminator = block.terminator.take();
            let terminator_span = block.terminator_span.take();
            if let Some(term) = terminator {
                let span = terminator_span.unwrap_or_else(Span::unknown);
                let rewritten = match term {
                    I::TypeCheck {
                        dst,
                        value,
                        expected_type,
                    } => {
                        let ty = crate::mir::MirType::Box(expected_type);
                        stats.intrinsic_optimizations += 1;
                        I::TypeOp {
                            dst,
                            op: TypeOpKind::Check,
                            value,
                            ty,
                        }
                    }
                    I::Cast {
                        dst,
                        value,
                        target_type,
                    } => {
                        let ty = target_type;
                        stats.intrinsic_optimizations += 1;
                        I::TypeOp {
                            dst,
                            op: TypeOpKind::Cast,
                            value,
                            ty,
                        }
                    }
                    I::WeakNew { dst, box_val } => {
                        stats.intrinsic_optimizations += 1;
                        I::WeakRef {
                            dst,
                            op: WeakRefOp::New,
                            value: box_val,
                        }
                    }
                    I::WeakLoad { dst, weak_ref } => {
                        stats.intrinsic_optimizations += 1;
                        I::WeakRef {
                            dst,
                            op: WeakRefOp::Load,
                            value: weak_ref,
                        }
                    }
                    I::BarrierRead { ptr } => {
                        stats.intrinsic_optimizations += 1;
                        I::Barrier {
                            op: BarrierOp::Read,
                            ptr,
                        }
                    }
                    I::BarrierWrite { ptr } => {
                        stats.intrinsic_optimizations += 1;
                        I::Barrier {
                            op: BarrierOp::Write,
                            ptr,
                        }
                    }
                    I::Print { value, .. } => {
                        stats.intrinsic_optimizations += 1;
                        I::ExternCall {
                            dst: None,
                            iface_name: "env.console".to_string(),
                            method_name: "log".to_string(),
                            args: vec![value],
                            effects: crate::mir::EffectMask::PURE,
                        }
                    }
                    I::ArrayGet { dst, array, index } if array_to_boxcall => {
                        stats.intrinsic_optimizations += 1;
                        I::BoxCall {
                            dst: Some(dst),
                            box_val: array,
                            method: "get".to_string(),
                            method_id: None,
                            args: vec![index],
                            effects: crate::mir::EffectMask::READ,
                        }
                    }
                    I::ArraySet {
                        array,
                        index,
                        value,
                    } if array_to_boxcall => {
                        stats.intrinsic_optimizations += 1;
                        I::BoxCall {
                            dst: None,
                            box_val: array,
                            method: "set".to_string(),
                            method_id: None,
                            args: vec![index, value],
                            effects: crate::mir::EffectMask::WRITE,
                        }
                    }
                    other => other,
                };
                block.set_terminator_with_span(rewritten, span);
            }

            let (insts, spans): (Vec<_>, Vec<_>) =
                new_spanned.into_iter().map(|sp| (sp.inst, sp.span)).unzip();
            block.instructions = insts;
            block.instruction_spans = spans;
            block.effects = block
                .instructions
                .iter()
                .chain(block.terminator.iter())
                .fold(EffectMask::PURE, |mask, inst| mask | inst.effects());
        }
        if idemp_enabled() {
            let key = idemp_key(pass_name, &fname);
            idemp_mark(module, key);
        }
    }
    stats
}

pub fn normalize_ref_field_access(
    _opt: &mut MirOptimizer,
    module: &mut MirModule,
) -> OptimizationStats {
    use crate::mir::MirInstruction as I;
    let mut stats = OptimizationStats::new();
    let pass_name = "normalize.ref_field_access";
    let func_names: Vec<String> = module.functions.keys().cloned().collect();
    for fname in func_names {
        if idemp_enabled() {
            let key = idemp_key(pass_name, &fname);
            if idemp_already_done(module, &key) {
                continue;
            }
        }
        let function = match module.functions.get_mut(&fname) {
            Some(f) => f,
            None => continue,
        };
        for (_bb, block) in &mut function.blocks {
            let old_spanned: Vec<SpannedInstruction> = block.drain_spanned_instructions();
            let mut out: Vec<I> = Vec::with_capacity(old_spanned.len() + 2);
            let mut out_spans: Vec<Span> = Vec::with_capacity(old_spanned.len() + 2);
            for SpannedInstruction { inst, span } in old_spanned.into_iter() {
                match inst {
                    I::RefGet {
                        dst,
                        reference,
                        field,
                    } => {
                        let new_id = ValueId::new(function.next_value_id);
                        function.next_value_id += 1;
                        out.push(I::Const {
                            dst: new_id,
                            value: crate::mir::ConstValue::String(field),
                        });
                        out_spans.push(span);
                        out.push(I::BoxCall {
                            dst: Some(dst),
                            box_val: reference,
                            method: "getField".to_string(),
                            method_id: None,
                            args: vec![new_id],
                            effects: crate::mir::EffectMask::READ,
                        });
                        out_spans.push(span);
                        stats.intrinsic_optimizations += 1;
                    }
                    I::RefSet {
                        reference,
                        field,
                        value,
                    } => {
                        let new_id = ValueId::new(function.next_value_id);
                        function.next_value_id += 1;
                        out.push(I::Const {
                            dst: new_id,
                            value: crate::mir::ConstValue::String(field),
                        });
                        out_spans.push(span);
                        out.push(I::Barrier {
                            op: BarrierOp::Write,
                            ptr: reference,
                        });
                        out_spans.push(span);
                        out.push(I::BoxCall {
                            dst: None,
                            box_val: reference,
                            method: "setField".to_string(),
                            method_id: None,
                            args: vec![new_id, value],
                            effects: crate::mir::EffectMask::WRITE,
                        });
                        out_spans.push(span);
                        stats.intrinsic_optimizations += 1;
                    }
                    other => {
                        out.push(other);
                        out_spans.push(span);
                    }
                }
            }
            block.instructions = out;
            block.instruction_spans = out_spans;

            if let Some(term) = block.terminator.take() {
                let term_span = block.terminator_span.take().unwrap_or_else(Span::unknown);
                let rewritten = match term {
                    I::RefGet {
                        dst,
                        reference,
                        field,
                    } => {
                        let new_id = ValueId::new(function.next_value_id);
                        function.next_value_id += 1;
                        block.instructions.push(I::Const {
                            dst: new_id,
                            value: crate::mir::ConstValue::String(field),
                        });
                        block.instruction_spans.push(term_span);
                        I::BoxCall {
                            dst: Some(dst),
                            box_val: reference,
                            method: "getField".to_string(),
                            method_id: None,
                            args: vec![new_id],
                            effects: crate::mir::EffectMask::READ,
                        }
                    }
                    I::RefSet {
                        reference,
                        field,
                        value,
                    } => {
                        let new_id = ValueId::new(function.next_value_id);
                        function.next_value_id += 1;
                        block.instructions.push(I::Const {
                            dst: new_id,
                            value: crate::mir::ConstValue::String(field),
                        });
                        block.instruction_spans.push(term_span);
                        block.instructions.push(I::Barrier {
                            op: BarrierOp::Write,
                            ptr: reference,
                        });
                        block.instruction_spans.push(term_span);
                        I::BoxCall {
                            dst: None,
                            box_val: reference,
                            method: "setField".to_string(),
                            method_id: None,
                            args: vec![new_id, value],
                            effects: crate::mir::EffectMask::WRITE,
                        }
                    }
                    other => other,
                };
                block.set_terminator_with_span(rewritten, term_span);
            }

            block.effects = block
                .instructions
                .iter()
                .chain(block.terminator.iter())
                .fold(EffectMask::PURE, |mask, inst| mask | inst.effects());
        }
        if idemp_enabled() {
            let key = idemp_key(pass_name, &fname);
            idemp_mark(module, key);
        }
    }
    stats
}
