prometeu-runtime/crates/tools/pbxgen-stress/src/lib.rs

use anyhow::Result;
use prometeu_bytecode::assembler::assemble;
use prometeu_bytecode::model::{
    BytecodeModule, ConstantPoolEntry, DebugInfo, Export, FunctionMeta, SyscallDecl,
};
use prometeu_hal::asset::{
    AssetCodec, AssetEntry, BankType, PreloadEntry, SCENE_DECODED_LAYER_OVERHEAD_BYTES_V1,
    SCENE_LAYER_COUNT_V1, SCENE_PAYLOAD_MAGIC_V1, SCENE_PAYLOAD_VERSION_V1,
};
use prometeu_hal::cartridge::{
    AssetsPackHeader, AssetsPackPrelude, ASSETS_PA_MAGIC, ASSETS_PA_PRELUDE_SIZE,
    ASSETS_PA_SCHEMA_VERSION,
};
use prometeu_hal::color::Color;
use prometeu_hal::glyph::Glyph;
use prometeu_hal::glyph_bank::{
    TileSize, GLYPH_BANK_COLORS_PER_PALETTE, GLYPH_BANK_PALETTE_COUNT_V1,
};
use prometeu_hal::scene_bank::SceneBank;
use prometeu_hal::scene_layer::{ParallaxFactor, SceneLayer};
use prometeu_hal::tile::Tile;
use prometeu_hal::tilemap::TileMap;
use std::fs;
use std::mem::size_of;
use std::path::PathBuf;

fn asm(s: &str) -> Vec<u8> {
    assemble(s).expect("assemble")
}

pub fn generate() -> Result<()> {
    let mut rom: Vec<u8> = Vec::new();
    let syscalls = vec![
        SyscallDecl {
            module: "gfx".into(),
            name: "clear_565".into(),
            version: 1,
            arg_slots: 1,
            ret_slots: 0,
        },
        SyscallDecl {
            module: "gfx".into(),
            name: "draw_text".into(),
            version: 1,
            arg_slots: 4,
            ret_slots: 0,
        },
        SyscallDecl {
            module: "log".into(),
            name: "write".into(),
            version: 1,
            arg_slots: 2,
            ret_slots: 0,
        },
        SyscallDecl {
            module: "composer".into(),
            name: "bind_scene".into(),
            version: 1,
            arg_slots: 1,
            ret_slots: 1,
        },
        SyscallDecl {
            module: "composer".into(),
            name: "set_camera".into(),
            version: 1,
            arg_slots: 2,
            ret_slots: 0,
        },
        SyscallDecl {
            module: "composer".into(),
            name: "emit_sprite".into(),
            version: 1,
            arg_slots: 9,
            ret_slots: 1,
        },
    ];

    heavy_load(&mut rom);
    // light_load(&mut rom);

    let functions = vec![FunctionMeta {
        code_offset: 0,
        code_len: rom.len() as u32,
        param_slots: 0,
        local_slots: 2,
        return_slots: 0,
        max_stack_slots: 32,
    }];

    let module = BytecodeModule {
        version: 0,
        const_pool: vec![
            ConstantPoolEntry::String("stress".into()),
            ConstantPoolEntry::String("frame".into()),
            ConstantPoolEntry::String("overlay".into()),
            ConstantPoolEntry::String("composer".into()),
        ],
        functions,
        code: rom,
        debug_info: Some(DebugInfo {
            pc_to_span: vec![],
            function_names: vec![(0, "main".into())],
        }),
        exports: vec![Export { symbol: "main".into(), func_idx: 0 }],
        syscalls,
    };

    let bytes = module.serialize();

    let mut out_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    out_dir.pop(); // pbxgen-stress
    out_dir.pop(); // tools
    out_dir.pop(); // crates
    out_dir.push("test-cartridges");
    out_dir.push("stress-console");
    fs::create_dir_all(&out_dir)?;
    fs::write(out_dir.join("program.pbx"), bytes)?;
    fs::write(out_dir.join("assets.pa"), build_assets_pack()?)?;
    fs::write(out_dir.join("manifest.json"), b"{\n  \"magic\": \"PMTU\",\n  \"cartridge_version\": 1,\n  \"app_id\": 1,\n  \"title\": \"Stress Console\",\n  \"app_version\": \"0.1.0\",\n  \"app_mode\": \"Game\",\n  \"capabilities\": [\"gfx\", \"log\", \"asset\"]\n}\n")?;
    Ok(())
}

