#ifndef APU_H_
#define APU_H_

#include "nes.h"
#include "apu.h"
#include "blip-buf/blip_buf.h"


#define apu_gain_square     (563U)
#define apu_gain_triangle   (723U)
#define apu_gain_noise      (482U)
#define apu_gain_dmc        (240U)



static const uint8_t apu_length_lut[32] = {
    0x0A, 0xFE, 0x14, 0x02, 0x28, 0x04, 0x50, 0x06,
    0xA0, 0x08, 0x3C, 0x0A, 0x0E, 0x0C, 0x1A, 0x0E,
    0x0C, 0x10, 0x18, 0x12, 0x30, 0x14, 0x60, 0x16,
    0xC0, 0x18, 0x48, 0x1A, 0x10, 0x1C, 0x20, 0x1E,
};

static const uint16_t noise_period_lut[16] = {
      4,   8,  16,  32,  64,   96,  128,  160,
    202, 254, 380, 508, 762, 1016, 2034, 4068,
};

static const uint16_t dmc_period_lut[16] = {
    428, 380, 340, 320, 286, 254, 226, 214,
    190, 160, 142, 128, 106,  84,  72,  54,
};

static const void dmc_restart(nes_apu_Channel* channel) {
    channel->length = (channel->reg[3] * 16) + 1;
    channel->addr = 0x4000U | ((uint16_t)channel->reg[2] << 6);
}

static const void nes_apu_dmc_fetch(nes_apu* apu,
                                    nes_apu_Channel* channel) {
    channel->data = apu->mem_read(apu->arg_mem,
                                  0x8000U | channel->addr);
    channel->mask = 1;
    channel->addr = (channel->addr + 1) & 0x7FFFU;
    channel->length--;
    if (channel->length <= 0) {
        if (channel->reg[0] & apu_DMC_Loop) {
            dmc_restart(channel);
        } else if (channel->reg[0] & apu_DMC_IRQ_Enable) {
            channel->interrupt = 1;
        }
    }
}

static const void nes_apu_dmc_start(nes_apu* apu,
                                    nes_apu_Channel* channel) {
    dmc_restart(channel);
    nes_apu_dmc_fetch(apu, channel);
}

static inline int apu_channel_raw_timer_value(
        nes_apu_Channel *channel) {
    return (    channel->reg[2] |
                (((uint16_t)channel->reg[3] & 0x7U) << 8));
}

static inline void apu_channel_update_timer(
        nes_apu_Channel *channel) {
    channel->period = apu_channel_raw_timer_value(channel) *
                      nes_clock_cpu_div;
}

static void nes_apu_write_square(nes_apu_Channel* channel,
                                 int reg, uint8_t val) {
    APU_LOG("APU: Square %d < %02x\n", reg, val);

    if (0 == reg) {
        channel->env_period = (val & apu_Envelope_Volume);

    } else if (1 == reg) {
        channel->sweep_period = (
                ((val & apu_Square_Period) >> 4)
            );
        channel->flags |= apu_Channel_Reload;

    } else if (2 <= reg) {
        channel->reg[reg] = val;
        channel->period = (
                (apu_channel_raw_timer_value(channel) /*+ 1*/) * 2
            ) * nes_clock_cpu_div;
        if (3 == reg) {
//            channel->flags |= apu_Channel_Reload;
            channel->delay = 0;
            // Envelope and step are already restarted
        }
    }
}

static void nes_apu_write_triangle(nes_apu_Channel* channel,
                                   int reg, uint8_t val) {
    APU_LOG("APU: Triangle %d < %02x\n", reg, val);

    if (2 <= reg) {
        channel->reg[reg] = val;
        apu_channel_update_timer(channel);
        if (3 == reg) channel->flags |= apu_Channel_Reload;
    }
}

static void nes_apu_write_noise(nes_apu_Channel* channel,
                                int reg, uint8_t val) {
    APU_LOG("APU: Noise %d < %02x\n", reg, val);

    if (reg == 0) {
        channel->env_period = (val & apu_Envelope_Volume);

    } else if (reg == 2) {
        channel->period = noise_period_lut[val & apu_Noise_Period] *
                          nes_clock_cpu_div;
    }
}

