#include "timeline.h"
#include "rendercontext.h"
#include "simd.h"
#include "timing/composition.h"
#include "timing/global_loop.h"
#include "timing/operators.h"
#include "timing/piecewise_linear.h"
#include <catch2/catch.hpp>
#include <emscripten/atomic.h>
#include <emscripten/bind.h>
#include <emscripten/console.h>

Timeline::Timeline(
    int sampleRate,
    const std::shared_ptr<time_transform::TimelineTempoMap> &timing,
    const std::vector<std::shared_ptr<Track>> &tracks,
    const std::shared_ptr<AnalyticsObserver> &analyticsObserver,
    const std::shared_ptr<Limiter> &limiter,
    const std::shared_ptr<Meter> &masterMeter,
    const std::shared_ptr<TimelineSharedState> &sharedState)
    : timing(timing), sampleRate(sampleRate), tracks(tracks),
      tempBuffer(2, 4096), analyticsObserver(analyticsObserver),
      delayCompensationFrames(0), limiter(limiter), masterMeter(masterMeter),
      sharedState(sharedState) {
  assert(timing != nullptr);

  delayCompensationFrames = limiter ? limiter->getDelayFrames() : 0;
  assert(delayCompensationFrames >= 0);
}

double Timeline::getPosition() const {
  using namespace time_units;

  bool isPrerolled = emscripten_atomic_load_u32(
                         &sharedState->delayCompensationIsPrerolled) != 0;
  double currentPosition = emscripten_atomic_load_f64(&sharedState->position);
  if (isPrerolled) {
    return time_transform::findDomainPointForCodomainDelta(
               *timing, time_units::Beats(currentPosition),
               time_units::SecondsDelta(-delayCompensationFrames /
                                        (double)sampleRate))
        .raw();
  }
  return currentPosition;
}

void Timeline::setPosition(double position) {
  assert(std::isfinite(position));
  emscripten_atomic_store_f64(&sharedState->position, position);
  emscripten_atomic_store_u32(&sharedState->delayCompensationIsPrerolled, 0);
}

bool Timeline::isPlaying() const {
  return emscripten_atomic_load_u32(&sharedState->playing) != 0;
}

void Timeline::setPlaying(bool playing) {
  emscripten_atomic_store_u32(&sharedState->playing, playing ? 1 : 0);
}

double Timeline::getLoopStart() const {
  return emscripten_atomic_load_f64(&sharedState->loopStart);
}

void Timeline::setLoopStart(double start) {
  assert(std::isfinite(start));
  emscripten_atomic_store_f64(&sharedState->loopStart, start);
}

double Timeline::getLoopEnd() const {
  return emscripten_atomic_load_f64(&sharedState->loopEnd);
}

void Timeline::setLoopEnd(double end) {
  assert(!std::isnan(end));
  emscripten_atomic_store_f64(&sharedState->loopEnd, end);
}

bool Timeline::isLoopEnabled() const {
  return emscripten_atomic_load_u32(&sharedState->loopEnabled) != 0;
}

double Timeline::getFadeInStartBeats() const {
  return emscripten_atomic_load_f64(&sharedState->fadeInStartBeats);
}

void Timeline::setFadeInStartBeats(double start) {
  assert(!std::isnan(start));
  emscripten_atomic_store_f64(&sharedState->fadeInStartBeats, start);
}

double Timeline::getFadeInLengthBeats() const {
  return emscripten_atomic_load_f64(&sharedState->fadeInLengthBeats);
}

void Timeline::setFadeInLengthBeats(double length) {
  assert(!std::isnan(length));
  emscripten_atomic_store_f64(&sharedState->fadeInLengthBeats, length);
}

double Timeline::getFadeInExponent() const {
  return emscripten_atomic_load_f64(&sharedState->fadeInExponent);
}

void Timeline::setFadeInExponent(double exponent) {
  assert(std::isfinite(exponent));
  emscripten_atomic_store_f64(&sharedState->fadeInExponent, exponent);
}

