//! Register-mapped XMA context system — a faithful port of xenia-canary's
//! `apu/xma_decoder.cc` context-array + MMIO machinery, MINUS the audio
//! decoder itself (stage 3).
//!
//! The guest allocates XMA contexts via `XMACreateContext` (which hands back a
//! pointer into our 320-entry context array in physical guest memory), writes
//! the 64-byte `XMA_CONTEXT_DATA` struct, then *kicks* decode by writing the
//! per-context bit into the `0x7FEA0000` register aperture. This module
//! satisfies all of that without faulting and records which contexts the guest
//! kicked; stage 3 will consume the recorded `pending` flags to actually
//! produce PCM.
//!
//! ## Byte order
//! The guest accesses the aperture byte-reversed (`stwbrx`/`lwbrx`), so the raw
//! `u32` our MMIO boundary delivers is byte-swapped relative to the logical
//! register value — exactly the situation canary handles with `xe::byte_swap`.
//! So `write_register` swaps the incoming value before decoding and the
//! register file holds host-order values; `read_register` swaps on the way out.
//! This was proven empirically: the guest's Clear writes arrive as
//! `0x01000000`/`0x02000000`/`0x04000000`, i.e. byte-reversed `1`/`2`/`4`,
//! targeting contexts 0/1/2 (which it had just allocated) — NOT 24/25/26. The
//! register-index math (`(addr & 0xFFFF) / 4`) is the same as canary's.

use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};

use xenia_memory::access::MemoryAccess;
use xenia_memory::{GuestMemory, MmioRegion};

use crate::xma_decode::{self, ContextDecodeState, XmaContextData};

/// Size in bytes of an `XMA_CONTEXT_DATA` struct (canary `xma_context.h`).
/// Stage 1 does not decode the fields — only the stride matters.
pub const XMA_CONTEXT_SIZE: u32 = 64;
/// Number of XMA contexts the hardware exposes (canary `kContextCount`).
pub const XMA_CONTEXT_COUNT: usize = 320;

/// Register aperture base (guest physical). Canary maps the XMA decoder at
/// `0x7FEA0000` in `XmaDecoder::Setup`.
pub const APERTURE_BASE: u32 = 0x7FEA_0000;
/// Mask used by `MmioRegion::contains` so any `0x7FEAxxxx` address hits.
pub const APERTURE_MASK: u32 = 0xFFFF_0000;
/// Total aperture size in bytes (the low 16-bit register window).
pub const APERTURE_SIZE: u32 = 0x0001_0000;

// ----- Register indices (canary `XmaRegister` enum / xma_register_table.inc).
// Indices are dword indices: byte offset = index * 4.

/// `ContextArrayAddress` — physical base of the context array. byte 0x1800.
const REG_CONTEXT_ARRAY_ADDRESS: u32 = 0x600;
/// `CurrentContextIndex` — the context the HW is currently servicing. byte
/// 0x1818. Polled by the guest; we rotate it so a poll never sticks.
const REG_CURRENT_CONTEXT_INDEX: u32 = 0x606;

/// First of the 10 `ContextNKick` registers (`Context0Kick`..`Context9Kick`).
/// byte 0x1940. Each register's bit N kicks context `base*32 + N`.
const REG_CONTEXT_KICK_BASE: u32 = 0x650;
/// First of the 10 `ContextNLock` registers. byte 0x1A40.
const REG_CONTEXT_LOCK_BASE: u32 = 0x690;
/// First of the 10 `ContextNClear` registers. byte 0x1A80.
const REG_CONTEXT_CLEAR_BASE: u32 = 0x6A0;
/// Each group spans 10 registers (320 contexts / 32-per-register).
const CONTEXT_GROUP_LEN: u32 = 10;

/// Number of 32-bit words backing the register file. The highest index we
/// touch is `0x6A9`; round up generously so any in-aperture index is in range
/// (64 KB aperture / 4).
const REGISTER_FILE_WORDS: usize = 0x4000;

/// Register-mapped XMA context array. Owns the allocation bitmap, the register
/// file, and the per-context kick/enable bookkeeping that stage 3 consumes.
pub struct XmaDecoder {
    /// Guest virtual address of the context array (handed back by
    /// `allocate_context`).
    context_array_guest_va: u32,
    /// Physical address stored into `ContextArrayAddress` (reg 0x600).
    context_array_phys: u32,
    /// 320-slot allocation bitmap, one bit per context (`bitmap[i>>6]` bit
    /// `i & 63`). A set bit means *allocated*.
    bitmap: [u64; (XMA_CONTEXT_COUNT + 63) / 64],
    /// Flat register file, host-native values. Indexed by dword register index.
    registers: Vec<u32>,
    /// Per-context "decode requested" flag, set on Kick, cleared on Clear.
    /// Stage 3 drains this to produce PCM.
    pending: [bool; XMA_CONTEXT_COUNT],
    /// Per-context enable flag. A Lock disables; a Kick (re-)enables. Mirrors
    /// canary's "is_enabled" notion loosely — exact decode semantics are
    /// stage 3.
    enabled: [bool; XMA_CONTEXT_COUNT],
    /// Total kicks observed (diagnostic; lets headless logs show progress).
    kick_count: u64,
    /// Rotating value served for `CurrentContextIndex` reads so a guest poll
    /// can't spin forever on a fixed value. Atomic so the read path can stay
    /// `&self`.
    current_context_index: AtomicU32,
    /// Per-context stage-3 decode state (FFmpeg codec, staged PCM frame, ring
    /// bookkeeping). Lazily populated as contexts are decoded.
    decode_state: Vec<ContextDecodeState>,
    /// Total PCM bytes written to guest output buffers (diagnostic).
    pcm_bytes_total: u64,
    /// Stable pointer to the guest memory mapping, captured at init. Used to run
    /// `Work()` SYNCHRONOUSLY inside the kick MMIO write — exactly as canary's
    /// default `!use_dedicated_xma_thread` path does (`context.Work()` right in
    /// `WriteRegister`), so the game sees the updated context the instant its
    /// kick store retires. The mapping lives for the whole run; decode is
    /// deterministic and happens on the CPU thread, so this is determinism-safe.
    mem_ptr: *const GuestMemory,
}