static void nes_apu_write_dmc(nes_apu_Channel* channel,
                              int reg, uint8_t val) {
    APU_LOG("APU: DMC %d < %02x\n", reg, val);

    if (reg == 0) {
        channel->period = dmc_period_lut[val & apu_DMC_Period] *
                          nes_clock_cpu_div;
        if (!(val & apu_DMC_IRQ_Enable)) {
            channel->interrupt = 0;
        }

    } else if (reg == 1) {
        channel->sample = (val & 0x7FU);
    }
}

int nes_apu_init(nes_apu* apu, int clock, int frequency,
                 uint8_t(*mem_read)(void*, uint16_t),
                 void* arg_mem) {
    int ret = 0;

    // 20 ms buffer
    apu->blip = blip_new(frequency / 50);

    if (NULL == apu->blip) {
        APU_ERR("APU: Failed to create resampler\n");
        ret = -1;

    } else {
        blip_set_rates(apu->blip, clock, frequency);

        nes_apu_reset(apu);

        apu->mem_read = mem_read;
        apu->arg_mem = arg_mem;

        apu->channels[apu_Channel_Square_0].write = nes_apu_write_square;
        apu->channels[apu_Channel_Square_1].write = nes_apu_write_square;
        apu->channels[apu_Channel_Triangle].write = nes_apu_write_triangle;
        apu->channels[apu_Channel_Noise].write    = nes_apu_write_noise;
        apu->channels[apu_Channel_DMC].write      = nes_apu_write_dmc;

        apu->channels[apu_Channel_Square_0].gain = apu_gain_square;
        apu->channels[apu_Channel_Square_1].gain = apu_gain_square;
        apu->channels[apu_Channel_Triangle].gain = apu_gain_triangle;
        apu->channels[apu_Channel_Noise].gain    = apu_gain_noise;
        apu->channels[apu_Channel_DMC].gain      = apu_gain_dmc;

        apu->channels[apu_Channel_Noise].lfsr = 0x7FFFU;

        apu->frame_period = (clock / 240) +
                            (nes_clock_cpu_div / 2);
        apu->frame_delay = apu->frame_period;
    }

    return ret;
}

void nes_apu_done(nes_apu* apu) {
    blip_delete(apu->blip);
}

void nes_apu_reset(nes_apu* apu) {
    apu->status = 0;

    for (int chan = 0; chan < nes_apu_chan_count; ++chan) {
        apu->channels[chan].length = 0;
    }

    apu->channels[apu_Channel_Triangle].step = 31;
    apu->channels[apu_Channel_DMC].reg[1] &= 1;
    apu->channels[apu_Channel_DMC].sample = 0;
    apu->channels[apu_Channel_DMC].mask = 0;
}

uint8_t nes_apu_read(nes_apu* apu, uint16_t addr) {
    uint8_t val = 0;

    if (addr == nes_apu_reg_status) {
        for (int chan = 0; chan < nes_apu_chan_count; ++chan) {
            if (apu->channels[chan].length > 0) {
                val |= (1 << chan);
            }
        }

        if (apu->channels[apu_Channel_DMC].interrupt) {
            val |= apu_Status_DMC_Int;
        }

        if (apu->status & apu_Status_Frame_Int) {
            val |= apu_Status_Frame_Int;
            apu->status &= ~apu_Status_Frame_Int;
        }

        APU_LOG("APU: Status %02x @ %d\n", val, apu->frame_time_elapsed);
    }

    return val;
}

static inline void nes_apu_clock_length(nes_apu_Channel* channel,
                                        uint8_t halt_mask) {
    if (    channel->length > 0 &&
            !(channel->reg[0] & halt_mask)) {
        channel->length--;
        APU_LOG("APU: Clock Length %p -> %d\n", channel, channel->length);
    }
}