fn heavy_load(rom: &mut Vec<u8>) {
    // Single function 0: main
    // Global 0 = frame counter
    // Global 1 = scene bound flag
    // Local 0 = sprite row
    // Local 1 = sprite col

    rom.extend(asm("GET_GLOBAL 0\nPUSH_I32 1\nADD\nSET_GLOBAL 0"));

    rom.extend(asm("GET_GLOBAL 1\nPUSH_I32 0\nEQ"));
    let jif_bind_done_offset = rom.len() + 2;
    rom.extend(asm("JMP_IF_FALSE 0"));

    rom.extend(asm("PUSH_I32 0\nHOSTCALL 3\nPOP_N 1\nPUSH_I32 1\nSET_GLOBAL 1"));
    let bind_done_target = rom.len() as u32;

    rom.extend(asm("PUSH_I32 0\nHOSTCALL 0"));

    rom.extend(asm(
        "GET_GLOBAL 0\nPUSH_I32 2\nMUL\nPUSH_I32 192\nMOD\nGET_GLOBAL 0\nPUSH_I32 76\nMOD\nHOSTCALL 4",
    ));

    rom.extend(asm("PUSH_I32 0\nSET_LOCAL 0"));
    let row_loop_start = rom.len() as u32;
    rom.extend(asm("GET_LOCAL 0\nPUSH_I32 16\nLT"));
    let jif_row_end_offset = rom.len() + 2;
    rom.extend(asm("JMP_IF_FALSE 0"));

    rom.extend(asm("PUSH_I32 0\nSET_LOCAL 1"));
    let col_loop_start = rom.len() as u32;
    rom.extend(asm("GET_LOCAL 1\nPUSH_I32 32\nLT"));
    let jif_col_end_offset = rom.len() + 2;
    rom.extend(asm("JMP_IF_FALSE 0"));

    rom.extend(asm(
        "PUSH_I32 0\n\
         GET_LOCAL 0\nPUSH_I32 32\nMUL\nGET_LOCAL 1\nADD\nGET_GLOBAL 0\nADD\nPUSH_I32 15\nMOD\nPUSH_I32 1\nADD\n\
         GET_LOCAL 1\nPUSH_I32 10\nMUL\nGET_GLOBAL 0\nPUSH_I32 10\nMOD\nADD\nPUSH_I32 320\nMOD\n\
         GET_LOCAL 0\nPUSH_I32 10\nMUL\nGET_GLOBAL 0\nPUSH_I32 2\nMUL\nPUSH_I32 10\nMOD\nADD\nPUSH_I32 180\nMOD\n\
         GET_LOCAL 0\nGET_LOCAL 1\nADD\nGET_GLOBAL 0\nADD\nPUSH_I32 4\nMOD\n\
         PUSH_I32 0\n\
         GET_LOCAL 1\nGET_GLOBAL 0\nADD\nPUSH_I32 1\nBIT_AND\nPUSH_I32 0\nNEQ\n\
         GET_LOCAL 0\nGET_GLOBAL 0\nADD\nPUSH_I32 1\nBIT_AND\nPUSH_I32 0\nNEQ\n\
         GET_LOCAL 0\nGET_LOCAL 1\nADD\nPUSH_I32 4\nMOD\n\
         HOSTCALL 5\nPOP_N 1",
    ));

    rom.extend(asm("GET_LOCAL 1\nPUSH_I32 1\nADD\nSET_LOCAL 1"));
    let jmp_col_loop_offset = rom.len() + 2;
    rom.extend(asm("JMP 0"));
    let col_loop_end = rom.len() as u32;

    rom.extend(asm("GET_LOCAL 0\nPUSH_I32 1\nADD\nSET_LOCAL 0"));
    let jmp_row_loop_offset = rom.len() + 2;
    rom.extend(asm("JMP 0"));
    let row_loop_end = rom.len() as u32;

    rom.extend(asm("GET_GLOBAL 0\nPUSH_I32 3\nMUL\nPUSH_I32 220\nMOD\n\
         PUSH_I32 8\n\
         PUSH_CONST 0\n\
         GET_GLOBAL 0\nPUSH_I32 2047\nMUL\nPUSH_I32 65535\nBIT_AND\n\
         HOSTCALL 1"));
    rom.extend(asm("PUSH_I32 12\n\
         GET_GLOBAL 0\nPUSH_I32 2\nMUL\nPUSH_I32 120\nMOD\nPUSH_I32 24\nADD\n\
         PUSH_CONST 1\n\
         GET_GLOBAL 0\nPUSH_I32 4093\nMUL\nPUSH_I32 65535\nBIT_AND\n\
         HOSTCALL 1"));
    rom.extend(asm("PUSH_I32 220\n\
         GET_GLOBAL 0\nPUSH_I32 5\nMUL\nPUSH_I32 140\nMOD\n\
         PUSH_CONST 2\n\
         GET_GLOBAL 0\nPUSH_I32 1237\nMUL\nPUSH_I32 65535\nBIT_AND\n\
         HOSTCALL 1"));
    rom.extend(asm("GET_GLOBAL 0\nPUSH_I32 4\nMUL\nPUSH_I32 180\nMOD\nPUSH_I32 80\nADD\n\
         GET_GLOBAL 0\nPUSH_I32 3\nMUL\nPUSH_I32 90\nMOD\nPUSH_I32 70\nADD\n\
         PUSH_CONST 3\n\
         GET_GLOBAL 0\nPUSH_I32 3001\nMUL\nPUSH_I32 65535\nBIT_AND\n\
         HOSTCALL 1"));

    rom.extend(asm("GET_GLOBAL 0\nPUSH_I32 60\nMOD\nPUSH_I32 0\nEQ"));
    let jif_log_offset = rom.len() + 2;
    rom.extend(asm("JMP_IF_FALSE 0"));
    rom.extend(asm("PUSH_I32 2\nPUSH_CONST 0\nHOSTCALL 2"));
    let after_log = rom.len() as u32;

    rom.extend(asm("FRAME_SYNC\nRET"));

    let patch = |buf: &mut Vec<u8>, imm_offset: usize, target: u32| {
        buf[imm_offset..imm_offset + 4].copy_from_slice(&target.to_le_bytes());
    };

    patch(rom, jif_bind_done_offset, bind_done_target);
    patch(rom, jif_row_end_offset, row_loop_end);
    patch(rom, jif_col_end_offset, col_loop_end);
    patch(rom, jmp_col_loop_offset, col_loop_start);
    patch(rom, jmp_row_loop_offset, row_loop_start);
    patch(rom, jif_log_offset, after_log);
}

