use super::types::{Apu, ApuStateTail, ChannelOutputs};

impl Apu {
    pub fn new() -> Self {
        Self {
            io: [0; 0x20],
            frame_cycle: 0,
            frame_mode_5step: false,
            frame_irq_inhibit: false,
            frame_irq_pending: false,
            channel_enable_mask: 0,
            length_counters: [0; 4],
            dmc_bytes_remaining: 0,
            dmc_irq_enabled: false,
            dmc_irq_pending: false,
            dmc_cycle_counter: 0,
            dmc_current_addr: 0xC000,
            dmc_sample_buffer: 0,
            dmc_sample_buffer_valid: false,
            dmc_shift_reg: 0,
            dmc_bits_remaining: 8,
            dmc_silence: true,
            dmc_output_level: 0,
            dmc_dma_request: false,
            envelope_divider: [0; 3],
            envelope_decay: [0; 3],
            envelope_start_flags: 0,
            triangle_linear_counter: 0,
            triangle_linear_reload_flag: false,
            sweep_divider: [0; 2],
            sweep_reload_flags: 0,
            cpu_cycle_parity: false,
            frame_reset_pending: false,
            frame_reset_delay: 0,
            pending_frame_mode_5step: false,
            pending_frame_irq_inhibit: false,
            pulse_timer_counter: [0; 2],
            pulse_duty_step: [0; 2],
            triangle_timer_counter: 0,
            triangle_step: 0,
            noise_timer_counter: 0,
            noise_lfsr: 1, // LFSR initialized to 1 per NES hardware
        }
    }

    pub fn registers(&self) -> &[u8; 0x20] {
        &self.io
    }

    pub fn set_registers(&mut self, regs: [u8; 0x20]) {
        self.io = regs;
    }

    pub fn read(&mut self, addr: u16) -> u8 {
        match addr {
            0x4000..=0x4013 => self.io[(addr as usize) - 0x4000],
            0x4015 => self.read_status(),
            _ => 0,
        }
    }

    pub fn write(&mut self, addr: u16, value: u8) {
        match addr {
            0x4000..=0x4013 => {
                self.io[(addr as usize) - 0x4000] = value;
                match addr {
                    0x4003 => {
                        self.reload_length_counter(0, value >> 3);
                        self.envelope_start_flags |= 1 << 0;
                        self.pulse_duty_step[0] = 0;
                        self.pulse_timer_counter[0] = self.pulse_timer_period(0x02);
                    }
                    0x4001 => {
                        self.sweep_reload_flags |= 1 << 0;
                    }
                    0x4007 => {
                        self.reload_length_counter(1, value >> 3);
                        self.envelope_start_flags |= 1 << 1;
                        self.pulse_duty_step[1] = 0;
                        self.pulse_timer_counter[1] = self.pulse_timer_period(0x06);
                    }
                    0x4005 => {
                        self.sweep_reload_flags |= 1 << 1;
                    }
                    0x400B => {
                        self.reload_length_counter(2, value >> 3);
                        self.triangle_linear_reload_flag = true;
                    }
                    0x400F => {
                        self.reload_length_counter(3, value >> 3);
                        self.envelope_start_flags |= 1 << 2;
                    }
                    0x4010 => {
                        self.dmc_irq_enabled = (value & 0x80) != 0;
                        if !self.dmc_irq_enabled {
                            self.dmc_irq_pending = false;
                        }
                    }
                    0x4011 => {
                        self.dmc_output_level = value & 0x7F;
                    }
                    0x4012 => {
                        self.dmc_current_addr = self.dmc_sample_start_addr();
                    }
                    _ => {}
                }
            }
            0x4014 => {
                self.io[0x14] = value;
            }
            0x4015 => {
                self.io[(addr as usize) - 0x4000] = value;
                self.channel_enable_mask = value & 0x1F;
                self.dmc_irq_pending = false;
                for chan in 0..4usize {
                    if (self.channel_enable_mask & (1 << chan)) == 0 {
                        self.length_counters[chan] = 0;
                    }
                }
                if (self.channel_enable_mask & 0x10) == 0 {
                    self.dmc_bytes_remaining = 0;
                    self.dmc_cycle_counter = 0;
                    self.dmc_dma_request = false;
                } else if self.dmc_bytes_remaining == 0 {
                    self.dmc_bytes_remaining = self.dmc_sample_length_bytes();
                    self.dmc_cycle_counter = self.dmc_byte_period();
                    self.dmc_current_addr = self.dmc_sample_start_addr();
                    if !self.dmc_sample_buffer_valid {
                        self.dmc_dma_request = true;
                    }
                }
            }
            0x4017 => {
                self.io[(addr as usize) - 0x4000] = value;
                self.pending_frame_mode_5step = (value & 0x80) != 0;
                self.pending_frame_irq_inhibit = (value & 0x40) != 0;
                self.frame_reset_delay = if self.cpu_cycle_parity { 4 } else { 3 };
                self.frame_reset_pending = true;
                if self.pending_frame_irq_inhibit {
                    self.frame_irq_pending = false;
                }
            }
            _ => {}
        }
    }