static inline void nes_apu_clock_sweep(nes_apu_Channel* channel,
                                       int adjust) {
    int decrement = 1;

    if (channel->flags & apu_Channel_Reload) {
        channel->flags &= ~apu_Channel_Reload;
        decrement = 0;
    }

    if (channel->sweep_delay == 0) {
        decrement = 0;
        if (channel->period < 8 * nes_clock_cpu_div) {
            channel->period = 0;
            channel->length = 0;

        } else if (    (channel->reg[1] & apu_Square_Enable) &&
                0 != (channel->reg[1] & apu_Square_Shift)) {
            int delta = (
                    apu_channel_raw_timer_value(channel) >>
                    (channel->reg[1] & apu_Square_Shift)
                );
            if (channel->reg[1] & apu_Square_Negate) {
                delta = adjust - delta;
            }
            channel->period += delta * nes_clock_apu_div;
            if ( channel < 0 ||
                 channel->period > 0x7FFU * nes_clock_cpu_div) {
                channel->period = 0;
                channel->length = 0;
            }
        }
    }

    if (decrement) {
        channel->sweep_delay--;
    } else {
        channel->sweep_delay = channel->sweep_period;
    }
}

static inline void nes_apu_clock_envelope(nes_apu_Channel* channel) {
    if (channel->flags & apu_Channel_Start) {
        channel->flags &= ~apu_Channel_Start;
        channel->envelope = 15;
        channel->env_delay = channel->env_period;

    } else if (channel->env_delay <= 0) {
        channel->env_delay = channel->env_period;
        if (    channel->envelope > 0 ||
                (channel->reg[0] & apu_Envelope_Halt)) {
            channel->envelope = ((channel->envelope - 1) & 0xFU);
        }

    } else {
        channel->env_delay--;
    }
}

static inline void nes_apu_clock_linear(nes_apu_Channel* channel) {
    if (channel->flags & apu_Channel_Reload) {
        channel->counter = (channel->reg[0] & apu_Triangle_Count);
    } else if (channel->counter > 0) {
        channel->counter--;
    }
    if (!(channel->reg[0] & apu_Triangle_Halt)) {
        channel->flags &= ~apu_Channel_Reload;
    }
}

static void nes_apu_clock_quarter_frame(nes_apu* apu) {
    nes_apu_clock_envelope(&apu->channels[apu_Channel_Square_0]);
    nes_apu_clock_envelope(&apu->channels[apu_Channel_Square_1]);
    nes_apu_clock_envelope(&apu->channels[apu_Channel_Noise]);
    nes_apu_clock_linear(&apu->channels[apu_Channel_Triangle]);
}

static void nes_apu_clock_half_frame(nes_apu* apu) {
    nes_apu_clock_length(&apu->channels[apu_Channel_Square_0], apu_Envelope_Halt);
    nes_apu_clock_length(&apu->channels[apu_Channel_Square_1], apu_Envelope_Halt);
    nes_apu_clock_length(&apu->channels[apu_Channel_Triangle], apu_Triangle_Halt);
    nes_apu_clock_length(&apu->channels[apu_Channel_Noise], apu_Envelope_Halt);
    nes_apu_clock_sweep(&apu->channels[apu_Channel_Square_0], -1);
    nes_apu_clock_sweep(&apu->channels[apu_Channel_Square_1], 0);
}