// The decoder is owned behind an `Arc<Mutex<..>>` and only ever touched from the
// CPU scheduler thread (kick MMIO writes + the per-round pump). The raw `mem_ptr`
// is a stable whole-run mapping; access is single-threaded.
unsafe impl Send for XmaDecoder {}

impl XmaDecoder {
    /// Construct an un-initialized decoder. Call [`Self::init`] once the
    /// context-array memory has been reserved.
    pub fn new() -> Self {
        Self {
            context_array_guest_va: 0,
            context_array_phys: 0,
            bitmap: [0; (XMA_CONTEXT_COUNT + 63) / 64],
            registers: vec![0; REGISTER_FILE_WORDS],
            pending: [false; XMA_CONTEXT_COUNT],
            enabled: [false; XMA_CONTEXT_COUNT],
            kick_count: 0,
            current_context_index: AtomicU32::new(0),
            decode_state: (0..XMA_CONTEXT_COUNT).map(|_| ContextDecodeState::new()).collect(),
            pcm_bytes_total: 0,
            mem_ptr: std::ptr::null(),
        }
    }

    /// Capture the stable guest-memory mapping so the kick MMIO path can run
    /// `Work()` synchronously (canary semantics). Call once at boot, after the
    /// final `mem` is in its long-lived location.
    pub fn set_memory(&mut self, mem: &GuestMemory) {
        self.mem_ptr = mem as *const GuestMemory;
    }

    /// Wire in the context-array addresses (after the app reserves the buffer)
    /// and publish the physical base into `ContextArrayAddress` (reg 0x600),
    /// exactly as canary's `XmaDecoder::Setup` does.
    pub fn init(&mut self, context_array_guest_va: u32, context_array_phys: u32) {
        self.context_array_guest_va = context_array_guest_va;
        self.context_array_phys = context_array_phys;
        self.registers[REG_CONTEXT_ARRAY_ADDRESS as usize] = context_array_phys;
        tracing::info!(
            va = format_args!("{context_array_guest_va:#010x}"),
            phys = format_args!("{context_array_phys:#010x}"),
            "xma: context array initialized"
        );
    }

    /// Acquire a free context slot and return its guest pointer
    /// (`context_array_guest_va + i*64`), or 0 if all 320 slots are in use.
    /// Mirrors canary's `XmaDecoder::AllocateContext`.
    pub fn allocate_context(&mut self) -> u32 {
        for i in 0..XMA_CONTEXT_COUNT {
            let word = i >> 6;
            let bit = 1u64 << (i & 63);
            if self.bitmap[word] & bit == 0 {
                self.bitmap[word] |= bit;
                let ptr = self.context_array_guest_va + (i as u32) * XMA_CONTEXT_SIZE;
                tracing::info!(
                    index = i,
                    ptr = format_args!("{ptr:#010x}"),
                    "xma: allocate_context"
                );
                return ptr;
            }
        }
        tracing::warn!("xma: allocate_context — all {} slots in use", XMA_CONTEXT_COUNT);
        0
    }

    /// Free the slot backing `guest_ptr`. Mirrors canary's
    /// `XmaDecoder::ReleaseContext`. Out-of-range / unaligned pointers are
    /// ignored (the guest never faults).
    pub fn release_context(&mut self, guest_ptr: u32) {
        if guest_ptr < self.context_array_guest_va {
            return;
        }
        let offset = guest_ptr - self.context_array_guest_va;
        let i = (offset / XMA_CONTEXT_SIZE) as usize;
        if i >= XMA_CONTEXT_COUNT {
            return;
        }
        let word = i >> 6;
        let bit = 1u64 << (i & 63);
        self.bitmap[word] &= !bit;
        self.pending[i] = false;
        self.enabled[i] = false;
        tracing::info!(index = i, ptr = format_args!("{guest_ptr:#010x}"), "xma: release_context");
    }