    pub fn clock_cpu_cycle(&mut self) {
        let mut skip_frame_counter_clock = false;
        if self.frame_reset_pending {
            self.frame_reset_delay = self.frame_reset_delay.saturating_sub(1);
            if self.frame_reset_delay == 0 {
                self.frame_reset_pending = false;
                self.frame_mode_5step = self.pending_frame_mode_5step;
                self.frame_irq_inhibit = self.pending_frame_irq_inhibit;
                self.frame_cycle = 0;
                if self.frame_mode_5step {
                    self.clock_quarter_frame();
                    self.clock_half_frame();
                }
                skip_frame_counter_clock = true;
            }
        }
        if !skip_frame_counter_clock {
            self.clock_frame_counter();
        }
        self.clock_dmc();
        self.clock_pulse_timers();
        self.clock_triangle_timer();
        self.clock_noise_timer();
        self.cpu_cycle_parity = !self.cpu_cycle_parity;
    }

    pub fn poll_irq(&self) -> bool {
        self.dmc_irq_pending || (self.frame_irq_pending && !self.frame_irq_inhibit)
    }

    pub fn take_dmc_dma_request(&mut self) -> Option<u16> {
        if self.dmc_dma_request {
            Some(self.dmc_current_addr)
        } else {
            None
        }
    }

    pub fn provide_dmc_dma_byte(&mut self, byte: u8) {
        if !self.dmc_dma_request {
            return;
        }
        self.dmc_dma_request = false;
        self.dmc_sample_buffer = byte;
        self.dmc_sample_buffer_valid = true;
        self.dmc_current_addr = if self.dmc_current_addr == 0xFFFF {
            0x8000
        } else {
            self.dmc_current_addr.wrapping_add(1)
        };

        if self.dmc_bytes_remaining > 0 {
            self.dmc_bytes_remaining -= 1;
        }
        if self.dmc_bytes_remaining == 0 {
            if (self.io[0x10] & 0x40) != 0 {
                // Loop mode: reset address and byte counter.
                // Do NOT request another DMA here — the sample buffer is full
                // right now. clock_dmc will request the next fetch when the
                // output unit empties the buffer into the shift register, which
                // is the correct NES hardware behaviour (reader only fills an
                // empty buffer).  Requesting early would overwrite the valid
                // buffer and skip the last byte of each loop iteration.
                self.dmc_bytes_remaining = self.dmc_sample_length_bytes();
                self.dmc_current_addr = self.dmc_sample_start_addr();
            } else if self.dmc_irq_enabled {
                self.dmc_irq_pending = true;
            }
        }
    }

    pub fn save_state_tail(&self, out: &mut Vec<u8>) {
        out.extend_from_slice(&self.frame_cycle.to_le_bytes());
        out.push(u8::from(self.frame_mode_5step));
        out.push(u8::from(self.frame_irq_inhibit));
        out.push(u8::from(self.frame_irq_pending));
        out.push(self.channel_enable_mask);
        out.extend_from_slice(&self.length_counters);
        out.extend_from_slice(&self.dmc_bytes_remaining.to_le_bytes());
        out.push(u8::from(self.dmc_irq_enabled));
        out.push(u8::from(self.dmc_irq_pending));
        out.extend_from_slice(&self.dmc_cycle_counter.to_le_bytes());
        out.extend_from_slice(&self.dmc_current_addr.to_le_bytes());
        out.push(self.dmc_sample_buffer);
        out.push(u8::from(self.dmc_sample_buffer_valid));
        out.push(self.dmc_shift_reg);
        out.push(self.dmc_bits_remaining);
        out.push(u8::from(self.dmc_silence));
        out.push(self.dmc_output_level);
        out.push(u8::from(self.dmc_dma_request));
        out.extend_from_slice(&self.envelope_divider);
        out.extend_from_slice(&self.envelope_decay);
        out.push(self.envelope_start_flags);
        out.push(self.triangle_linear_counter);
        out.push(u8::from(self.triangle_linear_reload_flag));
        out.extend_from_slice(&self.sweep_divider);
        out.push(self.sweep_reload_flags);
        out.push(u8::from(self.cpu_cycle_parity));
        out.push(u8::from(self.frame_reset_pending));
        out.push(self.frame_reset_delay);
        out.push(u8::from(self.pending_frame_mode_5step));
        out.push(u8::from(self.pending_frame_irq_inhibit));
        out.extend_from_slice(&self.pulse_timer_counter[0].to_le_bytes());
        out.extend_from_slice(&self.pulse_timer_counter[1].to_le_bytes());
        out.extend_from_slice(&self.pulse_duty_step);
        out.extend_from_slice(&self.triangle_timer_counter.to_le_bytes());
        out.push(self.triangle_step);
        out.extend_from_slice(&self.noise_timer_counter.to_le_bytes());
        out.extend_from_slice(&self.noise_lfsr.to_le_bytes());
    }

