//! Structural, recompilation-invariant fingerprint of a function — integer features only.
//!
//! Walks the control-flow graph from the entry point (bounded), decoding with iced-x86, and
//! summarises shape into counts that survive a recompile: instruction/block/call/branch counts,
//! distinct callees, RIP-relative data references, and a mnemonic-category histogram. No
//! addresses or immediates enter the vector (those change every build); only structure does.

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

pub const NCATS: usize = 18;
/// Hash buckets for the referenced-string content sketch.
pub const NREF: usize = 32;

#[derive(Default, Clone)]
pub struct Fingerprint {
    pub size_bytes: u32,
    pub insns: u32,
    pub blocks: u32,
    pub calls: u32,
    pub distinct_callees: u32,
    pub cond_branches: u32,
    pub uncond_branches: u32,
    pub rets: u32,
    pub data_refs: u32,
    pub indirect: u32,
    // call-graph context: aggregate shape of this function's callees (a callee's structure is
    // build-invariant), which distinguishes otherwise-identical thunks by which function they hit.
    pub ctx_ins: u32,
    pub ctx_calls: u32,
    pub ctx_br: u32,
    pub cats: [u32; NCATS],
    // content sketch: hash-bucket histogram of the printable rodata strings this function
    // references. The string *content* is build-invariant and highly function-specific, so it adds
    // discriminative signal the pure structural counts lack.
    pub refs: [u32; NREF],
}

impl Fingerprint {
    pub fn to_vec(&self) -> Vec<u32> {
        let mut v = vec![
            self.size_bytes,
            self.insns,
            self.blocks,
            self.calls,
            self.distinct_callees,
            self.cond_branches,
            self.uncond_branches,
            self.rets,
            self.data_refs,
            self.indirect,
            self.ctx_ins,
            self.ctx_calls,
            self.ctx_br,
        ];
        v.extend_from_slice(&self.cats);
        v.extend_from_slice(&self.refs);
        v
    }
}

/// If `target` points at a printable C string in read-only data, hash it into a `[0, NREF)` bucket.
/// Only strings of a few printable chars count — this skips jump tables and pointer arrays (whose
/// bytes are addresses that move across builds), keeping the sketch build-invariant.
fn string_bucket(img: &CodeImage, target: u64) -> Option<usize> {
    let bytes = img.rodata_at(target, 32)?;
    let run: &[u8] = {
        let end = bytes
            .iter()
            .position(|&c| !(0x20..0x7f).contains(&c))
            .unwrap_or(bytes.len());
        &bytes[..end]
    };
    if run.len() < 3 {
        return None;
    }
    // FNV-1a over the printable run
    let mut h: u64 = 0xcbf29ce484222325;
    for &b in run {
        h = (h ^ b as u64).wrapping_mul(0x100000001b3);
    }
    Some((h % NREF as u64) as usize)
}

fn category(insn: &Instruction) -> usize {
    match insn.flow_control() {
        FlowControl::Call | FlowControl::IndirectCall => return 10,
        FlowControl::UnconditionalBranch | FlowControl::IndirectBranch => return 11,
        FlowControl::ConditionalBranch => return 12,
        FlowControl::Return => return 13,
        _ => {}
    }
    if insn.op0_register().is_xmm()
        || insn.op0_register().is_ymm()
        || insn.op1_register().is_xmm()
        || insn.op1_register().is_ymm()
    {
        return 17;
    }
    use Mnemonic::*;
    match insn.mnemonic() {
        Mov | Movzx | Movsx | Movsxd | Xchg => 0,
        Lea => 1,
        Push => 2,
        Pop => 3,
        Add | Adc | Sub | Sbb | Inc | Dec | Neg => 4,
        Imul | Mul | Idiv | Div => 5,
        And | Or | Xor | Not => 6,
        Shl | Shr | Sar | Rol | Ror | Shld | Shrd => 7,
        Cmp => 8,
        Test => 9,
        Nop | Int3 => 14,
        Leave => 15,
        _ => 16,
    }
}

fn near_target(insn: &Instruction) -> Option<u64> {
    match insn.op0_kind() {
        OpKind::NearBranch16 | OpKind::NearBranch32 | OpKind::NearBranch64 => {
            Some(insn.near_branch_target())
        }
        _ => None,
    }
}

fn in_span(t: u64, entry: u64, cap: usize) -> bool {
    t >= entry && (t - entry) as usize <= cap
}

struct Base {
    f: Fingerprint,
    targets: Vec<u64>, // outgoing call + tail-jump destinations
}