    /// Read a register. Returns the stored value, except `CurrentContextIndex`
    /// (0x606) which rotates `0..XMA_CONTEXT_COUNT` per read so a polling guest
    /// always sees forward progress. Out-of-range indices read 0.
    pub fn read_register(&self, reg_index: u32) -> u32 {
        // The guest accesses the aperture byte-reversed (`lwbrx`), so the
        // register file holds host-order values and we swap on the way out —
        // exactly as canary's `ReadRegister` returns `xe::byte_swap(reg)`.
        let host = if reg_index == REG_CURRENT_CONTEXT_INDEX {
            // Rotate mod context count on each read so a poll never sticks.
            let prev = self.current_context_index.fetch_add(1, Ordering::Relaxed);
            prev % XMA_CONTEXT_COUNT as u32
        } else {
            self.registers.get(reg_index as usize).copied().unwrap_or(0)
        };
        host.swap_bytes()
    }

    /// Write a register, then apply the side-effect of the Kick / Lock / Clear
    /// register groups. Each register in a group covers 32 contexts; bit N maps
    /// to `context_id = (reg_index - group_base) * 32 + N`. We iterate set bits
    /// with `trailing_zeros` + clear-lowest-bit, mirroring canary's
    /// `std::countr_zero` loop. The incoming value is byte-swapped first (see
    /// below).
    pub fn write_register(&mut self, reg_index: u32, value: u32) {
        // The guest writes the aperture byte-reversed (`stwbrx`); undo it so the
        // register file holds host-order values, mirroring canary's
        // `WriteRegister` which does `value = xe::byte_swap(value)` first. Proven
        // by the guest's Clear writes (`0x01000000` == context 0, not 24).
        let value = value.swap_bytes();
        if let Some(slot) = self.registers.get_mut(reg_index as usize) {
            *slot = value;
        }

        if (REG_CONTEXT_KICK_BASE..REG_CONTEXT_KICK_BASE + CONTEXT_GROUP_LEN).contains(&reg_index) {
            let base = (reg_index - REG_CONTEXT_KICK_BASE) * 32;
            let mut bits = value;
            while bits != 0 {
                let b = bits.trailing_zeros();
                bits &= bits - 1;
                let context_id = (base + b) as usize;
                if context_id < XMA_CONTEXT_COUNT {
                    self.pending[context_id] = true;
                    self.enabled[context_id] = true;
                    self.kick_count += 1;
                    tracing::debug!(
                        context_id,
                        kick_count = self.kick_count,
                        "xma: kick (decode requested)"
                    );
                    // Canary `!use_dedicated_xma_thread`: run Work() right here so
                    // the game observes the updated context when its kick store
                    // retires. Safe — `mem_ptr` is a stable whole-run mapping and
                    // we're on the CPU thread.
                    if !self.mem_ptr.is_null() {
                        let mem: &GuestMemory = unsafe { &*self.mem_ptr };
                        self.enabled[context_id] = false;
                        self.work_one(mem, context_id);
                    }
                }
            }
        } else if (REG_CONTEXT_LOCK_BASE..REG_CONTEXT_LOCK_BASE + CONTEXT_GROUP_LEN)
            .contains(&reg_index)
        {
            let base = (reg_index - REG_CONTEXT_LOCK_BASE) * 32;
            let mut bits = value;
            while bits != 0 {
                let b = bits.trailing_zeros();
                bits &= bits - 1;
                let context_id = (base + b) as usize;
                if context_id < XMA_CONTEXT_COUNT {
                    self.enabled[context_id] = false;
                    tracing::debug!(context_id, "xma: lock (context disabled)");
                }
            }
        } else if (REG_CONTEXT_CLEAR_BASE..REG_CONTEXT_CLEAR_BASE + CONTEXT_GROUP_LEN)
            .contains(&reg_index)
        {
            let base = (reg_index - REG_CONTEXT_CLEAR_BASE) * 32;
            let mut bits = value;
            while bits != 0 {
                let b = bits.trailing_zeros();
                bits &= bits - 1;
                let context_id = (base + b) as usize;
                if context_id < XMA_CONTEXT_COUNT {
                    self.pending[context_id] = false;
                    self.enabled[context_id] = false;
                    tracing::debug!(context_id, "xma: clear (context state reset)");
                }
            }
        }
    }

    /// Total kicks observed so far (diagnostic; stage 3 will consume `pending`).
    pub fn kick_count(&self) -> u64 {
        self.kick_count
    }

    /// Whether context `i` has a pending (un-serviced) kick. Stage-3 hook.
    pub fn is_pending(&self, i: usize) -> bool {
        self.pending.get(i).copied().unwrap_or(false)
    }

    /// Total PCM bytes the decoder has written to guest output buffers.
    pub fn pcm_bytes_total(&self) -> u64 {
        self.pcm_bytes_total
    }

    /// Stage-3 entry point. Called once per scheduler round from the CPU
    /// thread's per-round coordinator. For each context with a pending kick,
    /// run one `Work()` pass (canary `XmaContextNew::Work`): read the context,
    /// decode available input into PCM, drain into the output ring, and write
    /// the decoder-owned fields back. Deterministic — no host thread, no clock.
    pub fn decode_pending(&mut self, mem: &GuestMemory) {
        if self.context_array_guest_va == 0 {
            return;
        }
        for i in 0..XMA_CONTEXT_COUNT {
            if !self.pending[i] || !self.enabled[i] {
                continue;
            }
            // Canary `Work` clears is_enabled at entry; a fresh kick re-enables.
            self.enabled[i] = false;
            self.work_one(mem, i);
        }
    }

