//! Relocation and code-identity masking for PS2 level overlays.
//!
//! Ported from ReRAC (`vendor/rerac/crates/rc-formats/src/level_overlay.rs`)
//! and `font.rs`. Every level overlay carries the same engine code linked at
//! different addresses, so the same function differs only in relocated fields.
//! Masking those fields makes code-identity comparison work across all 27.
//!
//! # Layout
//!
//! Overlays (`extracted/levels/NN/overlay.bin`) open with a 16-byte header,
//! then a list of sections `{u32 dest, u32 size, u32 kind, u32 entry}` (the
//! list ends when the `entry` point stops matching the first). `kind == 8`
//! marks BSS. The `.text` section is the one that holds engine code.
//!
//! # Masking
//!
//! `mask` hides exactly the relocated fields: `j`/`jal` targets, `lui`
//! immediates of main-RAM addresses (0x0010..0x0040, so float and
//! small-integer `lui`s stay compared), and the 16-bit immediates of the
//! instructions that complete such a `lui` (`%lo`, tracked per register,
//! also through `move` or addressing `$gp`). Two functions are "the same
//! code" when their masked words are equal over the reference function's
//! extent.
//!
//! # Relocation
//!
//! `Relocation` maps a reference (level-01) function or data address to its
//! copy in the target overlay: `func` by masked-code lookup, `data` by
//! finding a reference function that forms the address with `lui`/`%lo` and
//! reading the same instruction pair in that function's copy.

use crate::buf::{Buf, Error, ErrorKind, Result};
use std::collections::{BTreeSet, HashMap};

/// Index of the `lvl.vtbl` section in the overlay's section list (0-based).
pub const VTBL_SECTION: usize = 3;
/// Index of the `lvl.camvtbl` section.
pub const CAMVTBL_SECTION: usize = 4;
/// Index of the `.text` section.
pub const TEXT_SECTION: usize = 6;
/// `$gp` in every level (the boot ELF sets it once; the small data is addressed
/// from it). Confirmed across all 27 overlays.
pub const GP: u32 = 0x1AEFF0;
/// Longest function extent compared (words). A longer reference extent is cut
/// here.
pub const MAX_EXTENT: usize = 0x8000 / 4;

/// One section of a level overlay (`{u32 dest, u32 size, u32 kind, u32 entry}`
/// then `size` bytes; ends when `entry` stops matching the first). Mirrors
/// `vendor/rerac/crates/rc-formats/src/font.rs::OverlaySection`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OverlaySection {
    /// Load (RAM) address of this section.
    pub dest: u32,
    /// ELF `sh_type`: 1 = PROGBITS, 8 = NOBITS (BSS).
    pub kind: u32,
    /// Entry point of the section; the table ends when this stops matching the
    /// first section's entry.
    pub entry: u32,
    /// The section's bytes (the 16-byte header included, since it is the first
    /// `size` bytes of the section in the table).
    pub data: Vec<u8>,
}

/// Splits the raw `overlay` lump into its sections.
/// RAM address at which the overlay's `.lit` section is loaded.
pub const LIT_BASE: u32 = 0x001A_7C80;

/// The bytes of the overlay at load address `addr`.
///
/// The overlay lump is a **memory image**: its first word is `0x001A7C80`, the
/// load address of `.lit`, and `overlay[addr - LIT_BASE]` is the byte at `addr`.
///
/// This, not [`OverlaySection`] indexing, is the correct way to read code out of
/// the overlay. `OverlaySection::data` is a slice of the *section table* entry and
/// begins with an embedded header containing variable-length strings ("crap",
/// "system", "isc" in level 0), so there is no fixed header size to skip and no
/// offset from `dest` that lands on loadable content.
///
/// Bounds-checked. Sixteen probes indexed `secs[n].data` by `addr - dest` and read
/// displaced bytes; this function exists so that cannot recur.
#[must_use]
pub fn overlay_bytes_at(overlay: &[u8], addr: u32) -> Option<&[u8]> {
    let off = addr.checked_sub(LIT_BASE)? as usize;
    overlay.get(off..)
}

