//! Offline Itanium C++ RTTI: locate a class's vtable in an ELF `.so` and read its slot array.
//!
//! Chain (Itanium ABI, LP64): the class name is stored length-prefixed+mangled (e.g.
//! "11CBaseEntity") as a `_ZTS` string in `.rodata`; the `_ZTI` typeinfo points to that string
//! at +8; the `_ZTV` vtable points to the typeinfo at +8, with `offset-to-top` at +0, so virtual
//! slots start at vtable+16. Those slot pointers live in `.data.rel.ro` and are supplied by
//! relocations, which `CodeImage::read_ptr` already resolves.
//!
//! This is the ELF/Itanium half; a Windows fork would add an MSVC-RTTI sibling behind the same
//! `find_vtable` shape (COL at vftable-8, TypeDescriptor `.?AV<name>@@`).

use crate::elf::{CodeImage, KindTag};
use std::collections::HashSet;

pub struct VTable {
    pub slot0: u64,      // vaddr of virtual slot index 0
    pub slots: Vec<u64>, // function vaddrs; gamedata offset of a method == its index here
}

/// One vtable discovered by the whole-binary sweep — the class inventory row.
pub struct ClassVtable {
    pub mangled: String,       // the raw `_ZTS` type name, e.g. "11CBaseEntity"
    pub name: String,          // demangled, e.g. "CBaseEntity"
    pub vtable_va: u64,        // vaddr of slot index 0
    pub offset_to_top: i64, // 0 for the primary (complete-object) vtable; <0 for sub-object tables
    pub typeinfo: u64,      // vaddr of the Itanium typeinfo struct
    pub slots: Vec<u64>,    // method vaddrs; a method's gamedata offset == its index here
    pub bases: Vec<BaseClass>, // direct base classes (the is-a graph edges)
}

/// A direct base class of a type, from its Itanium typeinfo.
pub struct BaseClass {
    pub name: String,       // demangled base class name
    pub offset: i64, // this-pointer adjustment to the base subobject (0 for the primary base)
    pub virtual_base: bool, // true if inherited virtually
}

/// The three Itanium `type_info` "kind" vtables (their in-object `+16` slot0 pointers). libc++abi is
/// statically bundled in CS2 libraries, so these resolve as WEAK symbols and let us classify each
/// typeinfo *exactly* — no heuristic guess of `__class` vs `__si` vs `__vmi`.
struct RttiKinds {
    class: u64, // __class_type_info  — no bases
    si: u64,    // __si_class_type_info — single public base at offset 0
    vmi: u64,   // __vmi_class_type_info — multiple / virtual / non-public bases
}

impl RttiKinds {
    fn detect(img: &CodeImage) -> Self {
        let k = |n: &str| img.symbol_addr(n).map_or(0, |a| a.wrapping_add(16));
        Self {
            class: k("_ZTVN10__cxxabiv117__class_type_infoE"),
            si: k("_ZTVN10__cxxabiv120__si_class_type_infoE"),
            vmi: k("_ZTVN10__cxxabiv121__vmi_class_type_infoE"),
        }
    }
    /// Is `p` (a typeinfo's `+0` field) one of the three kind vtables? When the kind symbols are
    /// stripped (all zero) we can't tell, so accept any pointer the caller already range-checked.
    fn is_kind(&self, p: u64) -> bool {
        if self.class == 0 && self.si == 0 && self.vmi == 0 {
            return true;
        }
        p == self.class || p == self.si || p == self.vmi
    }
}

/// Itanium length-prefixed name for a flat class, e.g. `CBaseEntity` -> `11CBaseEntity`.
/// (Namespaced/templated names need full mangling; our targets are flat class names.)
fn mangle(class: &str) -> String {
    format!("{}{}", class.len(), class)
}

