#pragma once
#include "timing/time_transform.h"
#include <cassert>
#include <optional>

namespace time_transform {

// Map for timeline loops, (co)domain are beats, fulfills conditions:
// - Times before loopStart are mapped by identity
// - For every i >= 0, the window
//   [loopStart + i * loopLength, loopStart + (i + 1) * loopLength)
//   is mapped to [loopStart, loopEnd)
// - Iterator iterates over infinitely many segments:
//     S0: Maps [-inf, loopStart) to [-inf, loopStart) with slope 1.0
//     S{i+1}: Maps
//       [loopStart + i * loopLength, loopStart + (i + 1) * loopLength)
//       to [loopStart, loopEnd)
// - Every segment after S1 has the mark SegmentMarks::SEEK_TO_LOOP_START
template <typename DomainUnit, TimeValue T = double> class GlobalLoopMap {
public:
  using DomainTime = TypedTime<DomainUnit, T>;
  using CodomainTime = TypedTime<DomainUnit, T>;
  using DomainDelta = TypedTimeDelta<DomainUnit, T>;
  using CodomainDelta = TypedTimeDelta<DomainUnit, T>;

  class iterator; // Forward declaration

private:
  DomainTime loopStart_;
  DomainTime loopEnd_;
  DomainDelta loopLength_;

public:
  explicit GlobalLoopMap(DomainTime loopStart, DomainTime loopEnd)
      : loopStart_(loopStart), loopEnd_(loopEnd),
        loopLength_(loopEnd - loopStart) {
    assert(loopStart_ < loopEnd_ && "Loop start must be before loop end.");
    assert(loopStart_.is_finite() && loopEnd_.is_finite() &&
           "Loop bounds must be finite.");
  }

  DomainTime getLoopStart() const { return loopStart_; }
  DomainTime getLoopEnd() const { return loopEnd_; }

  class iterator {
  public:
    using iterator_category = std::bidirectional_iterator_tag;
    using value_type = MappedSegment<DomainUnit, DomainUnit, T>;
    using difference_type = std::ptrdiff_t;
    using pointer = const value_type *;
    using reference = const value_type &;

  private:
    const GlobalLoopMap *map_ptr_;
    // segment_idx_ = 0 for S0 (pre-loop)
    // segment_idx_ = i+1 for S{i+1} loop iterations (i = 0, 1, ...)
    // segment_idx_ = -1 for end() sentinel
    int segment_idx_;
    mutable std::optional<value_type> current_mapped_segment_;

    void cache_current() const {
      if (current_mapped_segment_ || !map_ptr_) {
        // map_ptr_ is null for end()
        return;
      }
      assert(segment_idx_ >= 0 &&
             "segment_idx_ must be non-negative for non-end iterator");

      const T slope_one = 1.0;

      if (segment_idx_ == 0) {
        // S0: Maps [-inf, loopStart) to [-inf, loopStart)
        TimeRange<DomainUnit, T> source_range(DomainTime::neg_inf(),
                                              map_ptr_->loopStart_);
        TimeRange<DomainUnit, T> target_range(DomainTime::neg_inf(),
                                              map_ptr_->loopStart_);
        current_mapped_segment_.emplace(source_range, target_range, slope_one,
                                        SegmentMarks::NONE);
      } else {
        // S{i+1}, where loop_iteration_idx i = segment_idx_ - 1
        int loop_iteration_idx = segment_idx_ - 1; // i = 0, 1, 2...

        DomainTime domain_iteration_start =
            map_ptr_->loopStart_ + (loop_iteration_idx * map_ptr_->loopLength_);
        DomainTime domain_iteration_end =
            domain_iteration_start + map_ptr_->loopLength_;

        TimeRange<DomainUnit, T> source_range(domain_iteration_start,
                                              domain_iteration_end);
        TimeRange<DomainUnit, T> target_range(map_ptr_->loopStart_,
                                              map_ptr_->loopEnd_);

        current_mapped_segment_.emplace(source_range, target_range, slope_one,
                                        loop_iteration_idx == 0
                                            ? SegmentMarks::NONE
                                            : SegmentMarks::SEEK_TO_LOOP_START);
      }
    }

  public:
    // Constructor for valid iterators
    iterator(const GlobalLoopMap *map, int idx)
        : map_ptr_(map), segment_idx_(idx) {
      assert(map_ptr_ != nullptr && "Creating iterator with null map_ptr_");
      assert(idx >= 0 &&
             "Segment index must be non-negative for valid iterator");
    }

    // Default constructor for placeholder/end iterator (sentinel)
    iterator() : map_ptr_(nullptr), segment_idx_(-1) {}

    reference operator*() const {
      cache_current();
      assert(current_mapped_segment_.has_value() &&
             "Dereferencing invalid or end GlobalLoopMap iterator");
      return *current_mapped_segment_;
    }
    pointer operator->() const {
      cache_current();
      assert(current_mapped_segment_.has_value() &&
             "Dereferencing invalid or end GlobalLoopMap iterator");
      return &(*current_mapped_segment_);
    }

    iterator &operator++() {
      if (map_ptr_) {
        // Not an end iterator
        current_mapped_segment_.reset();
        if (segment_idx_ < std::numeric_limits<int>::max()) {
          segment_idx_++;
        }
        // If segment_idx_ is max_int, it stays max_int (conceptually at
        // infinity)
      }
      return *this;
    }

    iterator operator++(int) {
      iterator tmp = *this;
      ++(*this);
      return tmp;
    }

    iterator &operator--() {
      if (map_ptr_) { // Not an end iterator
        current_mapped_segment_.reset();
        if (segment_idx_ > 0) {
          segment_idx_--;
        }
        // If segment_idx_ is 0 (S0, begin()), it remains 0. Cannot decrement
        // past S0.
      }
      return *this;
    }

    iterator operator--(int) {
      iterator tmp = *this;
      --(*this);
      return tmp;
    }

    bool operator==(const iterator &other) const {
      // Both are canonical end iterators (default constructed)
      if (!map_ptr_ && segment_idx_ == -1 && !other.map_ptr_ &&
          other.segment_idx_ == -1) {
        return true;
      }
      // If one is canonical end and the other is not.
      if ((!map_ptr_ && segment_idx_ == -1) ||
          (!other.map_ptr_ && other.segment_idx_ == -1)) {
        return false;
      }
      // Both are "normal" iterators (or one is map.end() which is also
      // specific)
      return map_ptr_ == other.map_ptr_ && segment_idx_ == other.segment_idx_;
    }

    bool operator!=(const iterator &other) const { return !(*this == other); }
  };

  iterator begin() const { return iterator(this, 0); }
  iterator begin() { return iterator(this, 0); } // Non-const

  iterator end() const { return iterator(); } // Sentinel end iterator
  iterator end() { return iterator(); }       // Non-const

  iterator getSegmentIteratorAt(const DomainTime &p) const {
    if (!p.is_finite()) {
      if (p.raw() < 0) {          // Negative infinity
        return iterator(this, 0); // S0 segment
      } else {                    // Positive infinity
        // Corresponds to a loop iteration with a very large index
        return iterator(this, std::numeric_limits<int>::max());
      }
    }

    // p is finite
    if (p < loopStart_) {
      return iterator(this, 0); // S0 segment
    } else {
      // p >= loopStart_
      DomainDelta time_relative_to_loop_start = p - loopStart_;
      // loopLength_.raw() must be positive in constructor.
      double i_double = time_relative_to_loop_start.raw() / loopLength_.raw();
      int loop_iteration_idx = static_cast<int>(
          std::floor(i_double)); // loop iteration index i (0, 1, 2...)

      // segment_idx for S{i+1} is loop_iteration_idx + 1
      if (loop_iteration_idx == std::numeric_limits<int>::max()) {
        // If i is already max_int, segment_idx would overflow if we add 1.
        // So, the segment_idx is also max_int.
        return iterator(this, std::numeric_limits<int>::max());
      } else {
        return iterator(this, loop_iteration_idx + 1);
      }
    }
  }
};

using TimelineLoopMap = GlobalLoopMap<BeatsTag, double>;

static_assert(
    IsTimeTransformer<TimelineLoopMap>,
    "TimelineLoopMap does not satisfy the IsTimeTransformer concept.");

} // namespace time_transform