    /// One `Work()` pass for context `i`. Faithful to canary's orchestration but
    /// uses the mainline xma2 decoder (whole-packet driven) for the actual
    /// frame decode in place of canary's per-frame `Decode()`.
    fn work_one(&mut self, mem: &GuestMemory, i: usize) {
        let ctx_va = self.context_array_guest_va + (i as u32) * XMA_CONTEXT_SIZE;
        let data = XmaContextData::read(mem, ctx_va);
        let initial = data;

        if data.output_buffer_valid == 0 {
            return;
        }

        let mut data = data;
        self.decode_into_output(mem, i, ctx_va, &mut data, &initial);
    }

    /// Decode available input packets into PCM and drain into the output ring.
    fn decode_into_output(
        &mut self,
        mem: &GuestMemory,
        i: usize,
        ctx_va: u32,
        data: &mut XmaContextData,
        initial: &XmaContextData,
    ) {
        use xma_decode::*;

        let output_capacity = data.output_buffer_block_count * OUTPUT_BYTES_PER_BLOCK;
        if output_capacity == 0 {
            return;
        }
        let out_backing = xma_phys_to_backing(data.output_buffer_ptr);
        let mut write_off = data.output_buffer_write_offset * OUTPUT_BYTES_PER_BLOCK;
        let read_off = data.output_buffer_read_offset * OUTPUT_BYTES_PER_BLOCK;

        // write_count: free space in the ring from write to read.
        let free_bytes = ring_write_count(read_off, write_off, output_capacity);
        self.decode_state[i].remaining_subframe_blocks_in_output =
            (free_bytes / OUTPUT_BYTES_PER_BLOCK) as i32;

        let effective_sdc = data.subframe_decode_count.max(1);
        let min_blocks = effective_sdc as i32 + data.output_buffer_padding as i32;

        if min_blocks > self.decode_state[i].remaining_subframe_blocks_in_output {
            // No room — write back unchanged and wait for the game to drain.
            store_merged_pub(mem, ctx_va, data, initial);
            return;
        }

        let mut produced_any = false;

        // Ensure codec configured for current rate/channels.
        let rate = sample_rate_hz(data.sample_rate);
        let channels = if data.is_stereo != 0 { 2 } else { 1 };
        self.ensure_codec(i, rate, channels);

        // Main decode loop: while there's output ring room and valid input.
        loop {
            if self.decode_state[i].remaining_subframe_blocks_in_output < min_blocks {
                break;
            }

            // If we still have undrained subframes from a prior decode, consume
            // them first (canary Consume before next Decode).
            if self.decode_state[i].current_frame_remaining_subframes == 0 {
                // Need a fresh decoded frame. Pull from the codec, feeding input
                // packets as required.
                if !self.produce_frame(mem, i, data) {
                    break;
                }
            }

            // Consume: write up to `effective_sdc` subframes (256B blocks) of
            // the staged raw_frame into the output ring.
            let total_subframes =
                ((BYTES_PER_FRAME_CHANNEL / OUTPUT_BYTES_PER_BLOCK) << data.is_stereo) as u8;
            let remaining = self.decode_state[i].current_frame_remaining_subframes;
            let to_write = remaining.min(effective_sdc as u8);
            let frame_read_off = (total_subframes - remaining) as usize * OUTPUT_BYTES_PER_BLOCK as usize;
            let nbytes = to_write as u32 * OUTPUT_BYTES_PER_BLOCK;

            // Write into the output ring (handle wrap).
            let raw = &self.decode_state[i].raw_frame;
            write_off = ring_write(
                mem,
                out_backing,
                output_capacity,
                write_off,
                &raw[frame_read_off..frame_read_off + nbytes as usize],
            );
            self.pcm_bytes_total += nbytes as u64;
            produced_any = true;

            let headroom = if remaining - to_write == 0 {
                data.output_buffer_padding as i32
            } else {
                0
            };
            self.decode_state[i].remaining_subframe_blocks_in_output -=
                to_write as i32 + headroom;
            self.decode_state[i].current_frame_remaining_subframes -= to_write;
        }

        // Writeback offsets.
        data.output_buffer_write_offset = write_off / OUTPUT_BYTES_PER_BLOCK;

        if self.decode_state[i].remaining_subframe_blocks_in_output == 0
            && write_off == read_off
        {
            data.output_buffer_valid = 0;
        }
        if !produced_any && !data.is_any_input_buffer_valid() {
            data.output_buffer_valid = 0;
        }

        store_merged_pub(mem, ctx_va, data, initial);
    }

    /// Configure (or reconfigure) the FFmpeg xma2 codec for this context.
    fn ensure_codec(&mut self, i: usize, rate: u32, channels: u32) {
        let st = &mut self.decode_state[i];
        if st.codec.is_some() && st.codec_rate == rate && st.codec_channels == channels {
            return;
        }
        match crate::xma2_codec::Xma2Codec::new(rate, channels) {
            Ok(c) => {
                st.codec = Some(c);
                st.codec_rate = rate;
                st.codec_channels = channels;
                tracing::info!(ctx = i, rate, channels, "xma: xma2 codec configured");
            }
            Err(e) => {
                tracing::error!(ctx = i, rate, channels, error = %e, "xma: xma2 codec init failed");
                st.codec = None;
            }
        }
    }