fn build_assets_pack() -> Result<Vec<u8>> {
    let (glyph_entry, glyph_payload) = build_glyph_asset();
    let scene = build_scene_bank();
    let scene_payload = encode_scene_payload(&scene);
    let scene_entry = AssetEntry {
        asset_id: 1,
        asset_name: "stress_scene".into(),
        bank_type: BankType::SCENE,
        offset: glyph_payload.len() as u64,
        size: scene_payload.len() as u64,
        decoded_size: expected_scene_decoded_size(&scene) as u64,
        codec: AssetCodec::None,
        metadata: serde_json::json!({}),
    };

    let asset_table = vec![glyph_entry, scene_entry];
    let preload =
        vec![PreloadEntry { asset_id: 0, slot: 0 }, PreloadEntry { asset_id: 1, slot: 0 }];
    let payload_len = glyph_payload.len() + scene_payload.len();
    let header = serde_json::to_vec(&AssetsPackHeader { asset_table, preload })?;
    let payload_offset = (ASSETS_PA_PRELUDE_SIZE + header.len()) as u64;
    let prelude = AssetsPackPrelude {
        magic: ASSETS_PA_MAGIC,
        schema_version: ASSETS_PA_SCHEMA_VERSION,
        header_len: header.len() as u32,
        payload_offset,
        flags: 0,
        reserved: 0,
        header_checksum: 0,
    };

    let mut bytes = prelude.to_bytes().to_vec();
    bytes.extend_from_slice(&header);
    bytes.extend_from_slice(&glyph_payload);
    bytes.extend_from_slice(&scene_payload);
    debug_assert_eq!(bytes.len(), payload_offset as usize + payload_len);
    Ok(bytes)
}