pub fn parse_overlay_sections(overlay: &[u8]) -> Result<Vec<OverlaySection>> {
    let b = Buf::new(overlay);
    let mut out: Vec<OverlaySection> = Vec::new();
    let mut pos = 0usize;
    if overlay.len() < 16 {
        return Err(Error {
            offset: 0,
            context: "overlay section table",
            kind: ErrorKind::OutOfBounds {
                needed: 16,
                available: overlay.len(),
            },
        });
    }
    loop {
        if pos + 16 > overlay.len() {
            break;
        }
        let dest = b.peek_u32(pos, "overlay section dest")?;
        let size = b.peek_u32(pos + 4, "overlay section size")? as usize;
        let kind = b.peek_u32(pos + 8, "overlay section kind")?;
        let entry = b.peek_u32(pos + 12, "overlay section entry")?;
        if out.first().is_some_and(|f| f.entry != entry) {
            break;
        }
        let data = b
            .peek_bytes(pos + 16, size, "overlay section payload")?
            .to_vec();
        out.push(OverlaySection {
            dest,
            kind,
            entry,
            data,
        });
        pos += 16 + size;
    }
    if out.is_empty() {
        return Err(Error {
            offset: 0,
            context: "overlay section table",
            kind: ErrorKind::Invalid("overlay has no sections"),
        });
    }
    Ok(out)
}

/// `n` bytes at EE address `addr` of the loaded overlay.
pub fn read_overlay(sections: &[OverlaySection], addr: u32, n: usize) -> Option<&[u8]> {
    sections.iter().filter(|s| s.kind != 8).find_map(|s| {
        let off = addr.checked_sub(s.dest)? as usize;
        s.data.get(off..off.checked_add(n)?)
    })
}

/// Code-word convenience: four bytes at `addr` as a little-endian `u32`.
pub fn read_u32(sections: &[OverlaySection], addr: u32) -> Option<u32> {
    read_overlay(sections, addr, 4).map(|b| u32::from_le_bytes(b.try_into().unwrap()))
}

/// A parsed level overlay.
#[derive(Debug, Clone)]
pub struct LevelOverlay {
    pub sections: Vec<OverlaySection>,
    /// The `.text` words (little-endian), at `text_start`.
    pub text: Vec<u32>,
    /// The physical load address of `.text`.
    pub text_start: u32,
    /// Set of function starts, computed from `.text` (jal targets, vtbl
    /// updates, frame set-ups, lui/%lo references).
    pub starts: BTreeSet<u32>,
}

impl LevelOverlay {
    /// Build from already-parsed sections.
    pub fn from_sections(sections: Vec<OverlaySection>) -> Result<Option<LevelOverlay>> {
        let s = &sections[TEXT_SECTION];
        if s.kind != 1 {
            return Ok(None);
        }
        if s.data.len() % 4 != 0 {
            return Err(Error {
                offset: s.data.len(),
                context: "overlay .text size not a multiple of 4",
                kind: ErrorKind::OutOfBounds {
                    needed: s.data.len(),
                    available: s.data.len(),
                },
            });
        }
        let words: Vec<u32> = s
            .data
            .chunks(4)
            .map(|c| u32::from_le_bytes(c.try_into().unwrap()))
            .collect();
        let (text_start, text_words) = (s.dest, words);
        let end = text_start + (text_words.len() as u32) * 4;
        let mut ov = LevelOverlay {
            sections,
            text: text_words,
            text_start,
            starts: BTreeSet::new(),
        };
        let t = &ov.text;
        // .text function starts: every `jal` target inside .text.
        for (p, &w) in t.iter().enumerate() {
            if op(w) == 3 {
                let a = (w & 0x03ff_ffff) << 2;
                let a = text_start + (p as u32 + 1) * 4 + a;
                if a >= text_start && a < end && a % 4 == 0 {
                    ov.starts.insert(a);
                }
            }
        }
        // lvl.vtbl update targets inside .text.
        if let Some(v) = ov.sections.get(3) {
            // `chunks_exact` rather than `as_chunks`: the latter is 1.88 and the
            // workspace declares 1.85. Same twelve-byte records.
            for e in v.data.chunks_exact(12) {
                let u = u32::from_le_bytes(e[4..8].try_into().unwrap());
                if u >= text_start && u < end && u % 4 == 0 {
                    ov.starts.insert(u);
                }
            }
        }
        // A frame set-up right after a return (`jr ra` + delay slot, then
        // `addiu sp, sp, -n`).
        for p in 2..t.len() {
            let w = t[p];
            if t[p - 2] == 0x03e0_0008 && w >> 16 == 0x27bd && w & 0x8000 != 0 {
                ov.starts.insert(text_start + 4 * p as u32);
            }
        }
        // lui/%lo references (callbacks handed to other code).
        for (_, a) in address_refs(t) {
            if a >= text_start && a < end && a % 4 == 0 {
                ov.starts.insert(a);
            }
        }
        Ok(Some(ov))
    }