    /// Produce one decoded 512-sample frame into `raw_frame` (interleaved S16BE).
    ///
    /// Input-consumption model (faithful to canary's packet/buffer contract).
    ///
    /// The mainline xma2 decoder consumes whole 2 KB packets via `send_packet`
    /// and emits frames in bursts (internal FIFO + lookahead), so its intake
    /// position can't be read per-frame. We therefore keep TWO cursors:
    ///
    ///  1. A private FFmpeg *feed* cursor (`feed_buffer`/`feed_packet_index`)
    ///     that hands raw packets to FFmpeg only far enough ahead to keep the
    ///     PCM queue stocked. This follows the same buffer ping-pong as the
    ///     guest but is NOT what the guest observes.
    ///  2. The guest-visible `input_buffer_read_offset`, advanced by exactly
    ///     ONE compressed frame each time we emit a 512-sample frame to the
    ///     guest — via `advance_read_offset_one_frame`, a faithful port of the
    ///     offset arithmetic in canary's `Decode()`. This crosses packet and
    ///     buffer boundaries (and fires SwapInputBuffer, clearing the drained
    ///     buffer's valid bit) at canary's true per-frame cadence, which is
    ///     what the WMV demuxer polls to refill ADV.wmv.
    ///
    /// Decoupling the two means FFmpeg's whole-packet burst framing no longer
    /// freezes the guest-visible offset: the offset now tracks emitted output,
    /// so the input buffer is consumed and swapped as the movie actually plays.
    fn produce_frame(&mut self, mem: &GuestMemory, i: usize, data: &mut XmaContextData) -> bool {
        use xma_decode::*;
        let channels = if data.is_stereo != 0 { 2u32 } else { 1u32 };
        let frame_bytes = (BYTES_PER_FRAME_CHANNEL * channels) as usize;

        // Top up FFmpeg's internal FIFO (and our queue) just enough to satisfy
        // one frame, feeding raw packets via the private feed cursor.
        if self.decode_state[i].pcm_queue.len() < frame_bytes {
            self.feed_codec(mem, i, data);
        }

        // Pop exactly one 512-sample frame from the queue into raw_frame.
        if self.decode_state[i].pcm_queue.len() < frame_bytes {
            return false;
        }
        {
            let st = &mut self.decode_state[i];
            st.raw_frame.iter_mut().for_each(|b| *b = 0);
            for b in st.raw_frame[..frame_bytes].iter_mut() {
                *b = st.pcm_queue.pop_front().unwrap();
            }
            st.current_frame_remaining_subframes = (4u8) << data.is_stereo;
        }

        // We just emitted one frame to the guest — advance its visible read
        // offset by one compressed frame at canary's cadence (may swap buffer).
        self.advance_read_offset_one_frame(mem, data);
        true
    }

    /// Feed raw 2 KB packets to FFmpeg from the private feed cursor until the
    /// PCM queue holds at least one frame or the codec stops accepting input.
    /// The feed cursor follows the guest's `current_buffer` ping-pong but keeps
    /// its own packet index (`feed_packet_index`), so feeding ahead of the
    /// guest-visible read offset is fine — the offset advances separately per
    /// emitted frame.
    fn feed_codec(&mut self, mem: &GuestMemory, i: usize, data: &XmaContextData) {
        use xma_decode::*;
        let channels = if data.is_stereo != 0 { 2u32 } else { 1u32 };
        let frame_bytes = (BYTES_PER_FRAME_CHANNEL * channels) as usize;

        // Re-sync the feed buffer to the guest's current buffer if the guest has
        // swapped past us (the buffer we were feeding was consumed).
        if self.decode_state[i].feed_buffer != data.current_buffer
            && !data.is_input_buffer_valid(self.decode_state[i].feed_buffer)
        {
            self.decode_state[i].feed_buffer = data.current_buffer;
            self.decode_state[i].feed_packet_index = 0;
        }

        const MAX_FEED: u32 = 8;
        let mut fed = 0u32;
        while self.decode_state[i].pcm_queue.len() < frame_bytes && fed < MAX_FEED {
            let fb = self.decode_state[i].feed_buffer;
            if !data.is_input_buffer_valid(fb) {
                // Nothing to feed from this buffer; try the other if valid.
                let other = fb ^ 1;
                if data.is_input_buffer_valid(other) {
                    self.decode_state[i].feed_buffer = other;
                    self.decode_state[i].feed_packet_index = 0;
                    continue;
                }
                break;
            }
            let pkt_count = data.input_buffer_packet_count(fb);
            let pidx = self.decode_state[i].feed_packet_index;
            if pidx >= pkt_count {
                // Exhausted this buffer's packets at the feed cursor; advance to
                // the other buffer if it's valid (it was refilled), else wait.
                let other = fb ^ 1;
                if data.is_input_buffer_valid(other) {
                    self.decode_state[i].feed_buffer = other;
                    self.decode_state[i].feed_packet_index = 0;
                    continue;
                }
                break;
            }
            let backing = xma_phys_to_backing(data.input_buffer_address(fb));
            let pkt_va = backing + pidx * BYTES_PER_PACKET;
            let mut packet = vec![0u8; BYTES_PER_PACKET as usize];
            mem.read_bytes(pkt_va, &mut packet);
            let send_res = match self.decode_state[i].codec.as_mut() {
                Some(codec) => codec.send_packet(&packet),
                None => break,
            };
            match send_res {
                Ok(()) => {
                    self.decode_state[i].feed_packet_index += 1;
                    fed += 1;
                    self.drain_codec_frames(i);
                }
                // Decoder full — drain what it has and stop; re-offer this same
                // packet next time (don't advance the feed cursor).
                Err(ref e) if e == "EAGAIN" => {
                    self.drain_codec_frames(i);
                    break;
                }
                Err(e) => {
                    tracing::warn!(ctx = i, error = %e, "xma: send_packet failed");
                    break;
                }
            }
        }
    }

