//! MSVC RTTI recovery — the authoritative source of C++ class identity.
//!
//! [`crate::vtables`] finds vtables *bottom-up*, by looking for runs of words
//! that happen to be function entries, and only then tries the RTTI walk. That
//! misses every table whose head holds a null / pure-virtual / thunk slot, and
//! it cannot see a class that has no such run at all. This module works
//! *top-down* from the RTTI structures the linker emitted, which is exact:
//! a `CompleteObjectLocator` names its class, and the word that points at a
//! COL is by definition `vftable[-1]`.
//!
//! ## Structure layout (32-bit MSVC, big-endian on Xbox 360)
//!
//! ```text
//! TypeDescriptor           (in .data — it is written at startup)
//!   +0  void*  pVFTable    -> type_info's own vftable (identical for all TDs)
//!   +4  void*  spare
//!   +8  char   name[]      -> ".?AVFoo@Bar@@", NUL-terminated
//!
//! RTTICompleteObjectLocator (in .rdata)
//!   +0  u32    signature   -> 0 for 32-bit images
//!   +4  u32    offset      -> this-offset of the subobject this vftable serves
//!   +8  u32    cdOffset    -> constructor-displacement offset
//!   +12 TypeDescriptor*
//!   +16 RTTIClassHierarchyDescriptor*
//!
//! RTTIClassHierarchyDescriptor (in .rdata)
//!   +0  u32    signature
//!   +4  u32    attributes  -> bit 0 = multiple inheritance, bit 1 = virtual
//!   +8  u32    numBaseClasses
//!   +12 RTTIBaseClassDescriptor** pBaseClassArray
//!
//! RTTIBaseClassDescriptor  (in .rdata)
//!   +0  TypeDescriptor*
//!   +4  u32    numContainedBases
//!   +8  i32    PMD.mdisp   -> member displacement
//!   +12 i32    PMD.pdisp   -> vbtable displacement (-1 = not virtual)
//!   +16 i32    PMD.vdisp   -> displacement inside the vbtable
//!   +20 u32    attributes
//! ```
//!
//! A vtable is located at `col_ref + 4` for every word `col_ref` whose value is
//! a validated COL address. `offset` distinguishes the primary vftable
//! (`offset == 0`) from the extra vftables a multiply-inheriting class emits
//! for its secondary base subobjects — those are linked to the same class
//! rather than being reported as unrelated tables.
//!
//! ## Limits
//!
//! - Only statically-emitted RTTI is seen; a class whose RTTI the linker
//!   stripped stays anonymous and is left to [`crate::vtables`].
//! - Vtable *length* is measured by walking forward from `vftable[0]` while the
//!   words are plausible method pointers, stopping at the next COL reference or
//!   at a known label — the linker does not record it.

use std::collections::{BTreeMap, BTreeSet};

use sylpheed_xex::pe::PeSection;

use crate::demangle;

/// One `TypeDescriptor`: the mangled class name the compiler emitted.
#[derive(Debug, Clone)]
pub struct TypeDescriptor {
    /// VA of the descriptor (i.e. of its `pVFTable` word).
    pub address: u32,
    /// Raw decorated name, e.g. `.?AVSilph@silph@@`.
    pub mangled_name: String,
    /// Readable form, e.g. `silph::Silph`. Falls back to `mangled_name`.
    pub demangled_name: String,
}

/// One `RTTICompleteObjectLocator` and the vtable it labels.
#[derive(Debug, Clone)]
pub struct CompleteObjectLocator {
    pub address: u32,
    /// `this`-offset of the subobject whose vftable this is. 0 = primary.
    pub offset: u32,
    pub cd_offset: u32,
    pub type_descriptor: u32,
    pub class_hierarchy: u32,
    /// VA of `vftable[0]`, when a word pointing at this COL was found.
    pub vtable_address: Option<u32>,
}