    /// `n` words at `addr` of `.text`.
    pub fn code(&self, addr: u32, words: usize) -> Option<&[u32]> {
        let i = addr.checked_sub(self.text_start)?;
        if i % 4 != 0 {
            return None;
        }
        let i = (i / 4) as usize;
        self.text.get(i..i.checked_add(words)?)
    }

    /// The code extent of the function at `addr`, in words: up to the next
    /// known function start, at most `MAX_EXTENT` words and the end of `.text`.
    pub fn extent(&self, addr: u32) -> Option<usize> {
        if addr < self.text_start || addr >= self.text_end() {
            return None;
        }
        let next = self
            .starts
            .range(addr + 4..)
            .next()
            .copied()
            .unwrap_or(self.text_end());
        Some(
            (next
                .min(addr.saturating_add(MAX_EXTENT as u32))
                .min(self.text_end())
                - addr) as usize
                / 4,
        )
    }

    fn text_end(&self) -> u32 {
        self.text_start + self.text.len() as u32 * 4
    }
}

/// Op / rs / rt / rd field extractors (module doc).
const fn op(w: u32) -> u32 {
    w >> 26
}
const fn rs(w: u32) -> usize {
    ((w >> 21) & 31) as usize
}
const fn rt(w: u32) -> usize {
    ((w >> 16) & 31) as usize
}
const fn rd(w: u32) -> usize {
    ((w >> 11) & 31) as usize
}
fn sext(w: u32) -> u32 {
    (w & 0xffff) as u16 as i16 as i32 as u32
}

/// Branches (not relocated: relative offsets).
fn is_branch(o: u32) -> bool {
    matches!(o, 1 | 4..=7 | 0x14..=0x17)
}
/// I-type instructions that write GPR `rt`.
fn writes_rt(o: u32) -> bool {
    matches!(o, 8..=15 | 0x18..=0x1b | 0x1e | 0x20..=0x27 | 0x37)
}
/// `lui` of a main-RAM address (the only ones a relink changes).
fn lui_is_address(w: u32) -> bool {
    (0x0010..0x0040).contains(&(w & 0xffff))
}

/// Per-register "holds a `lui` of an address" state while walking code in
/// address order. Calls and branches do not reset it (a `lui` in a branch's
/// delay slot is used at the branch target, after code that may call): a
/// register keeps its state until an instruction writes it.
#[derive(Clone, Copy, Default)]
struct HiState([Option<u32>; 32]);

impl HiState {
    /// Masks `w` and updates the state; returns (masked word, the full address
    /// `w` forms from a `lui`, if any).
    fn step(&mut self, w: u32) -> (u32, Option<u32>) {
        self.step_inner(w)
    }

