#include "mixer.h"
#include "arena.h"
#include "scratch.h"
#include "sound.h"
#include "sys.h"
#include "math.h"

/* TODO: Cap max sounds playing. */

/* Terminology:
 *
 * `Sample`: Once "PCM" data point representing the smallest unit of audio available for a single channel at a point in time.
 * Examples:
 *   - Single 32 bit float output by mixer and consumed by playback API, that the API interprets as a sound sample for a single channel
 *   - Single 16 bit integer output by audio file decoder, that may represent a mono sound sample
 *
 * `Frame`: Represents a single data point of audio for all audio channels at a point in time.
 *  Examples:
 *   - Single 16 bit integer output by audio file decoder representing one mono sound sample
 *   - 2 16 bit integer samples output by audio file decoder representing two sound samples, one sample for each audio channel
 *   - 2 32 bit float samples output by mixer and consumed by playback API, one sample for each audio channel
 */

struct effect_data {
    /* Spatialization */
    f32 spatial_volume;
    f32 spatial_pan;
};

struct mix {
    struct mixer_track_handle track_handle;
    b32 track_finished;

    struct mixer_desc desc;
    struct effect_data effect_data;
    struct sound *source;
    u64 source_pos;
};

struct track {
    u64 gen;

    /* Controlled via interface */
    struct sound *sound;
    struct mixer_desc desc;

    /* Internal */
    struct mix mix;

    struct track *next;
    struct track *prev;
};

GLOBAL struct {
    struct sys_mutex mutex;

    /* Listener */
    struct v2 listener_pos;
    struct v2 listener_dir;

    /* Track list */
    struct arena track_arena;
    struct track *track_first_playing;
    struct track *track_last_playing;
    u64 track_playing_count;
    struct track *track_first_free;
} G = ZI, DEBUG_ALIAS(G, G_mixer);

/* ========================== *
 * Startup
 * ========================== */

struct mixer_startup_receipt mixer_startup(void)
{
    G.track_arena = arena_alloc(GIGABYTE(64));
    G.mutex = sys_mutex_alloc();
    G.listener_pos = V2(0, 0);
    G.listener_dir = V2(0, -1);

    return (struct mixer_startup_receipt) { 0 };
}

/* ========================== *
 * Track
 * ========================== */

INTERNAL struct mixer_track_handle track_to_handle(struct track *track)
{
    return (struct mixer_track_handle) {
        .gen = track->gen,
        .data = track
    };
}

INTERNAL struct track *track_from_handle(struct mixer_track_handle handle)
{
    struct track *track = (struct track *)handle.data;
    if (track && track->gen == handle.gen) {
        return track;
    } else {
        return NULL;
    }
}

INTERNAL struct track *track_alloc_locked(struct sys_lock *lock, struct sound *sound)
{
    sys_assert_locked_e(lock, &G.mutex);
    (UNUSED)lock;

    struct track *track = NULL;
    if (G.track_first_free) {
        /* Take from free list */
        track = G.track_first_free;
        struct track *next_free = track->next;
        G.track_first_free = next_free;
        if (next_free) {
            next_free->prev = NULL;
        }
        *track = (struct track) { .gen = track->gen + 1 };
    } else {
        /* Allocate new */
        track = arena_push_zero(&G.track_arena, struct track);
        track->gen = 1;
    }

    track->sound = sound;
    track->mix.source = sound;
    track->mix.track_handle = track_to_handle(track);

    /* Append to playing list */
    struct track *prev = G.track_last_playing;
    if (prev) {
        prev->next = track;
    } else {
        G.track_first_playing = track;
    }
    G.track_last_playing = track;
    track->prev = prev;
    ++G.track_playing_count;

    return track;
}

