#include <iostream>
#include <filesystem>
#include <memory>
#include <vector>
#include <ranges>

#include "fmt/format.h"
#include "miniaudio.h"

extern "C" {
#include "libavcodec/avcodec.h"
#include "libavformat/avformat.h"
#include "libavcodec/avcodec.h"
}

enum class track_type {
    unknown,
    video,
    audio,
};

class track {
public:
    virtual ~track() {
        av_frame_free(&m_frame);
    };

    virtual auto str() const -> std::string = 0;
    auto stream_index() const -> std::size_t { return m_stream_index; }
    auto type() const -> track_type { return m_type; }
    auto frame() -> AVFrame* { return m_frame;}
    auto codec() -> AVCodecContext* { return m_codec; }

protected:
    track(AVStream* stream, track_type const& type)
        : m_stream_index(stream->index), m_stream(stream), m_type(type) {
        auto* params = stream->codecpar;
        auto const* codec = avcodec_find_decoder(params->codec_id);

        m_codec = avcodec_alloc_context3(codec);
        if (!m_codec) throw std::runtime_error("Couldn't create AVCodecContext");

        if (avcodec_parameters_to_context(m_codec, params) < 0)
            throw std::runtime_error("Couldn't initialize AVCodecContext");

        if (avcodec_open2(m_codec, codec, nullptr) < 0)
            throw std::runtime_error("Couln't open codec");

        m_frame = av_frame_alloc();
        if (!m_frame) throw std::runtime_error("Couldn't allocate AVFrame");
    }

protected:
    std::size_t     m_stream_index;
    AVStream*       m_stream;
    track_type      m_type;
    AVCodecContext* m_codec = nullptr;
    AVFrame*        m_frame = nullptr;
};
using track_ref_t = std::shared_ptr<track>;

class audio_track : public track {
public:
    audio_track(AVStream* stream)
        : track(stream, track_type::audio) { }

    auto str() const -> std::string override {
        return fmt::format("[{}] audio_track", m_stream_index);
    }
    auto channels() const -> std::size_t { return m_stream->codecpar->ch_layout.nb_channels; }
    auto sample_rate() const -> std::size_t { return m_stream->codecpar->sample_rate; }
    auto empty() const -> bool { return m_frame_count == 0; }

    auto read_pcm_frame(void* output, std::size_t frame_count) -> std::size_t {
        if (m_frame_count == 0) m_frame_count = std::size_t(m_frame->nb_samples);

        auto const data_size = av_get_bytes_per_sample(m_codec->sample_fmt);
        auto* ptr = static_cast<std::uint8_t*>(output);

        auto const end = m_frame->nb_samples;
        std::size_t frame_read = 0;
        std::size_t inc = 0;
        for (int i = end - m_frame_count; i < end; ++i) {
            for (int ch = 0; ch < m_codec->channels; ++ch) {
                std::memcpy(ptr + (inc++) * data_size, m_frame->data[ch] + data_size * i, data_size);
            }
            if (++frame_read >= frame_count) break;
        }
        m_frame_count -= frame_read;
        return frame_read;
    }

private:
    std::size_t m_frame_count{0};
};
using audio_track_ref_t = std::shared_ptr<audio_track>;

auto ffmpeg_error_str(int err) {
    static char str[AV_ERROR_MAX_STRING_SIZE]{};
    return av_make_error_string(str, AV_ERROR_MAX_STRING_SIZE, err);
};

class media {
public:
    media() {
        m_format = avformat_alloc_context();
        if (!m_format) throw std::runtime_error("Failed to create AVFormatContext");
        m_packet = av_packet_alloc();
        if (!m_packet) throw std::runtime_error("Couldn't allocate AVPacket");
    }
    ~media() {
        cleanup();
        avformat_free_context(m_format);
    }

    auto load(std::string const& path) -> void {
        cleanup();
        m_path = path;
        m_filename = m_path.filename().string();

        if (avformat_open_input(&m_format, m_path.string().c_str(), nullptr, nullptr) != 0)
            throw std::runtime_error(fmt::format("Failed to open file {}\n", m_filename));

        if (avformat_find_stream_info(m_format, nullptr) < 0)
            throw std::runtime_error("Error loading stream info");

        for (std::uint32_t i = 0; i < m_format->nb_streams; ++i) {
            auto* stream = m_format->streams[i];
            auto* params = stream->codecpar;
            auto const* codec = avcodec_find_decoder(params->codec_id);
            if (!codec) continue;

            switch (params->codec_type) {
                case AVMEDIA_TYPE_AUDIO:
                    m_tracks.push_back(std::make_shared<audio_track>(stream));
                    ++m_audio_track_size;
                    break;
                // case AVMEDIA_TYPE_VIDEO:
                //     m_tracks.push_back(std::make_shared<video_track>(std::size_t(i), stream));
                //     break;
                default: continue;
            }
        }
    }