/// One entry of a class's `RTTIBaseClassArray`, in linearised order.
#[derive(Debug, Clone)]
pub struct BaseClass {
    /// VA of the deriving class's `RTTIClassHierarchyDescriptor`.
    pub class_hierarchy: u32,
    /// Position in the base-class array (index 0 is the class itself).
    pub index: u32,
    pub type_descriptor: u32,
    pub name: String,
    pub num_contained_bases: u32,
    pub mdisp: i32,
    pub pdisp: i32,
    pub vdisp: i32,
    pub attributes: u32,
}

/// Everything the RTTI walk recovered.
#[derive(Debug, Default)]
pub struct RttiResult {
    pub type_descriptors: Vec<TypeDescriptor>,
    pub locators: Vec<CompleteObjectLocator>,
    pub base_classes: Vec<BaseClass>,
    /// `vftable[0]` VA → the COL that labels it.
    pub vtable_to_locator: BTreeMap<u32, u32>,
}

impl RttiResult {
    /// Vtable base VAs the walk proved exist — the anchor set
    /// [`crate::vtables`] should treat as authoritative.
    pub fn vtable_anchors(&self) -> BTreeSet<u32> {
        self.vtable_to_locator.keys().copied().collect()
    }

    /// `vftable[0]` VA → `(demangled class name, subobject offset)`.
    pub fn vtable_class_names(&self) -> BTreeMap<u32, (String, u32)> {
        let td: BTreeMap<u32, &TypeDescriptor> =
            self.type_descriptors.iter().map(|t| (t.address, t)).collect();
        let mut out = BTreeMap::new();
        for col in &self.locators {
            if let (Some(vt), Some(t)) = (col.vtable_address, td.get(&col.type_descriptor)) {
                out.insert(vt, (t.demangled_name.clone(), col.offset));
            }
        }
        out
    }
}

// ── Scan ───────────────────────────────────────────────────────────────────