    /// Pull all currently-available decoded frames from the codec and append
    /// their interleaved S16BE PCM to the context's queue.
    fn drain_codec_frames(&mut self, i: usize) {
        loop {
            let out = match self.decode_state[i].codec.as_mut() {
                Some(c) => c.receive_frame(),
                None => None,
            };
            let Some((nb, bytes)) = out else { break };
            let st = &mut self.decode_state[i];
            st.frames_decoded += 1;
            if !st.first_frame_logged {
                st.first_frame_logged = true;
                tracing::info!(
                    ctx = i,
                    samples = nb,
                    pcm_bytes = bytes.len(),
                    "xma: first PCM frame decoded"
                );
            }
            st.pcm_queue.extend(bytes);
        }
    }

    /// Advance `input_buffer_read_offset` by exactly ONE compressed frame,
    /// faithfully mirroring the offset arithmetic in canary's
    /// `XmaContextNew::Decode` (frame-size parse + packet-boundary handling +
    /// SwapInputBuffer when the buffer's packets are exhausted). Called once per
    /// 512-sample frame we emit to the guest, so the guest-visible read offset
    /// crosses packet/buffer boundaries at canary's true cadence — independent
    /// of the mainline xma2 decoder's whole-packet burst framing. This is what
    /// lets `input_buffer_0_valid` toggle and the WMV demuxer refill ADV.wmv.
    fn advance_read_offset_one_frame(&mut self, mem: &GuestMemory, data: &mut XmaContextData) {
        use xma_decode::*;

        if !data.is_any_input_buffer_valid() {
            return;
        }
        if !data.is_current_input_buffer_valid() {
            self.swap_input_buffer(data);
            if !data.is_current_input_buffer_valid() {
                return;
            }
        }

        // Clamp a header-region offset (canary's Dirt-2 guard).
        if data.input_buffer_read_offset < BITS_PER_PACKET_HEADER {
            data.input_buffer_read_offset = BITS_PER_PACKET_HEADER;
        }

        let pkt_count = data.current_input_buffer_packet_count();
        let input_size = pkt_count * BYTES_PER_PACKET;
        let Some(packet_index) = packet_number(input_size, data.input_buffer_read_offset) else {
            return;
        };
        let buf_backing = xma_phys_to_backing(data.current_input_buffer_address());
        let pkt_va = buf_backing + packet_index * BYTES_PER_PACKET;
        let mut packet = vec![0u8; BYTES_PER_PACKET as usize];
        mem.read_bytes(pkt_va, &mut packet);

        let first_frame_offset = packet_frame_offset(&packet);
        let mut relative_offset = data.input_buffer_read_offset % BITS_PER_PACKET;
        if relative_offset < first_frame_offset {
            // Tail of a split frame — skip to this packet's first frame.
            data.input_buffer_read_offset =
                packet_index * BITS_PER_PACKET + first_frame_offset;
            relative_offset = first_frame_offset;
        }

        let skip_count = packet_skip_count(&packet);
        // Full-packet skip (0xFF): no frames begin here — advance to the next
        // packet that does, swapping the buffer if exhausted.
        if skip_count == 0xFF {
            let next_packet_index = packet_index + 1;
            let next_off =
                self.next_packet_read_offset(mem, data, next_packet_index, pkt_count);
            if next_packet_index >= pkt_count || next_off == BITS_PER_PACKET_HEADER {
                self.swap_input_buffer(data);
            }
            data.input_buffer_read_offset = next_off;
            return;
        }

        let info = get_packet_info(&packet, relative_offset);
        let packet_to_skip = (skip_count as u32) + 1;
        let next_packet_index = packet_index + packet_to_skip;

        // Frame size: clamp to the bits remaining in the packet stream (canary
        // GetAmountOfBitsToRead over the (packet_index+1)*kBitsPerPacket stream).
        let stream_remaining =
            ((packet_index + 1) * BITS_PER_PACKET).saturating_sub(data.input_buffer_read_offset);
        let frame_size = if info.current_frame_size == 0 {
            // Split header we can't resolve from this packet alone; fall back to
            // advancing past the rest of this packet so we don't stall.
            stream_remaining
        } else {
            info.current_frame_size
        };
        let bits_to_copy = amount_of_bits_to_read(stream_remaining, frame_size);

        if !info.is_last_frame_in_packet() {
            let next_frame_offset =
                (data.input_buffer_read_offset + bits_to_copy) % BITS_PER_PACKET;
            data.input_buffer_read_offset =
                packet_index * BITS_PER_PACKET + next_frame_offset;
            return;
        }

        // Last frame in this packet: move to the next packet's first frame, or
        // swap the input buffer if the packets are exhausted (canary's
        // `next_packet_index >= current_input_packet_count`).
        let mut next_off =
            self.next_packet_read_offset(mem, data, next_packet_index, pkt_count);
        if next_packet_index >= pkt_count || next_off == BITS_PER_PACKET_HEADER {
            self.swap_input_buffer(data);
        }
        if next_off == BITS_PER_PACKET_HEADER && data.is_any_input_buffer_valid() {
            // At the start of the next buffer: jump to its first frame offset.
            let nb_backing = xma_phys_to_backing(data.current_input_buffer_address());
            let mut hdr = [0u8; 4];
            mem.read_bytes(nb_backing, &mut hdr);
            let fo = packet_frame_offset(&hdr);
            if fo <= MAX_FRAME_SIZE_IN_BITS {
                next_off = fo;
            }
        }
        data.input_buffer_read_offset = next_off;
    }