void nes_apu_write(nes_apu* apu, uint16_t addr, uint8_t val) {
    if (addr < nes_apu_reg_base +
               (nes_apu_chan_count * nes_apu_chan_size)) {
        int chan = (addr - nes_apu_reg_base) / nes_apu_chan_size;
        int reg = (addr - nes_apu_reg_base) % nes_apu_chan_size;

        nes_apu_Channel* channel = &apu->channels[chan];
        if (3 == reg) {
            if (    chan != apu_Channel_DMC &&
                    (apu->status & (1 << chan))) {
                channel->length = apu_length_lut[val >> 3];
                channel->step = (uint8_t)-1;
                channel->flags |= apu_Channel_Start;
            }
        }
        channel->write(channel, reg, val);
        channel->reg[reg] = val;

    } else if (addr == nes_apu_reg_status) {
        APU_LOG("APU: Enable %02x\n", val);

        // Clear this now in case nes_apu_dmc_start resets it
        apu->channels[apu_Channel_DMC].interrupt = 0;

        for (int chan = 0; chan < nes_apu_chan_count; ++chan) {
            if (!(val & (1 << chan))) {
                apu->channels[chan].length = 0;
            }
        }

        if (!(val & (1 << apu_Channel_DMC))) {
            apu->channels[apu_Channel_DMC].mask = 0;

        } else if ( (val & (1 << apu_Channel_DMC)) &&
                    apu->channels[apu_Channel_DMC].length <= 0) {
            nes_apu_dmc_start(apu,
                              &apu->channels[apu_Channel_DMC]);
        }

        apu->status &= ~((1 << nes_apu_chan_count) - 1);
        apu->status |= (val & ((1 << nes_apu_chan_count) - 1));

    } else if (addr == nes_apu_reg_frame) {
        APU_LOG("APU: Frame %02x\n", val);

        apu->frame_reg = val &
                         (apu_Frame_Mode | apu_Frame_Inhibit);

        apu->frame = 0;
        apu->frame_time_elapsed = 0;
        apu->frame_delay = apu->frame_period;

        if (val & apu_Frame_Mode) {
            nes_apu_clock_quarter_frame(apu);
            nes_apu_clock_half_frame(apu);
        }

        if (val & apu_Frame_Inhibit) {
            apu->status &= ~apu_Status_Frame_Int;
        }
    }
}

static int channel_update(nes_apu_Channel* channel) {
    int output = (channel->gain * channel->sample);
    int delta = output - channel->output;
    channel->output = output;
    return delta;
}

static void nes_apu_run_dmc(nes_apu* apu,
                            nes_apu_Channel* channel,
                            int cycles) {
    if (channel->length <= 0 || channel->period <= 0) return;

    int time = apu->time;
    int delta = channel_update(channel);
    if (delta) {
        blip_add_delta(apu->blip, time, delta);
    }

    while (cycles > 0) {
        int run = cycles;
        if (run > channel->delay) {
            run = channel->delay;
        }

        channel->delay -= run;
        time += run;

        if (channel->delay <= 0) {
            channel->delay = channel->period;

            delta = (channel->mask & channel->data) ? 2 : -2;
            int sample = channel->sample + delta;
            if ((unsigned)sample <= 0x7FU) {
                channel->sample = sample;
                delta = channel_update(channel);
                if (delta) {
                    blip_add_delta(apu->blip, time, delta);
                }
            }

            channel->mask <<= 1;
            if (!channel->mask) {
                nes_apu_dmc_fetch(apu, channel);
            }
        }

        cycles -= run;
    }
}

static inline int apu_envelope_volume(nes_apu_Channel* channel) {
    return (    (channel->reg[0] & apu_Envelope_Constant) ?
                (channel->reg[0] & apu_Envelope_Volume) :
                channel->envelope );
}

static void nes_apu_run_noise(nes_apu* apu,
                              nes_apu_Channel* channel,
                              int cycles) {
    if (channel->length <= 0 || channel->period <= 0) return;

    int time = apu->time;

    while (cycles > 0) {
        int run = cycles;
        if (run > channel->delay) {
            run = channel->delay;
        }

        channel->delay -= run;
        time += run;

        if (channel->delay <= 0) {
            channel->delay = channel->period;

            uint16_t feedback = (1 & (channel->lfsr ^ (
                    channel->lfsr >> (
                            (channel->reg[2] & apu_Noise_Mode) ?
                            6 : 1
                        )
                    )
                ));
            channel->lfsr = (feedback << 14) |
                            (channel->lfsr >> 1);
            channel->sample = (channel->lfsr & 1) ?
                              apu_envelope_volume(channel) : 0;
            int delta = channel_update(channel);
            if (delta) {
                blip_add_delta(apu->blip, time, delta);
            }
        }

        cycles -= run;
    }
}

static uint8_t triangle_steps[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,
};