/// Walk the image's RTTI. `sections` must be the full PE section list.
#[tracing::instrument(skip_all, fields(image_base = format_args!("{:#010x}", image_base)))]
pub fn analyze(pe: &[u8], image_base: u32, sections: &[PeSection]) -> RttiResult {
    let started = std::time::Instant::now();

    let read = |va: u32| -> Option<u32> {
        let off = va.wrapping_sub(image_base) as usize;
        if off.checked_add(4)? > pe.len() { return None; }
        Some(u32::from_be_bytes([pe[off], pe[off + 1], pe[off + 2], pe[off + 3]]))
    };

    // Byte ranges actually backed by file data (a section's tail beyond
    // `raw_size` is BSS: reading it yields zeros, never a real structure).
    let backed = |s: &PeSection| -> (u32, u32) {
        let start = image_base + s.virtual_address;
        let len = s.virtual_size.min(s.raw_size);
        (start, start + len)
    };
    let ranges: Vec<(String, u32, u32)> = sections
        .iter()
        .map(|s| { let (a, b) = backed(s); (s.name.clone(), a, b) })
        .collect();
    let range_of = |name: &str| -> Option<(u32, u32)> {
        ranges.iter().find(|(n, _, _)| n == name).map(|&(_, a, b)| (a, b))
    };

    // 1. TypeDescriptors. The decorated name lives at descriptor+8 and always
    //    starts with ".?A". MSVC places these in writable data.
    let mut type_descriptors: Vec<TypeDescriptor> = Vec::new();
    let mut td_addrs: BTreeSet<u32> = BTreeSet::new();
    for (name, start, end) in &ranges {
        if !matches!(name.as_str(), ".data" | ".rdata") { continue; }
        let s = (*start).wrapping_sub(image_base) as usize;
        let e = (*end).wrapping_sub(image_base) as usize;
        if e > pe.len() || s >= e { continue; }
        let bytes = &pe[s..e];
        let mut i = 0usize;
        while i + 3 < bytes.len() {
            if &bytes[i..i + 3] != b".?A" { i += 1; continue; }
            let name_va = start.wrapping_add(i as u32);
            // The descriptor head sits 8 bytes before the name.
            let Some(td_va) = name_va.checked_sub(8) else { i += 1; continue };
            if td_va < *start { i += 1; continue; }
            let Some(decorated) = read_cstr(bytes, i, 512) else { i += 1; continue };
            i += decorated.len() + 1;
            if td_addrs.insert(td_va) {
                type_descriptors.push(TypeDescriptor {
                    address: td_va,
                    demangled_name: demangle::demangle_type_descriptor(&decorated)
                        .unwrap_or_else(|| decorated.clone()),
                    mangled_name: decorated,
                });
            }
        }
    }

    // 2. CompleteObjectLocators. Scan read-only data on a 4-byte grid for the
    //    5-word shape whose `pTypeDescriptor` hits a descriptor we just found
    //    and whose `pClassDescriptor` points back into read-only data.
    let rdata = range_of(".rdata");
    let mut locators: Vec<CompleteObjectLocator> = Vec::new();
    let mut col_addrs: BTreeSet<u32> = BTreeSet::new();
    if let Some((rd_start, rd_end)) = rdata {
        let mut va = rd_start;
        while va + 20 <= rd_end {
            let (Some(sig), Some(off), Some(cd), Some(ptd), Some(pchd)) = (
                read(va), read(va + 4), read(va + 8), read(va + 12), read(va + 16),
            ) else { break };
            if sig == 0 && td_addrs.contains(&ptd) && pchd >= rd_start && pchd < rd_end {
                col_addrs.insert(va);
                locators.push(CompleteObjectLocator {
                    address: va,
                    offset: off,
                    cd_offset: cd,
                    type_descriptor: ptd,
                    class_hierarchy: pchd,
                    vtable_address: None,
                });
            }
            va += 4;
        }
    }

    // 3. `vftable[-1]` sites: any word in initialised data whose value is a COL.
    let mut vtable_to_locator: BTreeMap<u32, u32> = BTreeMap::new();
    for (name, start, end) in &ranges {
        if !matches!(name.as_str(), ".rdata" | ".data") { continue; }
        let mut va = *start;
        while va + 4 <= *end {
            if let Some(w) = read(va)
                && col_addrs.contains(&w)
            {
                vtable_to_locator.insert(va + 4, w);
            }
            va += 4;
        }
    }
    let locator_to_vtable: BTreeMap<u32, u32> =
        vtable_to_locator.iter().map(|(&vt, &col)| (col, vt)).collect();
    for col in &mut locators {
        col.vtable_address = locator_to_vtable.get(&col.address).copied();
    }

    // 4. Class hierarchies: for each distinct CHD, read its base-class array.
    let td_by_addr: BTreeMap<u32, &TypeDescriptor> =
        type_descriptors.iter().map(|t| (t.address, t)).collect();
    let mut base_classes: Vec<BaseClass> = Vec::new();
    let chds: BTreeSet<u32> = locators.iter().map(|c| c.class_hierarchy).collect();
    if let Some((rd_start, rd_end)) = rdata {
        for chd in chds {
            let (Some(n_bases), Some(p_array)) = (read(chd + 8), read(chd + 12)) else { continue };
            // A malformed or misidentified descriptor would blow the scan up;
            // real hierarchies are small.
            if n_bases == 0 || n_bases > 64 { continue; }
            if p_array < rd_start || p_array >= rd_end { continue; }
            for i in 0..n_bases {
                let Some(bcd) = read(p_array + i * 4) else { break };
                if bcd < rd_start || bcd >= rd_end { break; }
                let (Some(ptd), Some(ncb), Some(md), Some(pd), Some(vd), Some(attr)) = (
                    read(bcd), read(bcd + 4), read(bcd + 8),
                    read(bcd + 12), read(bcd + 16), read(bcd + 20),
                ) else { break };
                let Some(td) = td_by_addr.get(&ptd) else { break };
                base_classes.push(BaseClass {
                    class_hierarchy: chd,
                    index: i,
                    type_descriptor: ptd,
                    name: td.demangled_name.clone(),
                    num_contained_bases: ncb,
                    mdisp: md as i32,
                    pdisp: pd as i32,
                    vdisp: vd as i32,
                    attributes: attr,
                });
            }
        }
    }

    let elapsed_ms = started.elapsed().as_millis() as f64;
    metrics::histogram!("analysis.phase_ms", "phase" => "rtti").record(elapsed_ms);
    tracing::info!(
        type_descriptors = type_descriptors.len(),
        locators = locators.len(),
        vtables = vtable_to_locator.len(),
        base_class_records = base_classes.len(),
        elapsed_ms,
        "RTTI walk complete",
    );

    RttiResult { type_descriptors, locators, base_classes, vtable_to_locator }
}

