//! M11.5 — static-initialiser driver detection.
//!
//! MSVC's CRT static-init driver (`_initterm` / `_initterm_e` style)
//! is a tight loop that walks a function-pointer array between two
//! addresses, calling each non-null entry:
//!
//! ```text
//!   loop_top:
//!     cmpw[l]  rA, rB              ; compare cursor vs end
//!     beq      done
//!     lwz      rN, 0(rA)           ; load fn ptr
//!     cmpwi    rN, 0               ; null-skip (optional)
//!     beq      skip
//!     mtctr    rN
//!     bcctrl
//!   skip:
//!     addi     rA, rA, 4
//!     b        loop_top
//!   done:
//! ```
//!
//! Two static addresses (`rA` and `rB` at loop start) bracket the
//! function-pointer array. Detection strategy: scan every function for
//! the canonical pattern; when found, extract the array bounds and
//! emit one row in `function_pointer_arrays` with `kind='static_init'`.
//!
//! ### What this layer does
//!
//! - Walks each function looking for an `lwz; mtctr; bcctrl` sequence
//!   inside a loop bounded by a comparison against another constant.
//! - When the loop's cursor register is observed to be incremented by
//!   exactly 4 per iteration, classifies it as a static-init driver
//!   and records the (start, end) array bounds.
//!
//! ### What this layer does NOT do
//!
//! - No support for back-to-back drivers sharing a common loop trampoline.
//! - No detection of the M11 prologue-style heuristic; M11.5 is
//!   structure-grounded and replaces the prior heuristic where it fires.
//! - Does not handle CRT-style `_initterm_e` (the `_e` variant returns
//!   a status); detection works for both as long as the loop shape
//!   matches.
//!
//! Reference: Microsoft CRT `crt0.c::_initterm` source pattern.

use std::collections::{BTreeSet, HashMap, HashSet};

use crate::func::FuncAnalysis;
use crate::funcptr_arrays::FuncPtrArray;
use xenia_xex::pe::PeSection;

#[derive(Debug, Clone, Copy)]
pub struct StaticInitDriver {
    /// VA of the driver function (the one containing the loop).
    pub driver_function: u32,
    /// VA of the array start.
    pub array_start: u32,
    /// VA one-past-end of the array.
    pub array_end: u32,
    /// Detected length in slots.
    pub length: u32,
}

#[derive(Debug, Default)]
pub struct StaticInitResult {
    pub drivers: Vec<StaticInitDriver>,
    /// Newly-detected static-init arrays, ready to be merged into the
    /// `function_pointer_arrays` table with `kind='static_init'`.
    pub arrays: Vec<FuncPtrArray>,
}

const OP_ADDI:  u32 = 14;
const OP_ADDIS: u32 = 15;
const OP_BCCTR: u32 = 19;
const OP_LWZ:   u32 = 32;
const OP_X_FORM: u32 = 31;

#[derive(Debug, Clone, Copy)]
enum RegVal {
    Const(u32),
}

#[tracing::instrument(skip_all, fields(image_base = format_args!("{:#010x}", image_base)))]
pub fn analyze(
    pe: &[u8],
    image_base: u32,
    sections: &[PeSection],
    func_analysis: &FuncAnalysis,
    function_starts: &BTreeSet<u32>,
    labels: &HashMap<u32, String>,
) -> StaticInitResult {
    let started = std::time::Instant::now();
    let block_boundaries: HashSet<u32> = labels.keys().copied().collect();

    let mut drivers: Vec<StaticInitDriver> = Vec::new();

    for (&fn_start, fi) in &func_analysis.functions {
        if fi.is_saverestore { continue; }
        if let Some(d) = scan_function_for_driver(
            pe, image_base, fn_start, fi.end, &block_boundaries,
        ) {
            drivers.push(d);
        }
    }

    // Build arrays from the discovered drivers + section data.
    let mut arrays: Vec<FuncPtrArray> = Vec::new();
    for d in &drivers {
        if let Some(entries) = read_array(pe, image_base, sections, d.array_start, d.array_end, function_starts) {
            arrays.push(FuncPtrArray {
                address: d.array_start,
                length: entries.len() as u32,
                kind: "static_init",
                entries,
            });
        }
    }

    let elapsed_ms = started.elapsed().as_millis() as f64;
    metrics::histogram!("analysis.phase_ms", "phase" => "static_init").record(elapsed_ms);
    tracing::info!(
        drivers = drivers.len(),
        arrays = arrays.len(),
        elapsed_ms,
        "M11.5 static-init driver scan complete",
    );

    StaticInitResult { drivers, arrays }
}

