#include <psemek/audio/stream.hpp>
#include <psemek/audio/duration.hpp>

namespace psemek::audio
{

	namespace
	{

		template <typename EnvelopeFn>
		struct envelope_base_impl
			: stream
		{
			envelope_base_impl(stream_ptr stream, EnvelopeFn envelope_fn, std::optional<duration> truncate)
				: stream_(std::move(stream))
				, envelope_fn_(std::move(envelope_fn))
				, remaining_(truncate ? std::optional{truncate->samples()} : std::nullopt)
				, total_length_(truncate ? stream_->played() + *remaining_ : stream_->length())
			{}

			std::optional<std::size_t> length() const override
			{
				return total_length_;
			}

			std::size_t read(util::span<float> samples) override
			{
				std::size_t const max_count = remaining_ ? std::min(*remaining_, samples.size()) : samples.size();

				auto const count = stream_->read(samples.prefix(max_count));
				if (remaining_)
					*remaining_ -= count;

				auto out = samples.data();
				for (std::size_t i = 0; i < count; i += 2)
				{
					float gain = envelope_fn_(played_);
					out[0] *= gain;
					out[1] *= gain;
					out += 2;
					played_ += duration::from_frames(1);
				}

				return count;
			}

			std::size_t played() const override
			{
				return stream_->played();
			}

		private:
			stream_ptr stream_;
			EnvelopeFn envelope_fn_;
			duration played_;
			std::optional<std::size_t> remaining_;
			std::optional<std::size_t> total_length_;
		};

	}

	// Guaranteed to call envelope_fn exactly once per audio frame
	template <typename EnvelopeFn>
	stream_ptr envelope_base(stream_ptr stream, EnvelopeFn envelope_fn, std::optional<duration> truncate = {})
	{
		return std::make_shared<envelope_base_impl<EnvelopeFn>>(std::move(stream), std::move(envelope_fn), truncate);
	}

}