    pub fn load_state_tail(&mut self, state: ApuStateTail) {
        self.frame_cycle = state.frame_cycle;
        self.frame_mode_5step = state.frame_mode_5step;
        self.frame_irq_inhibit = state.frame_irq_inhibit;
        self.frame_irq_pending = state.frame_irq_pending;
        self.channel_enable_mask = state.channel_enable_mask;
        self.length_counters = state.length_counters;
        self.dmc_bytes_remaining = state.dmc_bytes_remaining;
        self.dmc_irq_enabled = state.dmc_irq_enabled;
        self.dmc_irq_pending = state.dmc_irq_pending;
        self.dmc_cycle_counter = state.dmc_cycle_counter;
        self.dmc_current_addr = state.dmc_current_addr;
        self.dmc_sample_buffer = state.dmc_sample_buffer;
        self.dmc_sample_buffer_valid = state.dmc_sample_buffer_valid;
        self.dmc_shift_reg = state.dmc_shift_reg;
        self.dmc_bits_remaining = state.dmc_bits_remaining.max(1);
        self.dmc_silence = state.dmc_silence;
        self.dmc_output_level = state.dmc_output_level & 0x7F;
        self.dmc_dma_request = state.dmc_dma_request;
        self.envelope_divider = state.envelope_divider;
        self.envelope_decay = state.envelope_decay;
        self.envelope_start_flags = state.envelope_start_flags & 0x07;
        self.triangle_linear_counter = state.triangle_linear_counter & 0x7F;
        self.triangle_linear_reload_flag = state.triangle_linear_reload_flag;
        self.sweep_divider = state.sweep_divider;
        self.sweep_reload_flags = state.sweep_reload_flags & 0x03;
        self.cpu_cycle_parity = state.cpu_cycle_parity;
        self.frame_reset_pending = state.frame_reset_pending;
        self.frame_reset_delay = state.frame_reset_delay;
        self.pending_frame_mode_5step = state.pending_frame_mode_5step;
        self.pending_frame_irq_inhibit = state.pending_frame_irq_inhibit;
        self.pulse_timer_counter = state.pulse_timer_counter;
        self.pulse_duty_step = [state.pulse_duty_step[0] & 0x07, state.pulse_duty_step[1] & 0x07];
        self.triangle_timer_counter = state.triangle_timer_counter;
        self.triangle_step = state.triangle_step & 0x1F;
        self.noise_timer_counter = state.noise_timer_counter;
        self.noise_lfsr = if state.noise_lfsr == 0 { 1 } else { state.noise_lfsr };
    }

    pub fn channel_outputs(&self) -> ChannelOutputs {
        const PULSE_DUTY_TABLE: [[u8; 8]; 4] = [
            [0, 1, 0, 0, 0, 0, 0, 0],
            [0, 1, 1, 0, 0, 0, 0, 0],
            [0, 1, 1, 1, 1, 0, 0, 0],
            [1, 0, 0, 1, 1, 1, 1, 1],
        ];
        const TRIANGLE_SEQUENCE: [u8; 32] = [
            15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0,
            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
        ];

        let pulse1 = {
            let duty = (self.io[0x00] >> 6) as usize;
            let step = self.pulse_duty_step[0] as usize;
            let volume = if (self.io[0x00] & 0x10) != 0 {
                self.io[0x00] & 0x0F
            } else {
                self.envelope_decay[0]
            };
            let active = (self.channel_enable_mask & 0x01) != 0
                && self.length_counters[0] > 0
                && PULSE_DUTY_TABLE[duty][step] != 0
                && !self.sweep_mutes_channel(0, 0x02);
            if active { volume } else { 0 }
        };

        let pulse2 = {
            let duty = (self.io[0x04] >> 6) as usize;
            let step = self.pulse_duty_step[1] as usize;
            let volume = if (self.io[0x04] & 0x10) != 0 {
                self.io[0x04] & 0x0F
            } else {
                self.envelope_decay[1]
            };
            let active = (self.channel_enable_mask & 0x02) != 0
                && self.length_counters[1] > 0
                && PULSE_DUTY_TABLE[duty][step] != 0
                && !self.sweep_mutes_channel(1, 0x06);
            if active { volume } else { 0 }
        };

        let triangle = {
            // Timer period < 2 produces ultrasonic output (~28-56 kHz) that aliases
            // to audible frequencies when sampled at 48 kHz. Real hardware filters
            // this via the RC output stage; mute here to match that behaviour.
            let active = (self.channel_enable_mask & 0x04) != 0
                && self.length_counters[2] > 0
                && self.triangle_linear_counter > 0
                && self.triangle_timer_period() >= 2;
            if active {
                TRIANGLE_SEQUENCE[self.triangle_step as usize & 0x1F]
            } else {
                0
            }
        };

        let noise = {
            let volume = if (self.io[0x0C] & 0x10) != 0 {
                self.io[0x0C] & 0x0F
            } else {
                self.envelope_decay[2]
            };
            let active = (self.channel_enable_mask & 0x08) != 0
                && self.length_counters[3] > 0
                && (self.noise_lfsr & 1) == 0;
            if active { volume } else { 0 }
        };

        let dmc = self.dmc_output_level;

        ChannelOutputs { pulse1, pulse2, triangle, noise, dmc, expansion: 0.0 }
    }
}