fn build_glyph_asset() -> (AssetEntry, Vec<u8>) {
    let pixel_indices = vec![1_u8; 8 * 8];
    let mut payload = pack_4bpp(&pixel_indices);
    payload.extend_from_slice(&build_palette_bytes());

    let entry = AssetEntry {
        asset_id: 0,
        asset_name: "stress_square".into(),
        bank_type: BankType::GLYPH,
        offset: 0,
        size: payload.len() as u64,
        decoded_size: (8 * 8 + GLYPH_BANK_PALETTE_COUNT_V1 * GLYPH_BANK_COLORS_PER_PALETTE * 2)
            as u64,
        codec: AssetCodec::None,
        metadata: serde_json::json!({
            "tile_size": 8,
            "width": 8,
            "height": 8,
            "palette_count": GLYPH_BANK_PALETTE_COUNT_V1,
            "palette_authored": GLYPH_BANK_PALETTE_COUNT_V1
        }),
    };

    (entry, payload)
}

fn build_palette_bytes() -> Vec<u8> {
    let mut bytes =
        Vec::with_capacity(GLYPH_BANK_PALETTE_COUNT_V1 * GLYPH_BANK_COLORS_PER_PALETTE * 2);
    for palette_id in 0..GLYPH_BANK_PALETTE_COUNT_V1 {
        for color_index in 0..GLYPH_BANK_COLORS_PER_PALETTE {
            let color = if color_index == 1 { stress_color(palette_id) } else { Color::BLACK };
            bytes.extend_from_slice(&color.raw().to_le_bytes());
        }
    }
    bytes
}

fn stress_color(palette_id: usize) -> Color {
    let r = ((palette_id * 53) % 256) as u8;
    let g = ((palette_id * 97 + 64) % 256) as u8;
    let b = ((palette_id * 29 + 128) % 256) as u8;
    Color::rgb(r, g, b)
}

fn pack_4bpp(indices: &[u8]) -> Vec<u8> {
    let mut packed = Vec::with_capacity(indices.len().div_ceil(2));
    for chunk in indices.chunks(2) {
        let hi = chunk[0] & 0x0f;
        let lo = chunk.get(1).copied().unwrap_or(0) & 0x0f;
        packed.push((hi << 4) | lo);
    }
    packed
}

fn build_scene_bank() -> SceneBank {
    let mut layers = std::array::from_fn(|layer_index| {
        let mut tiles = vec![
            Tile {
                active: false,
                glyph: Glyph { glyph_id: 0, palette_id: layer_index as u8 + 1 },
                flip_x: false,
                flip_y: false,
            };
            64 * 32
        ];

        for step in 0..8 {
            let x = 4 + step * 7 + layer_index * 2;
            let y = 2 + step * 3 + layer_index * 2;
            let index = y * 64 + x;
            tiles[index] = Tile {
                active: true,
                glyph: Glyph { glyph_id: 0, palette_id: layer_index as u8 + 1 },
                flip_x: false,
                flip_y: false,
            };
        }

        SceneLayer {
            active: true,
            glyph_bank_id: 0,
            tile_size: TileSize::Size8,
            parallax_factor: match layer_index {
                0 => ParallaxFactor { x: 1.0, y: 1.0 },
                1 => ParallaxFactor { x: 0.75, y: 0.75 },
                2 => ParallaxFactor { x: 0.5, y: 0.5 },
                _ => ParallaxFactor { x: 0.25, y: 0.25 },
            },
            tilemap: TileMap { width: 64, height: 32, tiles },
        }
    });

    // Keep the farthest layer a bit sparser so the diagonal remains visually readable.
    for step in 0..4 {
        let x = 10 + step * 12;
        let y = 4 + step * 5;
        let index = y * 64 + x;
        layers[3].tilemap.tiles[index].active = false;
    }

    SceneBank { layers }
}