/// Read a NUL-terminated ASCII string starting at `off` in `bytes`.
fn read_cstr(bytes: &[u8], off: usize, max: usize) -> Option<String> {
    let end = (off + max).min(bytes.len());
    let slice = &bytes[off..end];
    let nul = slice.iter().position(|&b| b == 0)?;
    let s = &slice[..nul];
    if s.is_empty() || !s.iter().all(|&b| (0x20..0x7F).contains(&b)) {
        return None;
    }
    Some(String::from_utf8_lossy(s).into_owned())
}

#[cfg(test)]
mod tests {
    use super::*;

    const BASE: u32 = 0x8200_0000;
    const RDATA_RVA: u32 = 0x1000;
    const DATA_RVA: u32 = 0x2000;
    const SEC_SIZE: u32 = 0x1000;

    fn sections() -> Vec<PeSection> {
        vec![
            PeSection {
                name: ".rdata".into(),
                virtual_address: RDATA_RVA, virtual_size: SEC_SIZE,
                raw_offset: RDATA_RVA, raw_size: SEC_SIZE,
                flags: 0x4000_0040,
            },
            PeSection {
                name: ".data".into(),
                virtual_address: DATA_RVA, virtual_size: SEC_SIZE,
                raw_offset: DATA_RVA, raw_size: SEC_SIZE,
                flags: 0xC000_0040,
            },
        ]
    }

    struct Image(Vec<u8>);
    impl Image {
        fn new() -> Self { Image(vec![0u8; (DATA_RVA + SEC_SIZE) as usize]) }
        fn put_u32(&mut self, va: u32, v: u32) {
            let o = (va - BASE) as usize;
            self.0[o..o + 4].copy_from_slice(&v.to_be_bytes());
        }
        fn put_str(&mut self, va: u32, s: &str) {
            let o = (va - BASE) as usize;
            self.0[o..o + s.len()].copy_from_slice(s.as_bytes());
            self.0[o + s.len()] = 0;
        }
    }

    /// Lay down one class: TypeDescriptor in .data, COL + CHD + BCD in .rdata,
    /// and the `vftable[-1]` word that points at the COL.
    #[allow(clippy::too_many_arguments)]
    fn emit_class(
        img: &mut Image, td: u32, name: &str,
        col: u32, offset: u32, chd: u32, bcd_array: u32, bcd: u32, base_name_td: Option<u32>,
        vtable_minus_one: u32,
    ) {
        img.put_u32(td, 0xDEAD_BEEF); // type_info vftable — value is irrelevant
        img.put_str(td + 8, name);

        img.put_u32(col, 0);        // signature
        img.put_u32(col + 4, offset);
        img.put_u32(col + 8, 0);    // cdOffset
        img.put_u32(col + 12, td);
        img.put_u32(col + 16, chd);

        let n_bases = if base_name_td.is_some() { 2 } else { 1 };
        img.put_u32(chd, 0);
        img.put_u32(chd + 4, 0);
        img.put_u32(chd + 8, n_bases);
        img.put_u32(chd + 12, bcd_array);

        // Base-class array: entry 0 is the class itself.
        img.put_u32(bcd_array, bcd);
        img.put_u32(bcd, td);
        img.put_u32(bcd + 4, n_bases - 1);
        img.put_u32(bcd + 8, 0);            // mdisp
        img.put_u32(bcd + 12, u32::MAX);    // pdisp = -1
        img.put_u32(bcd + 16, 0);           // vdisp
        img.put_u32(bcd + 20, 0x40);        // attributes
        if let Some(base_td) = base_name_td {
            let bcd2 = bcd + 24;
            img.put_u32(bcd_array + 4, bcd2);
            img.put_u32(bcd2, base_td);
            img.put_u32(bcd2 + 4, 0);
            img.put_u32(bcd2 + 8, 4);        // mdisp = 4
            img.put_u32(bcd2 + 12, u32::MAX);
            img.put_u32(bcd2 + 16, 0);
            img.put_u32(bcd2 + 20, 0);
        }

        img.put_u32(vtable_minus_one, col);
    }

