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

namespace time_transform {
// Map for converting between time domains, where codomain time passes at a
// piecewise fixed rate in terms of domain time and T(Domain(0)) = Codomain(0)
template <typename DomainUnit, typename CodomainUnit, TimeValue T = double>
class PiecewiseLinearTimeMap {
public:
  using DomainTime = TypedTime<DomainUnit, T>;
  using CodomainTime = TypedTime<CodomainUnit, T>;
  using DomainDelta = TypedTimeDelta<DomainUnit, T>;
  using CodomainDelta = TypedTimeDelta<CodomainUnit, T>;

  // Forward declaration for iterator
  class iterator;

public:
  struct LinearSegment {
    // Codomain time flows at rate 1/dpcu *after* this time
    DomainTime start_time;
    double domain_per_codomain_unit; // Must be > 0

    // Codomain time at start time (keep this up to date as the map is modified)
    CodomainTime codomain_time_at_start;

    LinearSegment(DomainTime st, double dpcu, CodomainTime cts)
        : start_time(st), domain_per_codomain_unit(dpcu),
          codomain_time_at_start(cts) {
      assert(dpcu > 0.0 && std::isfinite(dpcu) &&
             "domain_per_codomain_unit must be positive and finite.");
    }

    LinearSegment(DomainTime st, double dpcu)
        : start_time(st), domain_per_codomain_unit(dpcu),
          codomain_time_at_start(CodomainTime(0.0)) {
      assert(dpcu > 0.0 && std::isfinite(dpcu) &&
             "domain_per_codomain_unit must be positive and finite.");
    }

    bool operator<(const LinearSegment &other) const {
      if (start_time != other.start_time) {
        return start_time < other.start_time;
      }
      return codomain_time_at_start < other.codomain_time_at_start;
    }
  };

private:
  std::vector<LinearSegment> segments_;
  T default_dpcu_;

  // Helper to calculate slope, 1.0 / dpcu
  static T calculate_slope(T dpcu) {
    assert(dpcu > 0.0 && std::isfinite(dpcu));
    return 1.0 / dpcu;
  }

  // Helper to find the last segment with start_time <= 0
  typename std::vector<LinearSegment>::iterator find_anchor_segment() {
    return std::find_if(segments_.rbegin(), segments_.rend(),
                        [](const LinearSegment &seg) {
                          return seg.start_time.raw() <= 0.0;
                        })
        .base(); // .base() converts reverse_iterator to forward iterator
  }

  typename std::vector<LinearSegment>::const_iterator
  find_anchor_segment() const {
    // Need const version as well
    auto rev_it = std::find_if(
        segments_.rbegin(), segments_.rend(),
        [](const LinearSegment &seg) { return seg.start_time.raw() <= 0.0; });
    return rev_it.base();
  }

public:
  // If the map has no segments, the default dpcu is used for all domain times
  explicit PiecewiseLinearTimeMap(T default_dpcu_val)
      : default_dpcu_(default_dpcu_val) {
    assert(default_dpcu_ > 0.0 && std::isfinite(default_dpcu_) &&
           "Default dpcu must be positive and finite.");
    // segments_ remains empty
  }