fn expected_scene_decoded_size(scene: &SceneBank) -> usize {
    scene
        .layers
        .iter()
        .map(|layer| {
            SCENE_DECODED_LAYER_OVERHEAD_BYTES_V1 + layer.tilemap.tiles.len() * size_of::<Tile>()
        })
        .sum()
}

fn encode_scene_payload(scene: &SceneBank) -> Vec<u8> {
    let mut data = Vec::new();
    data.extend_from_slice(&SCENE_PAYLOAD_MAGIC_V1);
    data.extend_from_slice(&SCENE_PAYLOAD_VERSION_V1.to_le_bytes());
    data.extend_from_slice(&(SCENE_LAYER_COUNT_V1 as u16).to_le_bytes());
    data.extend_from_slice(&0_u32.to_le_bytes());

    for layer in &scene.layers {
        let layer_flags = if layer.active { 0b0000_0001 } else { 0 };
        data.push(layer_flags);
        data.push(layer.glyph_bank_id);
        data.push(layer.tile_size as u8);
        data.push(0);
        data.extend_from_slice(&layer.parallax_factor.x.to_le_bytes());
        data.extend_from_slice(&layer.parallax_factor.y.to_le_bytes());
        data.extend_from_slice(&(layer.tilemap.width as u32).to_le_bytes());
        data.extend_from_slice(&(layer.tilemap.height as u32).to_le_bytes());
        data.extend_from_slice(&(layer.tilemap.tiles.len() as u32).to_le_bytes());
        data.extend_from_slice(&0_u32.to_le_bytes());

        for tile in &layer.tilemap.tiles {
            let mut tile_flags = 0_u8;
            if tile.active {
                tile_flags |= 0b0000_0001;
            }
            if tile.flip_x {
                tile_flags |= 0b0000_0010;
            }
            if tile.flip_y {
                tile_flags |= 0b0000_0100;
            }
            data.push(tile_flags);
            data.push(tile.glyph.palette_id);
            data.extend_from_slice(&tile.glyph.glyph_id.to_le_bytes());
        }
    }

    data
}

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

    #[test]
    fn assets_pack_contains_preloaded_glyph_and_scene_assets() {
        let bytes = build_assets_pack().expect("assets pack");
        let prelude =
            AssetsPackPrelude::from_bytes(&bytes[..ASSETS_PA_PRELUDE_SIZE]).expect("prelude");
        assert_eq!(prelude.magic, ASSETS_PA_MAGIC);
        assert_eq!(prelude.schema_version, ASSETS_PA_SCHEMA_VERSION);

        let header_start = ASSETS_PA_PRELUDE_SIZE;
        let header_end = header_start + prelude.header_len as usize;
        let header: AssetsPackHeader =
            serde_json::from_slice(&bytes[header_start..header_end]).expect("header");

        assert_eq!(header.asset_table.len(), 2);
        assert_eq!(header.preload.len(), 2);
        assert_eq!(header.asset_table[0].bank_type, BankType::GLYPH);
        assert_eq!(header.asset_table[1].bank_type, BankType::SCENE);
        assert_eq!(header.preload[0].slot, 0);
        assert_eq!(header.preload[1].slot, 0);
        assert_eq!(header.asset_table[0].offset, 0);
        assert_eq!(header.asset_table[1].offset, header.asset_table[0].size);
        assert_eq!(
            bytes.len(),
            prelude.payload_offset as usize
                + header.asset_table[0].size as usize
                + header.asset_table[1].size as usize
        );
    }
}