    #[test]
    fn recovers_class_name_vtable_and_bases() {
        let mut img = Image::new();
        let rd = BASE + RDATA_RVA;
        let da = BASE + DATA_RVA;

        // Base class Foo, then Derived : Foo.
        emit_class(&mut img, da + 0x100, ".?AVFoo@ns@@",
                   rd + 0x100, 0, rd + 0x200, rd + 0x280, rd + 0x300, None,
                   rd + 0x000);
        emit_class(&mut img, da + 0x200, ".?AVDerived@ns@@",
                   rd + 0x400, 0, rd + 0x500, rd + 0x580, rd + 0x600, Some(da + 0x100),
                   rd + 0x040);

        let r = analyze(&img.0, BASE, &sections());

        assert_eq!(r.type_descriptors.len(), 2);
        let derived = r.type_descriptors.iter()
            .find(|t| t.mangled_name.contains("Derived")).unwrap();
        assert_eq!(derived.demangled_name, "ns::Derived");

        assert_eq!(r.locators.len(), 2);
        // vftable[0] is one word past the word holding the COL pointer.
        assert_eq!(r.vtable_to_locator.get(&(rd + 0x044)), Some(&(rd + 0x400)));
        assert!(r.vtable_anchors().contains(&(rd + 0x004)));

        let names = r.vtable_class_names();
        assert_eq!(names.get(&(rd + 0x044)), Some(&("ns::Derived".to_string(), 0)));

        // Derived's hierarchy lists itself at index 0 and Foo at index 1.
        let mut bases: Vec<_> = r.base_classes.iter()
            .filter(|b| b.class_hierarchy == rd + 0x500)
            .collect();
        bases.sort_by_key(|b| b.index);
        assert_eq!(bases.len(), 2);
        assert_eq!(bases[1].name, "ns::Foo");
        assert_eq!(bases[1].mdisp, 4);
        assert_eq!(bases[1].pdisp, -1);
    }

    #[test]
    fn secondary_base_vftable_keeps_its_subobject_offset() {
        let mut img = Image::new();
        let rd = BASE + RDATA_RVA;
        let da = BASE + DATA_RVA;
        emit_class(&mut img, da + 0x100, ".?AVMulti@@",
                   rd + 0x100, 0x8, rd + 0x200, rd + 0x280, rd + 0x300, None,
                   rd + 0x000);

        let r = analyze(&img.0, BASE, &sections());
        let names = r.vtable_class_names();
        assert_eq!(names.get(&(rd + 0x004)), Some(&("Multi".to_string(), 0x8)));
    }

    #[test]
    fn ignores_data_that_merely_looks_like_a_locator() {
        // A 5-word run with signature 0 but a `pTypeDescriptor` that hits no
        // descriptor must not be reported.
        let mut img = Image::new();
        let rd = BASE + RDATA_RVA;
        img.put_u32(rd + 0x100, 0);
        img.put_u32(rd + 0x104, 0);
        img.put_u32(rd + 0x108, 0);
        img.put_u32(rd + 0x10C, BASE + DATA_RVA + 0x900); // no TD there
        img.put_u32(rd + 0x110, rd + 0x200);

        let r = analyze(&img.0, BASE, &sections());
        assert!(r.locators.is_empty());
        assert!(r.type_descriptors.is_empty());
    }
}