  // Constructor taking segments, calculates CTS values to enforce T(0)=0.
  explicit PiecewiseLinearTimeMap(T default_dpcu_val,
                                  std::vector<LinearSegment> &&segments)
      : default_dpcu_(default_dpcu_val) {
    assert(default_dpcu_ > 0.0 && std::isfinite(default_dpcu_) &&
           "Default dpcu must be positive and finite.");

    segments_ = std::move(segments);
    if (segments_.empty()) {
      return; // Nothing more to do if no segments
    }

    // 1. Order the segments and remove duplicates
    std::sort(segments_.begin(), segments_.end());
    segments_.erase(
        std::unique(segments_.begin(), segments_.end(),
                    [](const LinearSegment &a, const LinearSegment &b) {
                      return a.start_time == b.start_time;
                    }),
        segments_.end());

    // Ensure all DPCU values are valid before proceeding
    for (const auto &seg : segments_) {
      assert(seg.domain_per_codomain_unit > 0.0 &&
             std::isfinite(seg.domain_per_codomain_unit) &&
             "Segment dpcu must be positive and finite.");
    }

    // 2. Find anchor segment (last segment with start_time <= 0)
    auto anchor_it = find_anchor_segment();

    if (anchor_it != segments_.begin()) {
      // Case 1: Anchor segment exists (might be the first if its start <= 0)
      // The anchor segment is the one *before* the iterator returned by
      // find_if.base()
      auto zero_segment_it = std::prev(anchor_it);

      // Calculate CTS for the anchor segment ensuring its line passes through
      // (0,0)
      zero_segment_it->codomain_time_at_start =
          CodomainTime(zero_segment_it->start_time.raw() /
                       zero_segment_it->domain_per_codomain_unit);

      // Integrate forward from the anchor
      auto prev_it = zero_segment_it;
      for (auto current_it = std::next(prev_it); current_it != segments_.end();
           ++current_it) {
        DomainDelta domain_diff = current_it->start_time - prev_it->start_time;
        current_it->codomain_time_at_start =
            prev_it->codomain_time_at_start +
            CodomainDelta(domain_diff.raw() /
                          prev_it->domain_per_codomain_unit);
        prev_it = current_it;
      }

      // Integrate backward from the anchor
      auto current_it = zero_segment_it;
      while (current_it != segments_.begin()) {
        auto prev_segment_it = std::prev(current_it);
        DomainDelta domain_diff =
            current_it->start_time - prev_segment_it->start_time;
        // Calculate CTS of previous based on current
        prev_segment_it->codomain_time_at_start =
            current_it->codomain_time_at_start -
            CodomainDelta(domain_diff.raw() /
                          prev_segment_it->domain_per_codomain_unit);
        current_it = prev_segment_it;
      }
    } else {
      // Case 2: No segment starts at or before 0 (all start_time > 0)
      // The anchor point is (0,0), use the first segment's DPCU for t<0
      auto first_it = segments_.begin();

      // Calculate CTS for the first segment relative to (0,0)
      first_it->codomain_time_at_start = CodomainTime(
          first_it->start_time.raw() / first_it->domain_per_codomain_unit);

      // Integrate forward from the first segment
      auto prev_it = first_it;
      for (auto current_it = std::next(prev_it); current_it != segments_.end();
           ++current_it) {
        DomainDelta domain_diff = current_it->start_time - prev_it->start_time;
        current_it->codomain_time_at_start =
            prev_it->codomain_time_at_start +
            CodomainDelta(domain_diff.raw() /
                          prev_it->domain_per_codomain_unit);
        prev_it = current_it;
      }
      // No backward integration needed
    }
  }

  T getDefaultDpcu() const { return default_dpcu_; }

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

  private:
    const PiecewiseLinearTimeMap *map_ptr_;
    int effective_segment_idx_;
    mutable std::optional<value_type> current_mapped_segment_;

    void cache_current() const {
      if (current_mapped_segment_)
        return;

      const T inf = std::numeric_limits<T>::infinity();
      const DomainTime dom_0(0.0);
      const CodomainTime codom_0(0.0);

      TimeRange<DomainUnit, T> source_range(dom_0, dom_0);       // Placeholder
      TimeRange<CodomainUnit, T> target_range(codom_0, codom_0); // Placeholder
      T slope_val = 1.0;

      const auto &seg_vec = map_ptr_->segments_;

      if (seg_vec.empty()) {
        // Case 1: No actual segments, only default_dpcu_
        slope_val = calculate_slope(map_ptr_->default_dpcu_);
        if (effective_segment_idx_ == 0) {
          // First half: (-inf, 0)
          source_range = {DomainTime(-inf), dom_0};
          target_range = {CodomainTime(-inf), codom_0};
          // Right-anchored at (0,0)
        } else {
          // Second half: [0, +inf)
          source_range = {dom_0, DomainTime(inf)};
          target_range = {codom_0, CodomainTime(inf)};
          // Left-anchored at (0,0)
        }
      } else {
        // Case 2: Actual segments exist
        if (effective_segment_idx_ == 0) {
          // "Before" segment using first segment's dpcu
          const auto &first_ls = seg_vec.front();
          slope_val = calculate_slope(first_ls.domain_per_codomain_unit);
          source_range = {DomainTime(-inf), first_ls.start_time};
          target_range = {CodomainTime(-inf), first_ls.codomain_time_at_start};
          // Right-anchored at first_ls.start_time
        } else {
          // Segment from seg_vec
          // effective_segment_idx_ from 1 to seg_vec.size()
          int current_segment_vec_idx = effective_segment_idx_ - 1;
          const auto &current_ls = seg_vec[current_segment_vec_idx];
          slope_val = calculate_slope(current_ls.domain_per_codomain_unit);
          source_range.start = current_ls.start_time;
          target_range.start = current_ls.codomain_time_at_start;

          if (current_segment_vec_idx < seg_vec.size() - 1) {
            // Not the last actual segment
            const auto &next_ls = seg_vec[current_segment_vec_idx + 1];
            source_range.end = next_ls.start_time;
            target_range.end = next_ls.codomain_time_at_start;
          } else {
            // Last actual segment, extends to +inf
            source_range.end = DomainTime(inf);
            target_range.end = CodomainTime(inf);
          }
        }
      }
      current_mapped_segment_.emplace(source_range, target_range, slope_val,
                                      SegmentMarks::NONE);
    }