/// Cheap, build-invariant summary of a callee — (instructions, calls, branches) over a bounded
/// linear scan from its entry. Enough to tell one thunk's target from another's.
fn light_summary(img: &CodeImage, addr: u64) -> (u32, u32, u32) {
    let Some(code) = img.code_at(addr) else {
        return (0, 0, 0);
    };
    let mut dec = Decoder::with_ip(64, code, addr, DecoderOptions::NONE);
    let (mut ins, mut calls, mut br) = (0u32, 0u32, 0u32);
    while dec.can_decode() && ins < 48 {
        let insn = dec.decode();
        if insn.len() == 0 || insn.is_invalid() {
            break;
        }
        ins += 1;
        match insn.flow_control() {
            FlowControl::Call | FlowControl::IndirectCall => calls += 1,
            FlowControl::ConditionalBranch
            | FlowControl::UnconditionalBranch
            | FlowControl::IndirectBranch => br += 1,
            FlowControl::Return => break,
            _ => {}
        }
    }
    (ins, calls, br)
}

/// Extract the fingerprint of `entry` — base structure plus one hop of call-graph context.
pub fn extract(img: &CodeImage, entry: u64) -> Option<Fingerprint> {
    let base = extract_base(img, entry)?;
    let mut f = base.f;
    let mut seen = HashSet::new();
    for &t in &base.targets {
        if t != entry && img.is_code(t) && seen.insert(t) {
            let (ins, calls, br) = light_summary(img, t);
            f.ctx_ins = f.ctx_ins.wrapping_add(ins);
            f.ctx_calls = f.ctx_calls.wrapping_add(calls);
            f.ctx_br = f.ctx_br.wrapping_add(br);
        }
    }
    Some(f)
}

/// Base structural features + the function's outgoing call/tail-jump targets.
fn extract_base(img: &CodeImage, entry: u64) -> Option<Base> {
    let code = img.code_at(entry)?;
    const MAX_SPAN: usize = 96 * 1024;
    const MAX_INSNS: u32 = 8000;
    let cap = code.len().min(MAX_SPAN);

    let mut f = Fingerprint::default();
    let mut visited: HashSet<u64> = HashSet::new();
    let mut leaders: HashSet<u64> = HashSet::new();
    let mut callees: HashSet<u64> = HashSet::new();
    let mut targets: Vec<u64> = Vec::new();
    let mut hi = entry;
    let mut work = vec![entry];
    leaders.insert(entry);

    while let Some(start) = work.pop() {
        let mut ip = start;
        loop {
            if ip < entry || (ip - entry) as usize >= cap || f.insns >= MAX_INSNS {
                break;
            }
            if !visited.insert(ip) {
                break; // this path merged into already-decoded code
            }
            let off = (ip - entry) as usize;
            let mut dec = Decoder::with_ip(64, &code[off..], ip, DecoderOptions::NONE);
            if !dec.can_decode() {
                break;
            }
            let insn = dec.decode();
            let ilen = insn.len() as u64;
            if ilen == 0 || insn.is_invalid() {
                break;
            }
            f.insns += 1;
            f.cats[category(&insn)] += 1;
            if insn.is_ip_rel_memory_operand() {
                f.data_refs += 1;
                if let Some(bucket) = string_bucket(img, insn.memory_displacement64()) {
                    f.refs[bucket] += 1;
                }
            }
            if ip + ilen > hi {
                hi = ip + ilen;
            }
            let next = ip + ilen;
            match insn.flow_control() {
                FlowControl::Return => {
                    f.rets += 1;
                    break;
                }
                FlowControl::Call => {
                    f.calls += 1;
                    match near_target(&insn) {
                        Some(t) => {
                            callees.insert(t);
                            targets.push(t);
                        }
                        None => f.indirect += 1,
                    }
                    ip = next;
                }
                FlowControl::IndirectCall => {
                    f.calls += 1;
                    f.indirect += 1;
                    ip = next;
                }
                FlowControl::ConditionalBranch => {
                    f.cond_branches += 1;
                    if let Some(t) = near_target(&insn)
                        && in_span(t, entry, cap)
                    {
                        leaders.insert(t);
                        work.push(t);
                    }
                    leaders.insert(next);
                    ip = next;
                }
                FlowControl::UnconditionalBranch => {
                    f.uncond_branches += 1;
                    match near_target(&insn) {
                        Some(t) if in_span(t, entry, cap) => {
                            leaders.insert(t);
                            ip = t;
                        }
                        Some(t) => {
                            targets.push(t); // tail call (thunk target)
                            break;
                        }
                        None => break, // indirect jump
                    }
                }
                FlowControl::IndirectBranch => {
                    f.indirect += 1;
                    break;
                }
                _ => ip = next,
            }
        }
    }

    if f.insns == 0 {
        return None;
    }
    f.blocks = leaders.len() as u32;
    f.distinct_callees = callees.len() as u32;
    f.size_bytes = (hi - entry) as u32;
    Some(Base { f, targets })
}