INTERNAL void track_release_locked(struct sys_lock *lock, struct track *track)
{
    sys_assert_locked_e(lock, &G.mutex);
    (UNUSED)lock;

    /* Remove from playing list */
    struct track *prev = track->prev;
    struct track *next = track->next;
    if (prev) {
        prev->next = next;
    } else {
        /* Track was first in list */
        G.track_first_playing = next;
    }
    if (next) {
        next->prev = prev;
    } else {
        /* Track was last in list */
        G.track_last_playing = prev;
    }
    --G.track_playing_count;
    ++track->gen;

    /* Add to free list */
    track->prev = NULL;
    track->next = G.track_first_free;
    if (G.track_first_free) {
        G.track_first_free->prev = track;
    }
    G.track_first_free = track;
}

/* ========================== *
 * Interface
 * ========================== */

/* TODO: Rework interface to take "mixer_cmd"s instead of
 * directly modifying tracks. */

struct mixer_track_handle mixer_play(struct sound *sound)
{
    return mixer_play_ex(sound, MIXER_DESC());
}

struct mixer_track_handle mixer_play_ex(struct sound *sound, struct mixer_desc desc)
{
    struct track *track;
    {
        struct sys_lock lock = sys_mutex_lock_e(&G.mutex);
        {
            track = track_alloc_locked(&lock, sound);
            track->desc = desc;
        }
        sys_mutex_unlock(&lock);
    }
    return track_to_handle(track);
}

/* NOTE: This is quite inefficient. */
struct mixer_desc mixer_track_get(struct mixer_track_handle handle)
{
    struct mixer_desc res = ZI;

    struct track *track = track_from_handle(handle);
    if (track) {
        /* TODO: Only lock mutex on track itself or something */
        struct sys_lock lock = sys_mutex_lock_e(&G.mutex);
        {
            /* Confirm handle is still valid now that we're locked */
            track = track_from_handle(handle);
            if (track) {
                res = track->desc;
            }
        }
        sys_mutex_unlock(&lock);
    }

    return res;
}

/* NOTE: This is quite inefficient. */
void mixer_track_set(struct mixer_track_handle handle, struct mixer_desc desc)
{
    struct track *track = track_from_handle(handle);
    if (track) {
        /* TODO: Only lock mutex on track itself or something */
        struct sys_lock lock = sys_mutex_lock_e(&G.mutex);
        {
            /* Confirm handle is still valid now that we're locked */
            track = track_from_handle(handle);
            if (track) {
                track->desc = desc;
            }
        }
        sys_mutex_unlock(&lock);
    }
}

void mixer_set_listener(struct v2 pos, struct v2 dir)
{
    struct sys_lock lock = sys_mutex_lock_e(&G.mutex);
    {
        G.listener_pos = pos;
        G.listener_dir = v2_norm(dir);
    }
    sys_mutex_unlock(&lock);
}

/* ========================== *
 * Update
 * ========================== */

INTERNAL i16 sample_sound(struct sound *sound, u64 sample_pos, b32 wrap)
{
    if (wrap) {
        return sound->pcm.samples[sample_pos % sound->pcm.count];
    } else if (sample_pos < sound->pcm.count) {
        return sound->pcm.samples[sample_pos];
    } else {
        return 0;
    }
}

/* To be called once per audio playback interval */
struct mixed_pcm_f32 mixer_update(struct arena *arena, u64 frame_count)
{
    __prof;

    struct temp_arena scratch = scratch_begin(arena);

    struct mixed_pcm_f32 res = ZI;
    res.count = frame_count * 2;
    res.samples = arena_push_array_zero(arena, f32, res.count);

    struct v2 listener_pos = V2(0, 0);
    struct v2 listener_dir = V2(0, 0);

    /* Create temp array of mixes */
    struct mix **mixes = NULL;
    u64 mixes_count = 0;
    {
        struct sys_lock lock = sys_mutex_lock_e(&G.mutex);

        /* Read listener info */
        listener_pos = G.listener_pos;
        listener_dir = G.listener_dir;

        /* Update & read mixes */
        mixes = arena_push_array(scratch.arena, struct mix *, G.track_playing_count);
        for (struct track *track = G.track_first_playing; track; track = track->next) {
            __profscope(prepare_track);
            struct mix *mix = &track->mix;
            mix->desc = track->desc;
            mixes[mixes_count++] = mix;
        }

        sys_mutex_unlock(&lock);
    }