  public:
    iterator(const PiecewiseLinearTimeMap *map, int idx)
        : map_ptr_(map), effective_segment_idx_(idx) {}

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

    reference operator*() const {
      cache_current();
      return *current_mapped_segment_;
    }
    pointer operator->() const {
      cache_current();
      return &(*current_mapped_segment_);
    }

    iterator &operator++() {
      current_mapped_segment_.reset(); // Invalidate cache
      effective_segment_idx_++;
      // Ensure effective_segment_idx_ does not go past end()
      int end_idx = map_ptr_->segments_.empty()
                        ? 2
                        : static_cast<int>(map_ptr_->segments_.size()) + 1;
      if (effective_segment_idx_ >= end_idx) {
        effective_segment_idx_ = end_idx; // Canonical end state index
        // map_ptr_ remains valid for comparison with end()
      }
      return *this;
    }

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

    iterator &operator--() {
      current_mapped_segment_.reset();  // Invalidate cache
      if (effective_segment_idx_ > 0) { // Cannot decrement begin()
        effective_segment_idx_--;
      }
      // If effective_segment_idx_ was already 0 (begin), it remains 0.
      // This adheres to std::bidirectional_iterator: begin() is not
      // decrementable to a valid prior state.
      return *this;
    }

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

    bool operator==(const iterator &other) const {
      return map_ptr_ == other.map_ptr_ &&
             effective_segment_idx_ == other.effective_segment_idx_;
    }

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

    friend class PiecewiseLinearTimeMap<DomainUnit, CodomainUnit, T>;
  };

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

  // Add non-const overload for begin()
  iterator begin() { return iterator(this, 0); }

  iterator end() const {
    // N actual segments + 1 implicit "before" segment = N+1 effective segments
    // If segments_ is empty, there are 2 effective segments: (-inf, 0) and [0,
    // inf)
    return iterator(
        this, segments_.empty() ? 2 : static_cast<int>(segments_.size()) + 1);
  }

  // Add non-const overload for end()
  iterator end() {
    return iterator(
        this, segments_.empty() ? 2 : static_cast<int>(segments_.size()) + 1);
  }

  iterator getSegmentIteratorAt(const DomainTime &p) const {
    if (segments_.empty()) {
      // For an empty map, split point is DomainTime(0.0)
      bool is_before_split = p.is_finite() ? (p.raw() < 0.0) : true;
      if (p == DomainTime(-std::numeric_limits<T>::infinity()))
        is_before_split = true;

      return iterator(this, is_before_split ? 0 : 1);
    }

    // Find first LinearSegment S where S.start_time > p
    auto it_upper =
        std::upper_bound(segments_.begin(), segments_.end(), p,
                         [](const DomainTime &val, const LinearSegment &seg) {
                           return val < seg.start_time;
                         });

    if (it_upper == segments_.begin()) {
      // p is before or at the start_time of the very first segment.
      // If p < segments_.front().start_time, it's in the "before" segment (idx
      // 0). If p == segments_.front().start_time, it's in the first actual
      // segment (idx 1).
      if (p < segments_.front().start_time) {
        return iterator(this, 0); // "Before" segment
      } else {                    // p == segments_.front().start_time
        return iterator(this, 1); // First actual segment
      }
    } else {
      // it_upper points to S_k where S_k.start_time > p.
      // p must be >= (it_upper-1)->start_time.
      // So p belongs to the segment represented by (it_upper-1).
      // The index in segments_ is std::distance(segments_.begin(), it_upper-1).
      // The effective_segment_idx_ is this + 1.
      return iterator(
          this,
          static_cast<int>(std::distance(segments_.begin(), it_upper - 1) + 1));
    }
  }

  std::optional<T> js_transformBeatsToSeconds(T beats) const {
    auto dom_beats = DomainTime(beats);
    auto it = getSegmentIteratorAt(dom_beats);
    if (it == end()) {
      return std::nullopt;
    }
    auto res = it->map_point(dom_beats);
    if(res.has_value()) {
      return res->raw();
    }
    return std::nullopt;
  }
};

using TimelineTempoMap = PiecewiseLinearTimeMap<BeatsTag, SecondsTag, double>;

// Static assertion to ensure TimelineTempoMap conforms to
// IsTimeTransformer
static_assert(
    IsTimeTransformer<TimelineTempoMap>,
    "TimelineTempoMap does not satisfy the IsTimeTransformer concept.");

} // namespace time_transform