/// Find the class's primary (complete-object) vtable and read its function-pointer slots.
pub fn find_vtable(img: &CodeImage, class: &str, max_slots: usize) -> Option<VTable> {
    let mut candidates: Vec<u64> = Vec::new();

    // Fast path: an exported `_ZTV` symbol (uncommon for gameplay classes, but cheap).
    if let Some(ztv) = img.symbol_addr(&format!("_ZTV{}", mangle(class))) {
        candidates.push(ztv.wrapping_add(16));
    }

    // General path: name string -> typeinfo (points to name at +8) -> vtable (points to TI at +8).
    let mut needle = mangle(class).into_bytes();
    needle.push(0);
    for name_str in img.find_bytes(&needle) {
        for &ti_name_slot in img.ptrs_to(name_str) {
            if ti_name_slot < 8 {
                continue;
            }
            let typeinfo = ti_name_slot - 8;
            for &vt_ti_slot in img.ptrs_to(typeinfo) {
                candidates.push(vt_ti_slot.wrapping_add(8));
            }
        }
    }
    candidates.sort_unstable();
    candidates.dedup();

    for slot0 in candidates {
        // primary vtable has offset-to-top == 0 at slot0-16; filters typeinfo base-class lists
        if img.read_ptr(slot0.wrapping_sub(16)) != Some(0) {
            continue;
        }
        let slots = read_slots(img, slot0, max_slots);
        // Higher floor than `enumerate_vtables` (which admits `>= 2`): a 2-slot stub is too thin to trust
        // as the TARGET's real vtable when matching by name. A class whose primary vtable has exactly 2
        // code slots is still catalogued in the model but not re-located here, so its offsets flag
        // unresolved — a missed derivation for a rare class, never a wrong value.
        if slots.len() >= 3 {
            return Some(VTable { slot0, slots });
        }
    }
    None
}

/// Consecutive slot pointers that land in executable code; stops at the first that doesn't.
///
/// Returning exactly `max` slots is AMBIGUOUS — the vtable may genuinely end there, or may continue past
/// the cap with the tail silently dropped. Callers that care (the ones recording slot counts into the
/// model) should compare `len() == max` and warn; see `GameProfile::max_vtable_slots`.
fn read_slots(img: &CodeImage, slot0: u64, max: usize) -> Vec<u64> {
    let mut out = Vec::new();
    for i in 0..max {
        match img.read_ptr(slot0.wrapping_add((i as u64).wrapping_mul(8))) {
            Some(v) if img.is_code(v) => out.push(v),
            _ => break,
        }
    }
    out
}

/// Demangle an Itanium *type* name (the bare `_ZTS` payload, e.g. "11CBaseEntity") to a readable
/// class name. cpp_demangle wants a whole symbol, so we re-attach the `_ZTS` prefix and strip the
/// "typeinfo name for " decoration it produces. Falls back to the mangled form.
fn demangle_type(mangled: &str) -> String {
    let sym = format!("_ZTS{mangled}");
    cpp_demangle::Symbol::new(sym.as_bytes())
        .ok()
        .and_then(|s| s.demangle().ok())
        .map(|d| {
            d.strip_prefix("typeinfo name for ")
                .unwrap_or(&d)
                .to_string()
        })
        .unwrap_or_else(|| mangled.to_string())
}

/// If `ti` addresses a valid Itanium typeinfo, return its `(mangled, demangled)` class name.
/// A typeinfo is `[kind_vtable_ptr][name_ptr][ base-class data … ]`: `+0` points at one of the
/// C++ runtime's type_info-kind vtables, `+8` at the `_ZTS` name string.
fn typeinfo_name(img: &CodeImage, ti: u64, kinds: &RttiKinds) -> Option<(String, String)> {
    // +0 must be one of the three kind vtables. Prefer the symbol-name-derived tag (the only signal that
    // survives a DYNAMICALLY-linked C++ runtime, where the three kinds all resolve to the same offline
    // value); else fall back to the in-image value check (statically-linked / stripped builds).
    if img.kind_at(ti).is_none() {
        let kind = img.read_ptr(ti)?;
        if kind == 0 || !img.contains(kind) || !kinds.is_kind(kind) {
            return None;
        }
    }
    let name_ptr = img.read_ptr(ti.wrapping_add(8))?;
    let mangled = img.read_c_string(name_ptr)?;
    // Itanium type names start with a length digit (flat class) or a mangling sigil.
    let c0 = *mangled.as_bytes().first()?;
    if !(c0.is_ascii_digit() || matches!(c0, b'N' | b'I' | b'P' | b'K' | b'S')) {
        return None;
    }
    Some((mangled.clone(), demangle_type(&mangled)))
}