static void nes_apu_run_triangle(nes_apu* apu,
                                 nes_apu_Channel* channel,
                                 int cycles) {
    if (    channel->length <= 0 ||
            channel->period <= 0 ||
            channel->counter <= 0) {
        return;
    }

    int time = apu->time;

    while (cycles > 0) {
        int run = cycles;
        if (run > channel->delay) {
            run = channel->delay;
        }

        channel->delay -= run;
        time += run;

        if (channel->delay <= 0) {
            channel->delay = channel->period;

            channel->step = ((channel->step + 1) & 31U);

            channel->sample = triangle_steps[channel->step];

            int delta = channel_update(channel);
            if (delta) {
                blip_add_delta(apu->blip, time, delta);
            }
        }

        cycles -= run;
    }
}

static const uint8_t square_sequence[4] = {
    0b00000010,
    0b00000110,
    0b00011110,
    0b11111001,
};

static void nes_apu_run_square(nes_apu* apu,
                               nes_apu_Channel* channel,
                               int cycles) {
    if (    channel->length <= 0 ||
            channel->period <= 0) {
        return;
    }

    int time = apu->time;

    while (cycles > 0) {
        int run = cycles;
        if (run > channel->delay) {
            run = channel->delay;
        }

        channel->delay -= run;
        time += run;

        if (channel->delay <= 0) {
            channel->delay = channel->period;

            channel->step = ((channel->step + 1) & 7U);

            channel->sample = ((
                        square_sequence[channel->reg[0] >> 6] &
                        (1 << channel->step)
                    ) ? apu_envelope_volume(channel) : 0);

            int delta = channel_update(channel);
            if (delta) {
                blip_add_delta(apu->blip, time, delta);
            }
        }

        cycles -= run;
    }
}

nes_apu_Result nes_apu_run(nes_apu* apu, int cycles) {
    nes_apu_run_dmc(apu, &apu->channels[apu_Channel_DMC], cycles);

    while (cycles > 0) {
        int run = cycles;
        if (run > apu->frame_delay) {
            run = apu->frame_delay;
        }

        apu->frame_time_elapsed += run;

        nes_apu_run_square(
                apu, &apu->channels[apu_Channel_Square_0], run
            );
        nes_apu_run_square(
                apu, &apu->channels[apu_Channel_Square_1], run
            );
        nes_apu_run_triangle(
                apu, &apu->channels[apu_Channel_Triangle], run
            );
        nes_apu_run_noise(
                apu, &apu->channels[apu_Channel_Noise], run
            );

        apu->frame_delay -= run;

        if (apu->frame_delay <= 0) {
            APU_LOG("APU: End of quarter frame: %d\n", apu->frame_time_elapsed);

            int end = 0;
            int quarter_frame = 1;
            int half_frame = 0;

            apu->frame_delay += apu->frame_period;

            if (1 == apu->frame) {
                half_frame = 1;
            }

            if (apu->frame_reg & apu_Frame_Mode) {
                if (3 == apu->frame) {
                    quarter_frame = 0;
                } else if (4 <= apu->frame) {
                    half_frame = 1;
                    end = 1;
                }
            } else {
                if (3 <= apu->frame) {
                    half_frame = 1;
                    end = 1;
                }
            }

            if (half_frame) {
                nes_apu_clock_half_frame(apu);
            }

            if (quarter_frame) {
                nes_apu_clock_quarter_frame(apu);
            }

            if (end) {
                if (0 == apu->frame_reg) {
                    apu->status |= apu_Status_Frame_Int;
                    APU_LOG("APU: Frame Interrupt @ %d\n", apu->frame_time_elapsed);
                }
                apu->frame = 0;
                apu->frame_time_elapsed = 0;
            } else {
                apu->frame++;
            }
        }

        cycles -= run;
        apu->time += run;
    }

    return (    (apu->status & apu_Status_Frame_Int) ||
                apu->channels[apu_Channel_DMC].interrupt ) ?
            apu_Result_IRQ : apu_Result_Running;
}

#endif // APU_H_