double Timeline::getFadeOutEndBeats() const {
  return emscripten_atomic_load_f64(&sharedState->fadeOutEndBeats);
}

void Timeline::setFadeOutEndBeats(double end) {
  assert(!std::isnan(end));
  emscripten_atomic_store_f64(&sharedState->fadeOutEndBeats, end);
}

double Timeline::getFadeOutLengthBeats() const {
  return emscripten_atomic_load_f64(&sharedState->fadeOutLengthBeats);
}

void Timeline::setFadeOutLengthBeats(double length) {
  assert(!std::isnan(length));
  emscripten_atomic_store_f64(&sharedState->fadeOutLengthBeats, length);
}

double Timeline::getFadeOutExponent() const {
  return emscripten_atomic_load_f64(&sharedState->fadeOutExponent);
}

void Timeline::setFadeOutExponent(double exponent) {
  assert(std::isfinite(exponent));
  emscripten_atomic_store_f64(&sharedState->fadeOutExponent, exponent);
}

float Timeline::getMasterGain() const {
  return emscripten_atomic_load_f32(&sharedState->masterGain);
}

void Timeline::setMasterGain(float gain) {
  assert(std::isfinite(gain));
  emscripten_atomic_store_f32(&sharedState->masterGain, gain);
}

void Timeline::setLoopEnabled(bool enabled) {
  using namespace time_units;
  if (enabled) {
    // Only enable if loop points are valid
    const double start = getLoopStart();
    const double end = getLoopEnd();
    if (start >= end) {
      enabled = false;
    } else {
      // Check if loop length is at least 1 sample at current BPM
      const auto startSegment = timing->getSegmentIteratorAt(Beats(start));
      const auto endSegment = timing->getSegmentIteratorAt(Beats(end));
      assert(startSegment != timing->end() && endSegment != timing->end());
      const auto startMapped = startSegment->map_point(Beats(start));
      const auto endMapped = endSegment->map_point(Beats(end));
      assert(startMapped.has_value() && endMapped.has_value());
      if ((*endMapped - *startMapped).raw() < 1.0 / (double)sampleRate) {
        enabled = false;
      }
    }
  }
  emscripten_atomic_store_u32(&sharedState->loopEnabled, enabled ? 1 : 0);
}

std::shared_ptr<Track> Timeline::getTrack(int index) const {
  if (index < 0 || index >= tracks.size()) {
    return nullptr;
  }
  return tracks[index];
}

