//! Find the MIPS code that references a string in `overlay.bin`.
//!
//! The overlay has no symbol table, but it has a fixed load base and a fixed
//! `$gp`, so every string has a GP-relative offset. MIPS materialises that as a
//! `%hi`/`%lo` pair, and the `%lo` half is a 16-bit immediate that occurs in the
//! instruction stream. Searching for it turns a string into a code address.
//!
//! This is the step that precedes disassembly, and it needs no disassembler.
use anyhow::{anyhow, Context, Result};
use std::collections::BTreeSet;

/// Where a level program is loaded, and its global pointer. Both are stated in
/// OpenRAC/rac2-gc-decomp docs/LEVELS.md and hold for all 27 levels.
const LIT_BASE: u64 = 0x1A7C80;
const GP: u64 = 0x1AEFF0;

fn main() -> Result<()> {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let path = args
        .first()
        .context("usage: gc-overlay-anchors <overlay.bin> [needle] [--level N]")?;
    let overlay = std::fs::read(path).with_context(|| format!("read {path}"))?;

    if args.get(1).map(String::as_str) == Some("--disasm") {
        let ram = parse_ram(args.get(2).context("--disasm needs a RAM address")?)?;
        let before: usize = args.get(3).and_then(|v| v.parse().ok()).unwrap_or(64);
        let after: usize = args.get(4).and_then(|v| v.parse().ok()).unwrap_or(64);
        disasm(&overlay, ram, before, after)?;
        return Ok(());
    }
    match args.get(1) {
        None => {
            list_strings(&overlay);
            Ok(())
        }
        Some(needle) => {
            let needle = needle.as_bytes();
            let hits = find_all(&overlay, needle);
            if hits.is_empty() {
                return Err(anyhow!(
                    "{:?} not found in {path}",
                    String::from_utf8_lossy(needle)
                ));
            }
            for off in &hits {
                report(&overlay, needle, *off);
            }
            Ok(())
        }
    }
}

fn report(overlay: &[u8], needle: &[u8], off: usize) {
    let ram = LIT_BASE + off as u64;
    let delta = ram as i64 - GP as i64;
    let lo = (delta as u16).to_be_bytes();
    eprintln!("{:?}", String::from_utf8_lossy(needle));
    eprintln!("  overlay+{off:#x}  ->  RAM {ram:#x}");
    eprintln!(
        "  $gp {GP:#x}  ->  delta {delta:#x} ({delta}), $lo half {:#06x}",
        u16::from_be_bytes(lo)
    );

    // Count instruction-aligned occurrences of the $lo immediate.
    //
    // The PS2's MIPS is little-endian, and an I-type immediate occupies bits
    // 15..0 of the word, so the low byte of the immediate is word BYTE 0 and the
    // high byte is word BYTE 1. Reading bytes 1..3 instead -- which an earlier
    // version of this file did -- matches the opcode and $rs fields instead and
    // manufactures hits out of nothing.
    let lo_byte = (delta as u16) as u8;
    let hi_byte = ((delta as u16) >> 8) as u8;
    let mut hits = BTreeSet::new();
    for i in (0..overlay.len().saturating_sub(3)).step_by(4) {
        if overlay[i] == lo_byte && overlay[i + 1] == hi_byte {
            hits.insert(i);
        }
    }
    if hits.is_empty() {
        eprintln!("  no instruction-aligned $lo occurrence -- not referenced through $gp\n");
    } else {
        eprintln!("  {} instruction-aligned occurrence(s):", hits.len());
        for i in &hits {
            eprintln!("    {:#x}  (overlay+{:#x})", LIT_BASE + *i as u64, i);
        }
        eprintln!();
    }
}

fn parse_ram(s: &str) -> Result<u64> {
    let t = s.trim_start_matches("0x").replace('_', "");
    u64::from_str_radix(&t, 16).with_context(|| format!("parse address {s}"))
}

/// Disassemble `before` bytes before and `after` bytes after `ram`.
fn disasm(overlay: &[u8], ram: u64, before: usize, after: usize) -> Result<()> {
    use capstone::prelude::*;
    use capstone::Endian;
    let off = usize::try_from(
        ram.checked_sub(LIT_BASE)
            .context("address below .lit base")?,
    )?;
    if off + after > overlay.len() {
        anyhow::bail!("{ram:#x}+{after} runs past the end of the section");
    }
    let cs = capstone::Capstone::new()
        .mips()
        .mode(capstone::arch::mips::ArchMode::Mips32)
        .endian(Endian::Little)
        .detail(true)
        .build()
        .context("capstone mips32")?;

    // Disassemble word by word rather than one `disasm_all` over the window: a
    // single undecodable word stops the whole run, and code regions have data
    // mixed in. Skipping four bytes and continuing is the point of the tool --
    // without it a window around an interesting address is usually five
    // instructions long.
    let start = off.saturating_sub(before);
    let end = (off + after).min(overlay.len());
    let mut at = start;
    let mut bad = 0usize;
    while at + 4 <= end {
        let word = u32::from_le_bytes(overlay[at..at + 4].try_into().unwrap());
        match cs.disasm_all(&overlay[at..at + 4], LIT_BASE + at as u64) {
            Ok(v) if !v.is_empty() => {
                for i in v.iter() {
                    eprintln!(
                        "{} {:#010x}  {i}",
                        if i.address() == LIT_BASE + off as u64 {
                            ">>"
                        } else {
                            "  "
                        },
                        i.address()
                    );
                }
            }
            _ => {
                eprintln!("   {:#010x}  .word {word:#010x}", LIT_BASE + at as u64);
                bad += 1;
            }
        }
        at += 4;
    }
    eprintln!("({bad} of {} words did not decode)", (end - start) / 4);
    Ok(())
}

fn find_all(h: &[u8], needle: &[u8]) -> Vec<usize> {
    if needle.is_empty() || needle.len() > h.len() {
        return Vec::new();
    }
    (0..=h.len() - needle.len())
        .filter(|&i| &h[i..i + needle.len()] == needle)
        .collect()
}

fn list_strings(d: &[u8]) {
    let mut i = 0usize;
    while i < d.len() {
        if d[i].is_ascii_graphic() || d[i] == b' ' {
            let s = i;
            while i < d.len() && (d[i].is_ascii_graphic() || d[i] == b' ') {
                i += 1
            }
            if i - s >= 8 {
                println!(
                    "{:#010x}  {}",
                    LIT_BASE + s as u64,
                    String::from_utf8_lossy(&d[s..i])
                );
            }
        } else {
            i += 1;
        }
    }
}