    fn step_inner(&mut self, w: u32) -> (u32, Option<u32>) {
        let o = op(w);
        match o {
            // j / jal: the target.
            2 | 3 => (w & 0xfc00_0000, None),
            15 => {
                let r = rt(w);
                if r == 0 {
                    return (w, None);
                }
                if lui_is_address(w) {
                    self.0[r] = Some(w & 0xffff);
                    (w & 0xffff_0000, None)
                } else {
                    self.0[r] = None;
                    (w, None)
                }
            }
            0 => {
                let f = w & 63;
                let (s, t, d) = (rs(w), rt(w), rd(w));
                // move (addu / daddu / or with $zero) keeps the state; any other
                // write clears it.
                let moved = if matches!(f, 0x21 | 0x2d | 0x25) {
                    if t == 0 {
                        self.0[s]
                    } else if s == 0 {
                        self.0[t]
                    } else {
                        None
                    }
                } else {
                    None
                };
                if d != 0 {
                    self.0[d] = moved;
                }
                (w, None)
            }
            _ if is_branch(o) => (w, None),
            0x10..=0x13 => {
                // Coprocessor moves to a GPR (mfc / dmfc / qmfc / cfc: rs 0..2)
                // write rt.
                if rs(w) <= 2 && rt(w) != 0 {
                    self.0[rt(w)] = None;
                }
                (w, None)
            }
            0x1c => {
                // MMI: writes rd.
                if rd(w) != 0 {
                    self.0[rd(w)] = None;
                }
                (w, None)
            }
            _ => {
                let s = rs(w);
                // `ori` completes a `lui` with the zero-extended immediate, the
                // others sign-extend it.
                let lo = if o == 13 { w & 0xffff } else { sext(w) };
                let full = if s == 28 {
                    Some(GP.wrapping_add(sext(w)))
                } else {
                    self.0[s].map(|hi| (hi << 16).wrapping_add(lo))
                };
                let masked = if full.is_some() { w & 0xffff_0000 } else { w };
                if writes_rt(o) && rt(w) != 0 {
                    self.0[rt(w)] = None;
                }
                (masked, full)
            }
        }
    }
}

/// The masked words of `words` (module doc), walked from a function start.
pub fn mask(words: &[u32]) -> Vec<u32> {
    walk(words).into_iter().map(|(m, _)| m).collect()
}

/// The (word offset, address) of every `lui`/`%lo` pair completed in `words`,
/// and of every `$gp`-relative access.
pub fn address_refs(words: &[u32]) -> Vec<(usize, u32)> {
    walk(words)
        .into_iter()
        .enumerate()
        .filter_map(|(i, (_, a))| a.map(|a| (i, a)))
        .collect()
}

/// The masked first four words (the relocator's index key; equal to the first
/// four of `mask` of any longer run).
pub fn key(w: &[u32]) -> [u32; 4] {
    let m = walk(&w[..4]);
    [m[0].0, m[1].0, m[2].0, m[3].0]
}

/// Walks `words` from a function start in address order; at a forward branch
/// target the state also takes the registers the branch brought in (after its
/// delay slot), so a `lui` in a delay slot reaches its use.
fn walk(words: &[u32]) -> Vec<(u32, Option<u32>)> {
    let mut st = HiState::default();
    let mut pending: HashMap<usize, HiState> = HashMap::new();
    let mut branch_to: Option<usize> = None;
    let mut out = Vec::with_capacity(words.len());
    for (i, &w) in words.iter().enumerate() {
        if let Some(p) = pending.remove(&i) {
            for r in 0..32 {
                if st.0[r].is_none() {
                    st.0[r] = p.0[r];
                }
            }
        }
        out.push(st.step(w));
        if let Some(t) = branch_to.take() {
            // The delay slot has run: the branch's state reaches its target.
            let e = pending.entry(t).or_default();
            for r in 0..32 {
                if e.0[r].is_none() {
                    e.0[r] = st.0[r];
                }
            }
        }
        let o = op(w);
        if is_branch(o) || (o == 0x11 && rs(w) == 8) {
            let t = i as i64 + 1 + (w & 0xffff) as u16 as i16 as i64;
            if t > i as i64 + 1 && (t as usize) < words.len() {
                branch_to = Some(t as usize);
            }
        }
    }
    out
}

/// Reference → target address translation (module doc).
pub struct Relocation<'a> {
    reference: &'a LevelOverlay,
    target: &'a LevelOverlay,
    identity: bool,
    /// Target `.text` positions by the masked key of their first 4 words.
    index: HashMap<[u32; 4], Vec<u32>>,
    /// Reference data address → the (function start, word offset) that form it
    /// (built on the first [`Relocation::data`]).
    xrefs: std::cell::OnceCell<HashMap<u32, Vec<(u32, usize)>>>,
}

