//! Whole-binary cross-reference index — the layer that lets us *name and locate non-virtual
//! functions*, which have no vtable slot and (mostly) no symbol.
//!
//! Decode every function and record, for each referenced address, the instructions that reference
//! it: near `call`/`jmp` targets (code) and RIP-relative memory operands (`lea`/`mov` into `.rodata`
//! strings, globals, …). Each reference is attributed to its containing function, so we can answer
//! "which function uses this string?" (string-anchored location) and "who calls this function?".
//!
//! Function entries come from `locate::candidate_entries` (relocation code-pointers — every vtable
//! slot — ∪ decoded `call` targets) unioned with `.eh_frame` starts. This matters: CS2 strips
//! `.eh_frame` from the *game* code (unwind info survives only for the statically-linked runtime
//! tail), so an eh_frame-only index misses the entire gameplay region. Decoding from each entry to
//! the next avoids the misalignment a blind section-wide linear sweep suffers on data/padding.

use crate::elf::CodeImage;
use iced_x86::{Decoder, DecoderOptions, FlowControl, Instruction, OpKind};
use std::collections::HashMap;

pub struct XrefIndex {
    entries: Vec<u64>,            // sorted, de-duped function entry addresses
    refs: HashMap<u64, Vec<u64>>, // referenced VA -> source instruction VAs
    call_targets: Vec<u64>,       // sorted, de-duped near-call targets
}

impl XrefIndex {
    pub fn build(img: &CodeImage) -> Self {
        // Reliable gameplay entries (vtable slots + fn-pointers via relocations, plus call targets),
        // then add the eh_frame starts (the runtime tail). Union = coverage of the whole binary.
        let mut entries = crate::locate::candidate_entries(img);
        entries.extend(img.eh_frame_functions().into_iter().map(|(s, _)| s));
        entries.sort_unstable();
        entries.dedup();

        // Disassemble each function's [start, next) range independently across threads — this is the
        // single biggest decode in the tool and the ranges vary wildly in size, so the atomic work
        // scheduler load-balances them. Each task returns its (ref-pair, call-target) deltas; merging
        // them in entry order (parallel_map preserves input order) reproduces the serial build
        // byte-for-byte: refs[t] receives its srcs in the same (ascending entry, then instruction)
        // order and call_targets is sorted afterwards.
        type EntryData = (Vec<(u64, u64)>, Vec<u64>);
        let idxs: Vec<usize> = (0..entries.len()).collect();
        let per_entry: Vec<EntryData> =
            crate::par::parallel_map(&idxs, crate::par::default_threads(None), |&i| {
                let start = entries[i];
                let end = entries.get(i + 1).copied().unwrap_or(u64::MAX);
                let Some(code) = img.code_range(start, end) else {
                    return (Vec::new(), Vec::new());
                };
                let mut ref_pairs: Vec<(u64, u64)> = Vec::new();
                let mut call_targets: Vec<u64> = Vec::new();
                let mut insn = Instruction::default();
                let mut dec = Decoder::with_ip(64, code, start, DecoderOptions::NONE);
                while dec.can_decode() {
                    dec.decode_out(&mut insn);
                    let src = insn.ip();
                    // Near call/jmp: the target is code; call targets double as function entries.
                    if matches!(
                        insn.op0_kind(),
                        OpKind::NearBranch16 | OpKind::NearBranch32 | OpKind::NearBranch64
                    ) {
                        let t = insn.near_branch_target();
                        ref_pairs.push((t, src));
                        if insn.flow_control() == FlowControl::Call {
                            call_targets.push(t);
                        }
                    }
                    // RIP-relative memory operand: a reference to a string / global / code pointer.
                    if insn.is_ip_rel_memory_operand() {
                        let t = insn.ip_rel_memory_address();
                        ref_pairs.push((t, src));
                    }
                }
                (ref_pairs, call_targets)
            });

        let mut refs: HashMap<u64, Vec<u64>> = HashMap::new();
        let mut call_targets = Vec::new();
        for (ref_pairs, cts) in per_entry {
            for (t, src) in ref_pairs {
                refs.entry(t).or_default().push(src);
            }
            call_targets.extend(cts);
        }
        call_targets.sort_unstable();
        call_targets.dedup();
        Self {
            entries,
            refs,
            call_targets,
        }
    }

    /// The entry (function start) that contains `va`: the nearest entry at or below `va`.
    pub fn containing_func(&self, va: u64) -> Option<u64> {
        let i = self.entries.partition_point(|&s| s <= va);
        (i > 0).then(|| self.entries[i - 1])
    }

    /// Source instruction addresses that reference `target`.
    pub fn refs_to(&self, target: u64) -> &[u64] {
        self.refs.get(&target).map_or(&[], |v| v.as_slice())
    }

    /// Distinct functions that reference `target` (each referring instruction mapped to its
    /// containing function, so one function referencing `target` N times counts once).
    pub fn referrers(&self, target: u64) -> Vec<u64> {
        let mut fs: Vec<u64> = self
            .refs_to(target)
            .iter()
            .filter_map(|&s| self.containing_func(s))
            .collect();
        fs.sort_unstable();
        fs.dedup();
        fs
    }

    pub fn call_targets(&self) -> &[u64] {
        &self.call_targets
    }
}

/// Functions that reference the string `s` anywhere in read-only data — the canonical "find the
/// function by a string it uses" primitive. A string referenced by exactly one function names that
/// function unambiguously. We match `s` as a substring (a code `lea`/`mov` points at the string's
/// start whatever follows it — a trailing `\n`, format args, or a longer literal), so callers pass a
/// distinctive fragment without needing the whole literal.
pub fn funcs_using_string(img: &CodeImage, xref: &XrefIndex, s: &str) -> Vec<u64> {
    let mut out = Vec::new();
    for str_va in img.find_bytes(s.as_bytes()) {
        out.extend(xref.referrers(str_va));
    }
    out.sort_unstable();
    out.dedup();
    out
}