/// Read the function-pointer array between [start, end) from .rdata/.data.
/// NULL entries are skipped (CRT _initterm explicitly tolerates them).
/// Non-function-start entries cause us to bail (the driver bounds were
/// likely misidentified).
fn read_array(
    pe: &[u8],
    image_base: u32,
    sections: &[PeSection],
    start: u32,
    end: u32,
    function_starts: &BTreeSet<u32>,
) -> Option<Vec<u32>> {
    if end <= start || (end - start) > 4096 { return None; }
    let _section = sections.iter().find(|s| {
        let lo = image_base + s.virtual_address;
        let hi = lo + s.virtual_size;
        start >= lo && end <= hi && (s.name == ".rdata" || s.name == ".data")
    })?;
    let mut entries = Vec::new();
    let mut p = start;
    while p < end {
        let off = p.wrapping_sub(image_base) as usize;
        if off + 4 > pe.len() { return None; }
        let v = u32::from_be_bytes([pe[off], pe[off + 1], pe[off + 2], pe[off + 3]]);
        if v != 0 {
            if !function_starts.contains(&v) { return None; }
            entries.push(v);
        }
        p = p.wrapping_add(4);
    }
    if entries.is_empty() { return None; }
    Some(entries)
}

/// Walk one function looking for the canonical static-init driver shape.
/// Returns Some when the loop's cursor register starts at a known constant
/// `rA`, terminates at another known constant `rB` via a compare, and
/// increments by 4 per iteration with an `lwz; mtctr; bcctrl` body.
fn scan_function_for_driver(
    pe: &[u8],
    image_base: u32,
    fn_start: u32,
    fn_end: u32,
    block_boundaries: &HashSet<u32>,
) -> Option<StaticInitDriver> {
    let mut reg: [Option<RegVal>; 32] = [None; 32];
    // Pattern features observed during the walk.
    let mut cursor_reg: Option<usize> = None;
    let mut cursor_init: Option<u32> = None;
    let mut end_reg: Option<usize> = None;
    let mut end_init: Option<u32> = None;
    let mut saw_lwz_through_cursor = false;
    let mut saw_mtctr = false;
    let mut saw_bcctrl = false;
    let mut saw_addi_4 = false;

    let mut pc = fn_start;
    while pc < fn_end {
        if pc != fn_start && block_boundaries.contains(&pc) {
            // Heuristic: when we cross a basic-block boundary that
            // is not the loop-top, accumulated state remains valid for
            // pattern-matching purposes — but we drop register Const
            // tracking to be safe.
            reg = [None; 32];
        }
        let off = pc.wrapping_sub(image_base) as usize;
        if off + 4 > pe.len() { break; }
        let instr = u32::from_be_bytes([pe[off], pe[off + 1], pe[off + 2], pe[off + 3]]);
        let op = instr >> 26;
        let rd = ((instr >> 21) & 0x1F) as usize;
        let ra = ((instr >> 16) & 0x1F) as usize;
        let simm = ((instr & 0xFFFF) as i16) as i32;
        let uimm = instr & 0xFFFF;

        match op {
            OP_ADDIS if ra == 0 => reg[rd] = Some(RegVal::Const(uimm << 16)),
            OP_ADDIS => {
                if let Some(RegVal::Const(b)) = reg[ra] {
                    reg[rd] = Some(RegVal::Const(b.wrapping_add(uimm << 16)));
                } else { reg[rd] = None; }
            }
            OP_ADDI if ra != 0 => {
                let prev = reg[ra];
                if let Some(RegVal::Const(b)) = prev {
                    let v = b.wrapping_add(simm as u32);
                    reg[rd] = Some(RegVal::Const(v));
                    // Was this an `addi r, r, 4`? Mark cursor-increment.
                    if rd == ra && simm == 4 {
                        if Some(rd) == cursor_reg {
                            saw_addi_4 = true;
                        }
                    } else if cursor_reg.is_none() {
                        // First time we see a known-constant register that
                        // *could* be the cursor — defer the choice until we
                        // see a load through it.
                        cursor_init = Some(v);
                        cursor_reg = Some(rd);
                    } else if end_reg.is_none() && Some(rd) != cursor_reg {
                        end_init = Some(v);
                        end_reg = Some(rd);
                    }
                } else { reg[rd] = None; }
            }
            OP_LWZ => {
                if ra != 0 && Some(ra) == cursor_reg {
                    saw_lwz_through_cursor = true;
                }
                reg[rd] = None;
            }
            OP_X_FORM => {
                let xo = (instr >> 1) & 0x3FF;
                if xo == 467 {
                    let spr = (((instr >> 11) & 0x1F) << 5) | ((instr >> 16) & 0x1F);
                    if spr == 9 && saw_lwz_through_cursor { saw_mtctr = true; }
                }
                if xo != 444 && xo != 467 { reg[rd] = None; }
            }
            OP_BCCTR => {
                let xo = (instr >> 1) & 0x3FF;
                let lk = (instr & 1) != 0;
                if xo == 528 && lk && saw_mtctr {
                    saw_bcctrl = true;
                }
            }
            18 => {
                if (instr & 1) != 0 {
                    for r in 0..=12 { reg[r] = None; }
                }
            }
            16 => {
                if (instr & 1) != 0 {
                    for r in 0..=12 { reg[r] = None; }
                }
            }
            _ => {}
        }
        pc = pc.wrapping_add(4);
    }

    // Validate that all four pattern features fired.
    if !(saw_lwz_through_cursor && saw_mtctr && saw_bcctrl && saw_addi_4) {
        return None;
    }
    let cursor_init = cursor_init?;
    let end_init = end_init?;
    if end_init <= cursor_init { return None; }
    if end_init - cursor_init > 4096 { return None; }

    Some(StaticInitDriver {
        driver_function: fn_start,
        array_start: cursor_init,
        array_end: end_init,
        length: (end_init - cursor_init) / 4,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::func::FuncInfo;
    use std::collections::BTreeMap;
    use xenia_xex::pe::PeSection;

    fn mk_section(name: &str, va: u32, size: u32) -> PeSection {
        PeSection {
            name: name.into(),
            virtual_address: va, virtual_size: size,
            raw_offset: va, raw_size: size,
            flags: 0x4000_0040,
        }
    }
    fn write_be(pe: &mut [u8], at: usize, v: u32) {
        pe[at..at + 4].copy_from_slice(&v.to_be_bytes());
    }

    #[test]
    fn detects_canonical_initterm_loop() {
        // Build a tiny driver that loops over a 3-entry array.
        let image_base = 0x82000000u32;
        let mut pe = vec![0u8; 0x4000];

        // Array at .rdata + 0x800: 3 function pointers.
        let arr_va_lo = 0x800u32;
        let fns = [image_base + 0x2000, image_base + 0x2010, image_base + 0x2020];
        for (i, p) in fns.iter().enumerate() {
            write_be(&mut pe, arr_va_lo as usize + i * 4, *p);
        }
        let array_start = image_base + arr_va_lo;
        let array_end = array_start + 12;

        // Driver function at 0x82001000:
        //   lis  r3, hi(array_start)
        //   addi r3, r3, lo(array_start)
        //   lis  r4, hi(array_end)
        //   addi r4, r4, lo(array_end)
        //   lwz  r5, 0(r3)
        //   mtctr r5
        //   bcctrl
        //   addi r3, r3, 4
        //   blr
        let driver = 0x82001000u32;
        let off = (driver - image_base) as usize;
        let lis_r3 = (15u32 << 26) | (3 << 21) | ((array_start >> 16) as u32);
        let addi_r3 = (14u32 << 26) | (3 << 21) | (3 << 16) | ((array_start as u16) as u32);
        let lis_r4 = (15u32 << 26) | (4 << 21) | ((array_end >> 16) as u32);
        let addi_r4 = (14u32 << 26) | (4 << 21) | (4 << 16) | ((array_end as u16) as u32);
        let lwz = (32u32 << 26) | (5 << 21) | (3 << 16);
        let mtctr = (31u32 << 26) | (5 << 21) | (9 << 16) | (467 << 1);
        let bcctrl = (19u32 << 26) | (20 << 21) | (528 << 1) | 1;
        let addi_inc = (14u32 << 26) | (3 << 21) | (3 << 16) | 4;
        let blr = (19u32 << 26) | (20 << 21) | (16 << 1);
        for (i, w) in [lis_r3, addi_r3, lis_r4, addi_r4, lwz, mtctr, bcctrl, addi_inc, blr].iter().enumerate() {
            write_be(&mut pe, off + i * 4, *w);
        }

        let mut functions: BTreeMap<u32, FuncInfo> = BTreeMap::new();
        functions.insert(driver, FuncInfo {
            start: driver, end: driver + 0x40, frame_size: 0, saved_gprs: 0,
            is_leaf: false, is_saverestore: false,
            pdata_validated: false, pdata_length: None, has_eh: false,
        });
        let fa = FuncAnalysis {
            functions, save_gpr_base: None, restore_gpr_base: None, pdata_entries: Vec::new(),
        };

        let sections = vec![mk_section(".rdata", 0x800, 0x100)];
        let mut starts = BTreeSet::new();
        for &p in &fns { starts.insert(p); }
        let labels: HashMap<u32, String> = HashMap::new();

        let r = analyze(&pe, image_base, &sections, &fa, &starts, &labels);

        assert_eq!(r.drivers.len(), 1, "should detect one driver");
        let d = &r.drivers[0];
        assert_eq!(d.driver_function, driver);
        assert_eq!(d.array_start, array_start);
        assert_eq!(d.array_end, array_end);
        assert_eq!(d.length, 3);

        assert_eq!(r.arrays.len(), 1);
        assert_eq!(r.arrays[0].kind, "static_init");
        assert_eq!(r.arrays[0].entries.len(), 3);
    }

    #[test]
    fn rejects_function_without_pattern() {
        let image_base = 0x82000000u32;
        let mut pe = vec![0u8; 0x4000];
        let driver = 0x82001000u32;
        // Just a blr — no driver pattern.
        let blr = (19u32 << 26) | (20 << 21) | (16 << 1);
        write_be(&mut pe, (driver - image_base) as usize, blr);

        let mut functions: BTreeMap<u32, FuncInfo> = BTreeMap::new();
        functions.insert(driver, FuncInfo {
            start: driver, end: driver + 0x40, frame_size: 0, saved_gprs: 0,
            is_leaf: true, is_saverestore: false,
            pdata_validated: false, pdata_length: None, has_eh: false,
        });
        let fa = FuncAnalysis {
            functions, save_gpr_base: None, restore_gpr_base: None, pdata_entries: Vec::new(),
        };
        let sections = vec![mk_section(".rdata", 0x800, 0x100)];
        let starts: BTreeSet<u32> = BTreeSet::new();
        let labels: HashMap<u32, String> = HashMap::new();
        let r = analyze(&pe, image_base, &sections, &fa, &starts, &labels);
        assert_eq!(r.drivers.len(), 0);
    }
}