impl<'a> Relocation<'a> {
    /// Identity when the target is the reference itself (same bytes).
    pub fn new(reference: &'a LevelOverlay, target: &'a LevelOverlay) -> Relocation<'a> {
        let identity = reference.text_start == target.text_start
            && reference.text == target.text
            && reference.sections == target.sections;
        let mut r = Relocation {
            reference,
            target,
            identity,
            index: HashMap::new(),
            xrefs: std::cell::OnceCell::new(),
        };
        if identity {
            return r;
        }
        let t = &target.text;
        for p in 0..t.len().saturating_sub(3) {
            r.index
                .entry(key(&t[p..p + 4]))
                .or_default()
                .push(target.text_start + 4 * p as u32);
        }
        r
    }

    /// Whether the target is the reference overlay itself.
    pub fn is_identity(&self) -> bool {
        self.identity
    }

    /// The reference function at `ref_fn`, masked over its extent.
    fn ref_masked(&self, ref_fn: u32) -> Option<Vec<u32>> {
        let n = self.reference.extent(ref_fn)?;
        Some(mask(self.reference.code(ref_fn, n)?))
    }

    /// Whether the target code at `target_fn` is the reference function
    /// `ref_fn`.
    pub fn same_code(&self, ref_fn: u32, target_fn: u32) -> bool {
        if self.identity {
            return ref_fn == target_fn;
        }
        let Some(m) = self.ref_masked(ref_fn) else {
            return false;
        };
        self.target
            .code(target_fn, m.len())
            .is_some_and(|c| mask(c) == m)
    }

    /// Every copy of the reference function `ref_fn` in the target's `.text`.
    pub fn copies(&self, ref_fn: u32) -> Vec<u32> {
        if self.identity {
            return self
                .reference
                .code(ref_fn, 1)
                .map(|_| vec![ref_fn])
                .unwrap_or_default();
        }
        let Some(m) = self.ref_masked(ref_fn).filter(|m| m.len() >= 4) else {
            return Vec::new();
        };
        let Some(cands) = self.index.get(&[m[0], m[1], m[2], m[3]]) else {
            return Vec::new();
        };
        cands
            .iter()
            .copied()
            .filter(|&c| self.target.code(c, m.len()).is_some_and(|w| mask(w) == m))
            .collect()
    }

    /// The target's copy of the reference function `ref_fn` (None: not found,
    /// or found more than once).
    pub fn func(&self, ref_fn: u32) -> Option<u32> {
        if self.identity {
            return Some(ref_fn);
        }
        match self.copies(ref_fn).as_slice() {
            [a] => Some(*a),
            _ => None,
        }
    }

    /// The target's address for the reference data address `ref_addr` (module
    /// doc). None when no reference function forms it with `lui`/`%lo`, or none
    /// of those functions has a copy in the target.
    pub fn data(&self, ref_addr: u32) -> Option<u32> {
        if self.identity {
            return Some(ref_addr);
        }
        let xrefs = self.xrefs.get_or_init(|| {
            let mut x: HashMap<u32, Vec<(u32, usize)>> = HashMap::new();
            for &f in &self.reference.starts {
                let Some(code) = self
                    .reference
                    .extent(f)
                    .and_then(|n| self.reference.code(f, n))
                else {
                    continue;
                };
                for (off, a) in address_refs(code) {
                    x.entry(a).or_default().push((f, off));
                }
            }
            x
        });
        let sites = xrefs.get(&ref_addr)?;
        for &(f, off) in sites {
            let Some(g) = self.func(f) else {
                continue;
            };
            let n = self.reference.extent(f)?;
            let refs = address_refs(self.target.code(g, n)?);
            if let Some(&(_, a)) = refs.iter().find(|(o, _)| *o == off) {
                return Some(a);
            }
        }
        // No copy of a whole referencing function (its extent runs into level
        // code): the code around one of the references, from its function start
        // up to 16 words past it, found once in the target.
        for &(f, off) in sites {
            let n = (off + 16).min(self.reference.extent(f)?);
            let m = mask(self.reference.code(f, n)?);
            if m.len() < 4 {
                continue;
            }
            let Some(cands) = self.index.get(&[m[0], m[1], m[2], m[3]]) else {
                continue;
            };
            let hits: Vec<u32> = cands
                .iter()
                .copied()
                .filter(|&c| self.target.code(c, n).is_some_and(|w| mask(w) == m))
                .collect();
            if let [g] = hits.as_slice() {
                if let Some(&(_, a)) = address_refs(self.target.code(*g, n)?)
                    .iter()
                    .find(|(o, _)| *o == off)
                {
                    return Some(a);
                }
            }
        }
        None
    }
}

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

    #[test]
    fn masking_follows_lui_through_moves_and_keeps_constants() {
        // lui v0,0x17 ; addiu s1,v0,0x42c0 ; move s2,v0 ; addiu s1,s2,0x42c0 ; lui at,0x3f80 ; lw t0,8(a0) ; lw t1,-0x10(gp)
        let w = [
            0x3c02_0017,
            0x2451_42c0,
            0x0040_902d,
            0x2651_42c0,
            0x3c01_3f80,
            0x8c88_0008,
            0x8f89_fff0,
        ];
        let m = mask(&w);
        assert_eq!(
            m,
            vec![
                0x3c02_0000,
                0x2451_0000,
                0x0040_902d,
                0x2651_0000,
                0x3c01_3f80,
                0x8c88_0008,
                0x8f89_0000
            ]
        );
        assert_eq!(
            address_refs(&w),
            vec![(1, 0x1742c0), (3, 0x1742c0), (6, GP - 0x10)]
        );
        // lui a0,0x1b ; addiu a0,a0,-0x1000 → 0x1af000
        assert_eq!(
            address_refs(&[0x3c04_001b, 0x2484_f000]),
            vec![(1, 0x1a_f000)]
        );
    }

    /// A level-01 engine function, from the overlay's own `lvl.vtbl` update
    /// table, must relocate to exactly one copy on every other level's overlay.
    #[test]
    fn mask_finds_one_copy_of_a_level_01_function_on_all_27_overlays() {
        let data = std::env::var_os("GC_EXTRACTED")
            .map(PathBuf::from)
            .unwrap_or_else(|| {
                PathBuf::from(std::env::var("HOME").unwrap_or_default())
                    .join("Documents/gc-extracted")
            });
        if !data.exists() {
            eprintln!("skipped: unpack the disc first");
            return;
        }
        // Parse all 27 overlays and build their LevelOverlay.
        let mut overlays = Vec::with_capacity(27);
        for lv in 0..27usize {
            let ov = std::fs::read(data.join(format!("level{lv}/overlay.bin")))
                .unwrap_or_else(|_| panic!("level{lv}/overlay.bin"));
            let secs =
                parse_overlay_sections(&ov).unwrap_or_else(|_| panic!("level{lv} section parse"));
            let lvl = LevelOverlay::from_sections(secs).expect("LevelOverlay::from_sections");
            overlays.push(lvl.unwrap_or_else(|| panic!("level {lv}: no .text section")));
        }
        // Reference function: a level-01 engine function, found as the first
        // function start in level 0's .text whose masked extent is at least 4
        // words (short functions would index with fewer than 4 masked words).
        let ref_lv = &overlays[0];
        let mut ref_fn = 0;
        let mut m = Vec::new();
        for &fn_start in &ref_lv.starts {
            let ext = ref_lv
                .extent(fn_start)
                .expect("extent for a function start");
            let c = ref_lv
                .code(fn_start, ext)
                .expect("code for a function start");
            let masked = mask(c);
            if masked.len() >= 4 {
                ref_fn = fn_start;
                m = masked;
                break;
            }
        }
        assert!(
            !m.is_empty(),
            "no function with masked extent >= 4 in level 0 .text"
        );
        eprintln!("ref_fn 0x{ref_fn:x} masked to {} words", m.len());
        // On every other overlay, its relocated copy must be found exactly once.
        for (i, target) in overlays.iter().enumerate().skip(1) {
            let copies = Relocation::new(ref_lv, target).copies(ref_fn);
            assert_eq!(
                copies.len(),
                1,
                "level {i}: ref_fn 0x{ref_fn:x} found {} copies (expected 1), first={:x?}",
                copies.len(),
                copies.first().copied()
            );
        }
        eprintln!("mask: ref_fn 0x{ref_fn:x} found exactly once on all 26 other overlays");
    }
}