    auto filename() const -> std::string const& { return m_filename; }
    auto tracks() const -> std::vector<track_ref_t> const& { return m_tracks; }
    auto str() const -> std::string {
        using namespace std::string_literals;
        std::string str{"media: "};
        str += m_filename + "\n";
        str += "  tracks:\n";
        for (std::size_t i = 0; i < m_tracks.size(); ++i) {
            str += "    "s + m_tracks[i]->str();
            if (i < m_tracks.size() - 1) str += "\n";
        }
        return str;
    }

    auto read_pcm_frame(void* output, std::size_t frame_count) -> void {
        std::size_t frame_read = 0;
        while (frame_read < frame_count) {
            auto track = std::static_pointer_cast<audio_track>(m_tracks[0]);
            if (track->empty()) read(track);
            frame_read += track->read_pcm_frame(output, frame_count);
        }
    }

    auto read(track_ref_t const& track) -> void {
        while(av_read_frame(m_format, m_packet) >= 0) {
            if (m_packet->stream_index != int(track->stream_index())) {
                av_packet_unref(m_packet);
                continue;
            }

            auto ret = avcodec_send_packet(track->codec(), m_packet);
            if (ret < 0) {
                fmt::print(stderr, "Error submitting the packet to the decoder!\n");
                break;
            }

            ret = avcodec_receive_frame(track->codec(), track->frame());
            if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
                av_packet_unref(m_packet);
                continue;
            } else if (ret < 0) {
                fmt::print(stderr, "Failed to decode packet: %s\n", ffmpeg_error_str(ret));
                break;
            }

            av_packet_unref(m_packet);
            break;
        }

    }

    auto seek(track_ref_t const& track, std::int64_t timestamp) -> void {
        av_seek_frame(m_format, track->stream_index(), timestamp, AVSEEK_FLAG_BACKWARD);
        read(track);
    }

private:
    auto cleanup() -> void {
        m_audio_track_size = 0;
        m_tracks.clear();
        avformat_close_input(&m_format);
    }

private:
    std::filesystem::path    m_path{};
    std::string              m_filename{};
    std::vector<track_ref_t> m_tracks{};

    std::size_t m_audio_track_size{0};

private:
    AVFormatContext* m_format = nullptr;
    AVPacket*        m_packet = nullptr;
};

auto callback([[maybe_unused]]ma_device* device, [[maybe_unused]]void* output, [[maybe_unused]]void const* input, [[maybe_unused]]ma_uint32 frame_count) -> void {
    auto* ptr = static_cast<media*>(device->pUserData);
    ptr->read_pcm_frame(output, frame_count);
}

auto entry() -> void {
    media decoder{};
    decoder.load("C:/Users/miku/Downloads/Porter Robinson - Trying to Feel Alive (Official Audio).webm");
    fmt::print("{}\n", decoder.str());

    auto track = std::static_pointer_cast<audio_track>(decoder.tracks()[0]);

    ma_device device;
    auto device_config = ma_device_config_init(ma_device_type_playback);
    device_config.playback.format   = ma_format_f32;
    device_config.playback.channels = track->channels();
    device_config.sampleRate        = track->sample_rate();
    device_config.dataCallback      = callback;
    device_config.pUserData         = &decoder;

    fmt::print("Channels: {}\n", track->channels());
    fmt::print("Sample rate: {}\n", track->sample_rate());

    if (ma_device_init(nullptr, &device_config, &device) != MA_SUCCESS)
        throw std::runtime_error("Failed to open playback device\n");

    if (ma_device_start(&device) != MA_SUCCESS)
        throw std::runtime_error("Failed to start playback device.\n");

    fmt::print("Now playing: {}\n", decoder.filename());
    fmt::print("Press enter to quit...");
    std::cin.get();

    ma_device_uninit(&device);
}

auto main([[maybe_unused]]int argc, [[maybe_unused]]char const* argv[]) -> int {
    try {
        entry();
    } catch (std::exception const& e) {
        fmt::print(stderr, "{}\n", e.what());
    }
    return 0;
}