    for (u64 mix_index = 0; mix_index < mixes_count; ++mix_index) {
        __profscope(mix_track);
        struct mix *mix = mixes[mix_index];

        if (mix->source->pcm.count <= 0) {
            /* Skip empty sounds */
            continue;
        }

        struct sound *source = mix->source;
        struct mixer_desc desc = mix->desc;
        struct effect_data *effect_data = &mix->effect_data;
        b32 source_is_stereo = source->flags & SOUND_FLAG_STEREO;
        f32 speed = max_f32(0, desc.speed);

        /* Determine sample range */
        u64 source_samples_count = 0;
        if (source_is_stereo) {
            source_samples_count = frame_count * 2;
            /* Round <samples_count * speed> to nearest frame boundary (nearest multiple of 2) */
            source_samples_count = (u64)math_ceil_to_int((f32)source_samples_count * speed);
            source_samples_count &= ~1;
        } else {
            source_samples_count = frame_count;
            /* Round <samples_count * speed> to nearest sample */
            source_samples_count = (u64)math_round_to_int((f32)source_samples_count * speed);
        }

        u64 source_sample_pos_start = mix->source_pos;
        u64 source_sample_pos_end = source_sample_pos_start + source_samples_count;
        if (source_sample_pos_end >= source->pcm.count) {
            if (desc.looping) {
                source_sample_pos_end = source_sample_pos_end % source->pcm.count;
            } else {
                source_sample_pos_end = source->pcm.count;
                mix->track_finished = true;
            }
        }
        u64 source_frames_count = source_is_stereo ? source_samples_count / 2 : source_samples_count;
        u64 source_frame_pos_start = source_is_stereo ? source_sample_pos_start / 2 : source_sample_pos_start;

        mix->source_pos = source_sample_pos_end;

        struct mixed_pcm_f32 mix_pcm = {
            .count = res.count,
            .samples = arena_push_array_zero(scratch.arena, f32, res.count)
        };

        /* ========================== *
         * Resample
         * ========================== */

        /* Transform 16 bit source -> 32 bit stereo at output duration */
        {
            __profscope(resample);
            f32 *out_samples = mix_pcm.samples;

            u64 out_frames_count = mix_pcm.count / 2;

            /* TODO: Fast path for 1:1 copy when speed = 1.0? */
            /* TODO: Optimize */
            if (source_is_stereo) {
                /* 16 bit Stereo -> 32 bit Stereo */
                for (u64 out_frame_pos = 0; out_frame_pos < out_frames_count; ++out_frame_pos) {
                    f32 in_frame_pos_exact = source_frame_pos_start + (((f32)out_frame_pos / (f32)out_frames_count) * (f32)source_frames_count);
                    u32 in_frame_pos_prev = math_floor_to_int(in_frame_pos_exact);
                    u32 in_frame_pos_next = math_ceil_to_int(in_frame_pos_exact);

                    /* Sample source */
                    f32 sample1_prev = sample_sound(source, (in_frame_pos_prev * 2) + 0, desc.looping) * (1.f / 32768.f);
                    f32 sample1_next = sample_sound(source, (in_frame_pos_next * 2) + 0, desc.looping) * (1.f / 32768.f);
                    f32 sample2_prev = sample_sound(source, (in_frame_pos_prev * 2) + 1, desc.looping) * (1.f / 32768.f);
                    f32 sample2_next = sample_sound(source, (in_frame_pos_next * 2) + 1, desc.looping) * (1.f / 32768.f);

                    /* Lerp */
                    f32 t = in_frame_pos_exact - (f32)in_frame_pos_prev;
                    f32 sample1 = math_lerp_f32(sample1_prev, sample1_next, t);
                    f32 sample2 = math_lerp_f32(sample2_prev, sample2_next, t);

                    out_samples[(out_frame_pos * 2) + 0] += sample1;
                    out_samples[(out_frame_pos * 2) + 1] += sample2;
                }
            } else {
                /* 16 bit Mono -> 32 bit Stereo */
                for (u64 out_frame_pos = 0; out_frame_pos < out_frames_count; ++out_frame_pos) {
                    f32 in_frame_pos_exact = source_frame_pos_start + (((f32)out_frame_pos / (f32)out_frames_count) * (f32)source_frames_count);
                    u32 in_frame_pos_prev = math_floor_to_int(in_frame_pos_exact);
                    u32 in_frame_pos_next = math_ceil_to_int(in_frame_pos_exact);

                    /* Sample source */
                    f32 sample_prev = sample_sound(source, in_frame_pos_prev, desc.looping) * (1.f / 32768.f);
                    f32 sample_next = sample_sound(source, in_frame_pos_next, desc.looping) * (1.f / 32768.f);

                    /* Lerp */
                    f32 t = (f32)in_frame_pos_exact - in_frame_pos_prev;
                    f32 sample = math_lerp_f32(sample_prev, sample_next, t);

                    out_samples[(out_frame_pos * 2) + 0] += sample;
                    out_samples[(out_frame_pos * 2) + 1] += sample;
                }
            }
        }

        /* ========================== *
         * Spatialize
         * ========================== */

        if (desc.flags & MIXER_FLAG_SPATIALIZE) {
            __profscope(spatialize);

            /* Algorithm constants */
            const f32 rolloff_height = 1.2f;
            const f32 rolloff_scale = 6.0f;
            const f32 pan_scale = 0.75;

            struct v2 pos = desc.pos;

            /* If sound pos = listener pos, pretend sound is close in front of listener. */
            if (v2_eq(listener_pos, pos)) {
                pos = v2_add(listener_pos, v2_mul(listener_dir, 0.001f));
            }
            struct v2 sound_rel = v2_sub(pos, listener_pos);
            struct v2 sound_rel_dir = v2_norm(sound_rel);

            /* Calculate volume */
            f32 volume_start = effect_data->spatial_volume;
            f32 volume_end;
            {
                /* https://www.desmos.com/calculator/c2h941hobz
                 * h = `rolloff_height`
                 * s = `rolloff_scale`
                 */
                f32 dist = v2_len(sound_rel);
                f32 v = (dist / rolloff_scale) + 1.0f;
                volume_end = rolloff_height * (1.0f / (v * v));
            }
            effect_data->spatial_volume = volume_end;

            /* Calculate pan */
            f32 pan_start = effect_data->spatial_pan;
            f32 pan_end = v2_wedge(listener_dir, sound_rel_dir) * pan_scale;
            effect_data->spatial_pan = pan_end;

            /* Spatialize samples */
            for (u64 frame_pos = 0; frame_pos < frame_count; ++frame_pos) {
                f32 t = (f32)frame_pos / (f32)(frame_count - 1);
                f32 volume = math_lerp_f32(volume_start, volume_end, t);
                f32 pan = math_lerp_f32(pan_start, pan_end, t);

                u64 sample1_index = frame_pos * 2;
                u64 sample2_index = sample1_index + 1;

                f32 sample_mono = ((mix_pcm.samples[sample1_index + 0] / 2.0f) + (mix_pcm.samples[sample2_index] / 2.0f)) * volume;

                mix_pcm.samples[sample1_index] = sample_mono * (1.0f - pan);
                mix_pcm.samples[sample2_index] = sample_mono * (1.0f + pan);
            }
        }

        /* ========================== *
         * Mix into result
         * ========================== */

        for (u64 i = 0; i < mix_pcm.count; ++i) {
            res.samples[i] += mix_pcm.samples[i] * desc.volume;
        }
    }

    {
        __profscope(update_track_effect_data);
        struct sys_lock lock = sys_mutex_lock_e(&G.mutex);
        for (u64 i = 0; i < mixes_count; ++i) {
            struct mix *mix = mixes[i];
            struct track *track = track_from_handle(mix->track_handle);
            if (track) {
                if (mix->track_finished) {
                    /* Release finished tracks */
                    track_release_locked(&lock, track);
                }
            }
        }
        sys_mutex_unlock(&lock);
    }

    scratch_end(scratch);
    return res;
}