void Timeline::read(BufferF32 output) {
  auto &sharedStateRef = *sharedState;

  using namespace time_units;

  output.fill(0);
  if (!isPlaying()) {
    for (const auto &track : tracks) {
      track->progressSegment(output.getChannelCount(), output.getFrameCount());
    }
    if (masterMeter != nullptr) {
      masterMeter->update(output);
    }
    if (analyticsObserver != nullptr) {
      analyticsObserver->trackRenderCycleEnd();
    }
    return;
  }

  const auto initialPosition =
      emscripten_atomic_load_f64(&sharedStateRef.position);
  const auto isPrerolled =
      emscripten_atomic_load_u32(&sharedStateRef.delayCompensationIsPrerolled);
  const auto loopEnabled = isLoopEnabled();
  const auto loopStart = getLoopStart();
  const auto loopEnd = getLoopEnd();

  int prerollFramesLeft = (isPrerolled == 0) ? delayCompensationFrames : 0;

  const auto boundedSubdivideStep =
      [&](const time_transform::MappedSegment<
              time_units::BeatsTag, time_units::SecondsTag, double> &segment,
          time_units::Beats<double> position,
          time_units::Seconds<double> mappedPosition, size_t offset,
          size_t frameCount, time_units::BeatsDelta<double> segmentDuration)
      -> std::optional<time_units::Beats<double>> {
    auto outputSlice = output.slice(offset, offset + frameCount);
    auto workBuffer =
        tempBuffer.slice(0, frameCount, outputSlice.getChannelCount());
    for (const auto &track : tracks) {
      RenderContext renderContext(position, 1.0 / segment.slope, true, *timing,
                                  analyticsObserver.get(), segmentDuration);
      track->readSegment(&renderContext, workBuffer);
      outputSlice.sumWith(workBuffer);
    }
    applyFadesAndGain(position.raw(), 1.0 / segment.slope, outputSlice);
    if (segment.has_mark(time_transform::SegmentMarks::SEEK_TO_LOOP_START)) {
      return Beats(loopStart);
    }
    return std::nullopt;
  };

  // Render time map segments

  const auto boundedSubdivideAllSteps = [&](time_units::Beats<double> position,
                                            int frameCount) {
    if (loopEnabled && loopStart < loopEnd) {
      auto loopMap =
          time_transform::TimelineLoopMap(Beats(loopStart), Beats(loopEnd));
      auto composedMap = time_transform::CompositionMap(loopMap, *timing);
      return boundedSubdivide(boundedSubdivideStep, composedMap, position,
                              sampleRate, frameCount);
    } else {
      return boundedSubdivide(boundedSubdivideStep, *timing, position,
                              sampleRate, frameCount);
    }
  };

  time_units::Beats<double> finalPosition(initialPosition);

  // Preroll to delay compensate
  while (prerollFramesLeft > 0) {
    const auto frameCount =
        std::min(prerollFramesLeft, (int)output.getFrameCount());
    finalPosition = boundedSubdivideAllSteps(finalPosition, frameCount);
    prerollFramesLeft -= frameCount;
    auto outputPrefix = output.slice(0, frameCount);
    if (limiter != nullptr) {
      limiter->process(outputPrefix);
    }
    outputPrefix.fill(0.0);
  }

  // Render the buffer
  finalPosition =
      boundedSubdivideAllSteps(finalPosition, output.getFrameCount());
  if (masterMeter != nullptr) {
    masterMeter->update(output);
  }
  if (limiter != nullptr) {
    limiter->process(output);
  }

  // Use CAS to update position atomically
  const auto finalPositionF64 = finalPosition.raw();
  const auto initialPositionU64 =
      *reinterpret_cast<const uint64_t *>(&initialPosition);
  const auto finalPositionU64 =
      *reinterpret_cast<const uint64_t *>(&finalPositionF64);
  {
    bool success =
        initialPositionU64 ==
        emscripten_atomic_cas_u64(&sharedStateRef.position, initialPositionU64,
                                  finalPositionU64);
    if (isPrerolled == 0 && success) {
      emscripten_atomic_cas_u32(&sharedStateRef.delayCompensationIsPrerolled, 0,
                                1);
    }
  }

  if (analyticsObserver != nullptr) {
    analyticsObserver->trackRenderCycleEnd();
  }
}

void Timeline::applyFadesAndGain(double position, double bps,
                                 BufferF32 output) {
  const auto localFadeInStartBeats = getFadeInStartBeats();
  const auto localFadeOutEndBeats = getFadeOutEndBeats();

  const auto projectSamplesToBeats = bps / (double)sampleRate;
  const auto outputChannelCount = output.getChannelCount();

  const auto happyFadeInBeats = std::abs(getFadeInLengthBeats()) + 1.0e-6f;
  const auto happyFadeOutBeats = std::abs(getFadeOutLengthBeats()) + 1.0e-6f;

  const auto happyFadeInExponent = std::max(0.0, getFadeInExponent());
  const auto happyFadeOutExponent = std::max(0.0, getFadeOutExponent());

  const auto masterGain = getMasterGain();

  float sampleBeats[output.getFrameCount()];
  float gainCurve[output.getFrameCount()];

#pragma clang loop vectorize(enable)
  for (int i = 0; i < output.getFrameCount(); i++) {
    sampleBeats[i] = position + i * projectSamplesToBeats;
  }

  vfast_curve_fade_array(
      sampleBeats, masterGain, localFadeInStartBeats, localFadeOutEndBeats,
      happyFadeInBeats, happyFadeOutBeats, happyFadeInExponent,
      happyFadeOutExponent, gainCurve, output.getFrameCount());

  for (int channel = 0; channel < outputChannelCount; channel++) {
    float *const channelData = output.getChannelData(channel);
#pragma clang loop vectorize(enable)
    for (int i = 0; i < output.getFrameCount(); i++) {
      channelData[i] *= gainCurve[i];
    }
  }
}