/// Direct base classes of the typeinfo at `ti`, dispatched on its exact Itanium kind.
fn typeinfo_bases(img: &CodeImage, ti: u64, kinds: &RttiKinds) -> Vec<BaseClass> {
    // Classify the kind: prefer the symbol-name tag (dynamically-linked runtime), else compare the resolved
    // +0 pointer to the in-image kind vtables (statically-linked). Without the tag, an old build can't tell
    // __si from __vmi at all, and the base graph would silently come back empty.
    let tag = img.kind_at(ti).or_else(|| {
        let kind = img.read_ptr(ti).unwrap_or(0);
        if kind == 0 {
            None
        } else if kind == kinds.si {
            Some(KindTag::Si)
        } else if kind == kinds.vmi {
            Some(KindTag::Vmi)
        } else {
            None
        }
    });
    match tag {
        Some(KindTag::Si) => {
            // __si_class_type_info: one public, non-virtual base at offset 0; its typeinfo ptr at +16.
            if let Some(bp) = img.read_ptr(ti.wrapping_add(16))
                && let Some((_, name)) = typeinfo_name(img, bp, kinds)
            {
                return vec![BaseClass {
                    name,
                    offset: 0,
                    virtual_base: false,
                }];
            }
            Vec::new()
        }
        Some(KindTag::Vmi) => {
            // __vmi_class_type_info: flags@+16, base_count@+20, then 16-byte {typeinfo_ptr, offset_flags}.
            let Some(count) = img.read_u32(ti + 20) else {
                return Vec::new();
            };
            if count == 0 || count > 128 {
                return Vec::new();
            }
            let mut bases = Vec::new();
            for i in 0..count as u64 {
                let e = ti.wrapping_add(24).wrapping_add(i.wrapping_mul(16));
                let Some(bp) = img.read_ptr(e) else {
                    break;
                };
                if let Some((_, name)) = typeinfo_name(img, bp, kinds) {
                    let of = img.read_i64(e.wrapping_add(8)).unwrap_or(0);
                    bases.push(BaseClass {
                        name,
                        offset: of >> 8, // Itanium: high bits = this-pointer adjustment
                        virtual_base: of & 0x1 != 0, // low byte: 0x1 = virtual, 0x2 = public
                    });
                }
            }
            bases
        }
        _ => Vec::new(), // __class_type_info (no bases) or a kind we can't classify
    }
}

/// Enumerate EVERY class vtable in the image via Itanium RTTI — the whole-binary class inventory.
///
/// Reloc-driven (not a raw byte sweep): each vtable's typeinfo field at `vtable-8` is a relocation,
/// so we walk the reloc map, keep slots that point at a valid typeinfo, and recover the vtable just
/// above. Every pointer is read through the `.rela.dyn`-resolved `read_ptr`, so `.data.rel.ro` slots
/// (zero on disk) come back as their true as-loaded values.
pub fn enumerate_vtables(img: &CodeImage, max_slots: usize) -> Vec<ClassVtable> {
    let kinds = RttiKinds::detect(img);
    let mut out = Vec::new();
    let mut seen = HashSet::new();
    for (slot, val) in img.reloc_slots() {
        if slot < 8 {
            continue;
        }
        let Some((mangled, name)) = typeinfo_name(img, val, &kinds) else {
            continue;
        };
        let vtable_va = slot.wrapping_add(8);
        if !seen.insert(vtable_va) {
            continue;
        }
        // offset-to-top sits at vtable-16 (just below the typeinfo field): a plain, non-relocated,
        // pointer-aligned int, 0 for a primary table and a small negative for sub-object tables.
        let Some(ott) = img.read_i64(slot.wrapping_sub(8)) else {
            continue;
        };
        if !(-(1 << 24)..=0).contains(&ott) || ott % 8 != 0 {
            continue;
        }
        let slots = read_slots(img, vtable_va, max_slots);
        if slots.len() < 2 {
            continue;
        }
        let bases = typeinfo_bases(img, val, &kinds);
        out.push(ClassVtable {
            mangled,
            name,
            vtable_va,
            offset_to_top: ott,
            typeinfo: val,
            slots,
            bases,
        });
    }
    out.sort_by_key(|c| c.vtable_va);
    out
}