    /// Scan forward from `next_packet_index` (possibly into the *next* buffer)
    /// for the next packet that begins a frame and return its bit offset, or
    /// `BITS_PER_PACKET_HEADER` if none (canary `GetNextPacketReadOffset`).
    fn next_packet_read_offset(
        &self,
        mem: &GuestMemory,
        data: &XmaContextData,
        next_packet_index: u32,
        current_input_packet_count: u32,
    ) -> u32 {
        use xma_decode::*;
        // Resolve which buffer the packet lives in (current or the other).
        let (buffer_index, mut pidx) = if next_packet_index >= current_input_packet_count {
            (data.current_buffer ^ 1, next_packet_index - current_input_packet_count)
        } else {
            (data.current_buffer, next_packet_index)
        };
        if !data.is_input_buffer_valid(buffer_index) {
            return BITS_PER_PACKET_HEADER;
        }
        let addr = data.input_buffer_address(buffer_index);
        if addr == 0 {
            return BITS_PER_PACKET_HEADER;
        }
        let pkt_count = data.input_buffer_packet_count(buffer_index);
        let backing = xma_phys_to_backing(addr);
        while pidx < pkt_count {
            let mut hdr = [0u8; 4];
            mem.read_bytes(backing + pidx * BYTES_PER_PACKET, &mut hdr);
            let fo = packet_frame_offset(&hdr);
            if fo <= MAX_FRAME_SIZE_IN_BITS {
                return pidx * BITS_PER_PACKET + fo;
            }
            pidx += 1;
        }
        BITS_PER_PACKET_HEADER
    }

    fn swap_input_buffer(&mut self, data: &mut XmaContextData) {
        use xma_decode::*;
        tracing::debug!(
            from = data.current_buffer,
            to = data.current_buffer ^ 1,
            "xma: SwapInputBuffer (input buffer consumed)"
        );
        if data.current_buffer == 0 {
            data.input_buffer_0_valid = 0;
        } else {
            data.input_buffer_1_valid = 0;
        }
        data.current_buffer ^= 1;
        data.input_buffer_read_offset = BITS_PER_PACKET_HEADER;
    }
}

impl Default for XmaDecoder {
    fn default() -> Self {
        Self::new()
    }
}