EMSCRIPTEN_BINDINGS(timeline) {
  using namespace emscripten;

  class_<Timeline>("Timeline")
      .smart_ptr<std::shared_ptr<Timeline>>("Timeline")
      .property("position", &Timeline::getPosition, &Timeline::setPosition)
      .property("playing", &Timeline::isPlaying, &Timeline::setPlaying)
      .property("loopStart", &Timeline::getLoopStart, &Timeline::setLoopStart)
      .property("loopEnd", &Timeline::getLoopEnd, &Timeline::setLoopEnd)
      .property("loopEnabled", &Timeline::isLoopEnabled,
                &Timeline::setLoopEnabled)
      .property("fadeInStartBeats", &Timeline::getFadeInStartBeats,
                &Timeline::setFadeInStartBeats)
      .property("fadeInLengthBeats", &Timeline::getFadeInLengthBeats,
                &Timeline::setFadeInLengthBeats)
      .property("fadeInExponent", &Timeline::getFadeInExponent,
                &Timeline::setFadeInExponent)
      .property("fadeOutEndBeats", &Timeline::getFadeOutEndBeats,
                &Timeline::setFadeOutEndBeats)
      .property("fadeOutLengthBeats", &Timeline::getFadeOutLengthBeats,
                &Timeline::setFadeOutLengthBeats)
      .property("fadeOutExponent", &Timeline::getFadeOutExponent,
                &Timeline::setFadeOutExponent)
      .property("masterGain", &Timeline::getMasterGain,
                &Timeline::setMasterGain)
      .function("getTrackCount", &Timeline::getTrackCount)
      .function("getTrack(index)", &Timeline::getTrack);
};

TEST_CASE("timeline memory tests", "[timeline]") {
  auto meter = std::make_shared<Meter>(2, 44100, 10, 10, 0.99f);
  auto track =
      std::make_shared<Track>(44100, 0.99f, 1.0, 0.0, false, meter, nullptr,
                              std::vector<std::shared_ptr<AudioClip>>{});
  auto timing = std::make_shared<time_transform::TimelineTempoMap>(
      2.0, std::vector<time_transform::TimelineTempoMap::LinearSegment>{});
  auto limiter = std::make_shared<Limiter>(2, 44100, 0.01f, 0.01f, 0.01f, 0.f,
                                           false, 0.8f);
  auto timeline = std::make_shared<Timeline>(
      44100, timing, std::vector<std::shared_ptr<Track>>{track}, nullptr,
      limiter, meter, std::make_shared<TimelineSharedState>());

  timeline->read(BufferF32(2, 1024));
  REQUIRE(timeline->getPosition() == 0.0);

  timeline->setPlaying(true);
  timeline->read(BufferF32(2, 1024));
  REQUIRE(std::fabs(timeline->getPosition() - 1024.0 / 44100.0 * 2.0) < 1e-6);

  auto monoMeter = std::make_shared<Meter>(1, 44100, 10, 10, 0.99f);
  auto monoTrack =
      std::make_shared<Track>(44100, 0.99f, 1.0, 0.0, false, monoMeter, nullptr,
                              std::vector<std::shared_ptr<AudioClip>>{});
  auto monoLimiter = std::make_shared<Limiter>(1, 44100, 0.01f, 0.01f, 0.01f,
                                               0.f, false, 0.8f);
  auto monoTimeline = std::make_shared<Timeline>(
      44100, timing, std::vector<std::shared_ptr<Track>>{monoTrack}, nullptr,
      monoLimiter, monoMeter, std::make_shared<TimelineSharedState>());

  BufferF32 monoBuffer(1, 1024);
  monoTimeline->read(monoBuffer);
}