/// Build the [`MmioRegion`] for the XMA register aperture at `0x7FEA0000`.
/// Mirrors the GPU's `build_region`: the closures lock the shared decoder,
/// compute the dword register index, and dispatch to `read`/`write_register`.
pub fn build_mmio_region(dec: Arc<Mutex<XmaDecoder>>) -> MmioRegion {
    let read_dec = dec.clone();
    let write_dec = dec;

    MmioRegion {
        base_address: APERTURE_BASE,
        mask: APERTURE_MASK,
        size: APERTURE_SIZE,
        read_callback: Box::new(move |addr: u32| {
            let reg_index = (addr & 0xFFFF) / 4;
            read_dec.lock().unwrap().read_register(reg_index)
        }),
        write_callback: Box::new(move |addr: u32, value: u32| {
            let reg_index = (addr & 0xFFFF) / 4;
            write_dec.lock().unwrap().write_register(reg_index, value);
        }),
    }
}

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

    fn inited() -> XmaDecoder {
        let mut d = XmaDecoder::new();
        // Pick a plausible physical-window VA/phys pair.
        d.init(0xA010_0000, 0x0010_0000);
        d
    }

    /// The guest writes/reads the aperture byte-reversed; `wire(v)` is the raw
    /// bus value the guest sends to mean host-order `v` (and what a read of a
    /// host-order `v` returns). Equivalent to `lwbrx`/`stwbrx` semantics.
    fn wire(v: u32) -> u32 {
        v.swap_bytes()
    }

    /// (a) `allocate_context` hands back distinct, increasing pointers spaced by
    /// the 64-byte stride, exhausts at 320, and `release_context` frees the slot.
    #[test]
    fn allocate_distinct_then_exhaust_then_release() {
        let mut d = inited();
        let first = d.allocate_context();
        let second = d.allocate_context();
        assert_eq!(first, 0xA010_0000);
        assert_eq!(second, 0xA010_0000 + XMA_CONTEXT_SIZE);
        assert!(second > first);

        // Drain the remaining slots (2 already taken).
        for _ in 0..(XMA_CONTEXT_COUNT - 2) {
            assert_ne!(d.allocate_context(), 0);
        }
        // 321st allocation fails.
        assert_eq!(d.allocate_context(), 0);

        // Free the first slot and re-acquire it.
        d.release_context(first);
        assert_eq!(d.allocate_context(), first);
    }

    /// (b) A Kick to `Context0Kick` with host value `0b101` marks contexts 0
    /// and 2. The guest sends it byte-reversed (`wire`).
    #[test]
    fn kick_context0_marks_correct_contexts() {
        let mut d = inited();
        d.write_register(REG_CONTEXT_KICK_BASE, wire(0b101));
        assert!(d.is_pending(0));
        assert!(!d.is_pending(1));
        assert!(d.is_pending(2));
        assert_eq!(d.kick_count(), 2);
    }

    /// (c) A Kick to `Context1Kick` (0x651) bit 0 maps to context_id 32.
    #[test]
    fn kick_context1_bit0_is_context_32() {
        let mut d = inited();
        d.write_register(REG_CONTEXT_KICK_BASE + 1, wire(0b1));
        assert!(d.is_pending(32));
        assert!(!d.is_pending(0));
        assert_eq!(d.kick_count(), 1);
    }

    /// Regression for the byte-order fix: the guest's real Clear writes were
    /// `0x01000000`/`0x02000000`/`0x04000000` (bytes-reversed `1`/`2`/`4`),
    /// meaning contexts 0/1/2 — NOT 24/25/26. Verify the raw bus values decode
    /// to the low contexts.
    #[test]
    fn byte_reversed_clear_targets_low_contexts() {
        let mut d = inited();
        for i in 0..3 {
            d.write_register(REG_CONTEXT_KICK_BASE, wire(1 << i));
        }
        assert!(d.is_pending(0) && d.is_pending(1) && d.is_pending(2));
        // The exact bus values observed from the guest.
        d.write_register(REG_CONTEXT_CLEAR_BASE, 0x0100_0000);
        d.write_register(REG_CONTEXT_CLEAR_BASE, 0x0200_0000);
        d.write_register(REG_CONTEXT_CLEAR_BASE, 0x0400_0000);
        assert!(!d.is_pending(0) && !d.is_pending(1) && !d.is_pending(2));
    }

    /// (d) `read_register(0x600)` returns the base byte-reversed (the guest
    /// `lwbrx`-reverses it back to the host-order base on its side).
    #[test]
    fn context_array_address_reads_phys() {
        let d = inited();
        assert_eq!(
            d.read_register(REG_CONTEXT_ARRAY_ADDRESS),
            wire(0x0010_0000)
        );
    }

    /// (e) `CurrentContextIndex` rotates on each read and wraps at the count
    /// (values returned byte-reversed).
    #[test]
    fn current_context_index_rotates() {
        let d = inited();
        assert_eq!(d.read_register(REG_CURRENT_CONTEXT_INDEX), wire(0));
        assert_eq!(d.read_register(REG_CURRENT_CONTEXT_INDEX), wire(1));
        assert_eq!(d.read_register(REG_CURRENT_CONTEXT_INDEX), wire(2));
        // Advance to the wrap boundary.
        for _ in 3..XMA_CONTEXT_COUNT as u32 {
            d.read_register(REG_CURRENT_CONTEXT_INDEX);
        }
        // Next read wraps back to 0.
        assert_eq!(d.read_register(REG_CURRENT_CONTEXT_INDEX), wire(0));
    }

    /// Clear must drop a previously-kicked pending flag.
    #[test]
    fn clear_resets_pending() {
        let mut d = inited();
        d.write_register(REG_CONTEXT_KICK_BASE, wire(0b1));
        assert!(d.is_pending(0));
        d.write_register(REG_CONTEXT_CLEAR_BASE, wire(0b1));
        assert!(!d.is_pending(0));
    }

    /// The MMIO region routes a guest write at `BASE + 0x600*4` to reg 0x600
    /// and a read back through the same byte address, applying the byte swap.
    #[test]
    fn mmio_region_round_trips_register() {
        let dec = Arc::new(Mutex::new(inited()));
        let region = build_mmio_region(dec.clone());
        let kick_byte = APERTURE_BASE + REG_CONTEXT_KICK_BASE * 4;
        (region.write_callback)(kick_byte, wire(0b1));
        assert!(dec.lock().unwrap().is_pending(0));
        // ContextArrayAddress read-back via the bus (byte-reversed).
        let addr_byte = APERTURE_BASE + REG_CONTEXT_ARRAY_ADDRESS * 4;
        assert_eq!((region.read_callback)(addr_byte), wire(0x0010_0000));
    }
}
