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

namespace time_transform {

// Map for composition of maps
template <IsTimeTransformer Map1, IsTimeTransformer Map2> class CompositionMap {
public:
  static_assert(
      std::is_same_v<typename Map1::CodomainTime, typename Map2::DomainTime>,
      "Map1 CodomainTime must match Map2 DomainTime for composition.");
  static_assert(
      std::is_same_v<typename Map1::DomainTime::value_type,
                     typename Map2::CodomainTime::value_type>,
      "Underlying value_type (e.g., double) must match for both maps.");

  // Define types required by IsTimeTransformer
  using DomainTime = typename Map1::DomainTime;
  using CodomainTime = typename Map2::CodomainTime;
  using T = typename DomainTime::value_type;
  using DomainUnit = typename DomainTime::unit_tag;
  using CodomainUnit = typename CodomainTime::unit_tag;
  using IntermediateTime = typename Map1::CodomainTime;
  using IntermediateUnit = typename IntermediateTime::unit_tag;

  // Value type for iterator
  using MappedSegmentType = MappedSegment<DomainUnit, CodomainUnit, T>;
  static constexpr T EPSILON = time_units::EPSILON;

  class iterator {
  public:
    using iterator_category = std::bidirectional_iterator_tag;
    using value_type = MappedSegmentType;
    using difference_type = std::ptrdiff_t;
    using pointer = const value_type *;
    using reference = const value_type &;

  private:
    const CompositionMap *comp_map_ptr_;
    typename Map1::iterator it1_;
    typename Map2::iterator it2_;
    mutable std::optional<value_type> current_mapped_segment_;

    T calculate_composed_slope(T slope1, T slope2) const {
      if (slope1 == 0.0 || slope2 == 0.0)
        return 0.0;
      if (!std::isfinite(slope1) || !std::isfinite(slope2))
        return std::numeric_limits<T>::infinity();
      return slope1 * slope2;
    }

    // Mark iterator as end state
    void mark_as_end() {
      comp_map_ptr_ = nullptr;
      current_mapped_segment_.reset();
      // it1_, it2_ state doesn't matter when comp_map_ptr_ is null
    }

    // Ensure iterator points to a valid overlapping state or end
    void ensure_valid_state() {
      if (!comp_map_ptr_)
        return; // Already end

      constexpr int maxIterations = 1000;
      int iterations = 0;

      while (it1_ != comp_map_ptr_->map1_.end() &&
             it2_ != comp_map_ptr_->map2_.end()) {

        if (maxIterations > 0 && iterations++ >= maxIterations) {
          // Log details for debugging if possible/needed in a real environment
          // For example: log map types, current it1_/it2_ segment details.
          assert(false && "CompositionMap::iterator::ensure_valid_state hit "
                          "max iterations");
          mark_as_end(); // Give up and mark as end
          break;
        }

        const auto &seg1 = *it1_;
        const auto &seg2 = *it2_;
        auto intersection_b = seg1.target_range.intersect(seg2.source_range);

        // Check for valid, non-zero duration overlap
        if (intersection_b.has_value() &&
            intersection_b->duration().raw() > EPSILON) {
          break; // Found valid overlap
        }

        // No valid overlap. Decide which iterator(s) to advance.
        TimeRange<IntermediateUnit, T> target1_range = seg1.target_range;
        TimeRange<IntermediateUnit, T> source2_range = seg2.source_range;

        // Determine which iterator to advance based on their current segment
        // ends. This logic aims to "skip over" the one that finishes earlier in
        // the intermediate domain.
        bool try_advance_it1 =
            target1_range.end.raw() <= source2_range.end.raw() + EPSILON;
        bool try_advance_it2 =
            source2_range.end.raw() <= target1_range.end.raw() + EPSILON;

        typename Map1::iterator prev_it1_state = it1_;
        typename Map2::iterator prev_it2_state = it2_;
        bool it1_was_incremented_this_step = false;

        if (try_advance_it1 && it1_ != comp_map_ptr_->map1_.end()) {
          ++it1_;
          it1_was_incremented_this_step = true;
        }

        // If it1_ was incremented in this step (due to try_advance_it1 being
        // true), re-sync it2_ to the start of it1_'s new segment's target
        // range.
        if (it1_was_incremented_this_step &&
            it1_ != comp_map_ptr_->map1_.end()) {
          it2_ = comp_map_ptr_->map2_.getSegmentIteratorAt(
              (*it1_).target_range.start);
        } else if (try_advance_it2 && it2_ != comp_map_ptr_->map2_.end()) {
          ++it2_;
        }

        // Final check for termination after advancing and potential re-sync
        if (it1_ == comp_map_ptr_->map1_.end() ||
            it2_ == comp_map_ptr_->map2_.end()) {
          mark_as_end();
          break;
        }
      }
      current_mapped_segment_.reset();
    }

    // Calculate the current segment based on valid it1_ and it2_
    void cache_current() const {
      if (current_mapped_segment_ || !comp_map_ptr_)
        return; // Already cached or end

      // ensure_valid_state should guarantee we are not at end here unless map
      // is empty
      if (it1_ == comp_map_ptr_->map1_.end() ||
          it2_ == comp_map_ptr_->map2_.end()) {
        // Should not happen if ensure_valid_state was called correctly
        // and begin() didn't start at end unnecessarily.
        assert(false && "cache_current called on end iterator state");
        return;
      }

      const auto &seg1 = *it1_;
      const auto &seg2 = *it2_;

      auto intersection_b = seg1.target_range.intersect(seg2.source_range);
      assert(intersection_b.has_value() &&
             intersection_b->duration().raw() > EPSILON &&
             "cache_current called on non-overlapping iterator state");

      const auto &valid_intersection_b = intersection_b.value();

      // Map intersection B back to A using inverse map of seg1
      auto source_start_a_opt =
          seg1.inverse_map_point(valid_intersection_b.start);
      auto source_end_a_opt =
          seg1.inverse_map_point(valid_intersection_b.end, true);

      // Map intersection B forward to C using forward map of seg2
      auto target_start_c_opt = seg2.map_point(valid_intersection_b.start);
      auto target_end_c_opt = seg2.map_point(valid_intersection_b.end, true);

      // Assert successful mapping (should succeed if intersection is valid)
      assert(source_start_a_opt.has_value() && source_end_a_opt.has_value() &&
             target_start_c_opt.has_value() && target_end_c_opt.has_value() &&
             "Mapping failed for intersection points in cache_current");

      TimeRange<DomainUnit, T> source_range_a{*source_start_a_opt,
                                              *source_end_a_opt};
      TimeRange<CodomainUnit, T> target_range_c{*target_start_c_opt,
                                                *target_end_c_opt};

      // Ensure the calculated source range has positive duration
      if (source_range_a.duration().raw() <= EPSILON) {
        // This might happen with infinite slopes or zero duration
        // intersections. Treat this as an invalid segment for the purpose of
        // iteration? Or does it mean the iterator should advance again? For
        // now, let's assert against it, implies an issue upstream.
        assert(false &&
               "Zero duration source range calculated in cache_current");
        // If we wanted to handle it, we might need to trigger another ++
        // internally?
      }

      T composed_slope = calculate_composed_slope(seg1.slope, seg2.slope);
      SegmentMarks::MaskType composed_mask = seg1.get_mask() | seg2.get_mask();

      current_mapped_segment_.emplace(source_range_a, target_range_c,
                                      composed_slope, composed_mask);
    }

  public:
    // Constructor for begin/specific position
    iterator(const CompositionMap *map, typename Map1::iterator it1,
             typename Map2::iterator it2)
        : comp_map_ptr_(map), it1_(it1), it2_(it2) {
      assert(comp_map_ptr_ != nullptr);
      ensure_valid_state(); // Find first valid overlap or mark as end
    }

    // Default constructor for end iterator
    iterator() : comp_map_ptr_(nullptr) {}

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

    iterator &operator++() {
      if (!comp_map_ptr_)
        return *this; // Already end

      // Ensure we are in a valid state before advancing
      if (it1_ == comp_map_ptr_->map1_.end() ||
          it2_ == comp_map_ptr_->map2_.end()) {
        mark_as_end();
        return *this;
      }

      // Directly access current segments through iterators
      const auto &current_seg1 = *it1_;
      const auto &current_seg2 = *it2_;

      auto intersection_b =
          current_seg1.target_range.intersect(current_seg2.source_range);

      // This assertion should hold if the iterator was valid before operator++
      // If not, it indicates a deeper issue or misuse of the iterator.
      assert(intersection_b.has_value() &&
             intersection_b->duration().raw() > EPSILON &&
             "operator++ called on an iterator state without a valid current "
             "intersection");

      IntermediateTime intersection_end_b = intersection_b.value().end;

      // Determine if the current segments' boundaries align with the
      // intersection end
      bool s1_boundary_matches_intersection_end =
          std::abs((intersection_end_b - current_seg1.target_range.end).raw()) <
          EPSILON;
      bool s2_boundary_matches_intersection_end =
          std::abs((intersection_end_b - current_seg2.source_range.end).raw()) <
          EPSILON;

      assert(s1_boundary_matches_intersection_end ||
             s2_boundary_matches_intersection_end &&
                 "Intersection end must match at least one segment boundary");

      bool did_it1_advance = false;
      if (s1_boundary_matches_intersection_end) {
        if (it1_ != comp_map_ptr_->map1_.end()) {
          ++it1_;
          did_it1_advance = true;
        }
      }

      if (s2_boundary_matches_intersection_end) {
        // Avoid advancing it2_ if it1_ already advanced to its end AND s1 also
        // matched. This can happen if both segments end at the same
        // intermediate point.
        if (!(s1_boundary_matches_intersection_end && did_it1_advance &&
              it1_ == comp_map_ptr_->map1_.end()) &&
            it2_ != comp_map_ptr_->map2_.end()) {
          ++it2_;
        }
      }

      // If it1_ advanced (e.g., due to a loop in map1 causing a jump in
      // intermediate time), it2_ must be re-synchronized to the new starting
      // point in the intermediate domain defined by map1's new segment.
      if (did_it1_advance && it1_ != comp_map_ptr_->map1_.end()) {
        const auto &next_map1_segment = *it1_;
        // Reset it2_ to ensure it's positioned correctly relative to map1's new
        // segment. This handles cases where map1's target_range.start might
        // "jump back".
        it2_ = comp_map_ptr_->map2_.getSegmentIteratorAt(
            next_map1_segment.target_range.start);
      }

      // Find the next valid overlapping state or end
      ensure_valid_state();

      return *this;
    }

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

    iterator &operator--() {
      if (!comp_map_ptr_) {
        return *this; // Already canonical end
      }
      current_mapped_segment_.reset(); // Invalidate cache

      if (it1_ == comp_map_ptr_->map1_.begin() &&
          it2_ == comp_map_ptr_->map2_.begin()) {
        mark_as_end();
        return *this;
      }

      if (it1_ == comp_map_ptr_->map1_.end() ||
          it2_ == comp_map_ptr_->map2_.end()) {
        mark_as_end();
        return *this;
      }

      typename Map1::iterator prev_it1_for_check = it1_;
      typename Map2::iterator prev_it2_for_check = it2_;
      bool it1_was_decremented_this_step = false;

      // Logic to decide which iterator(s) to decrement.
      // This should ideally use the *current* segment's start to decide.
      // For simplicity and symmetry with operator++, we try to decrement it1_
      // first, then re-sync it2_, and let ensure_valid_state find the actual
      // previous segment. The original logic based on
      // sX_defines_intersection_start is more precise for identifying which
      // underlying iterator *caused* the current segment start.

      // Get current segment's intersection details to determine how to
      // decrement.
      const auto &s1_current = *it1_;
      const auto &s2_current = *it2_;
      auto intersection =
          s1_current.target_range.intersect(s2_current.source_range);

      if (!intersection.has_value() ||
          intersection->duration().raw() <= EPSILON) {
        assert(false && "operator-- called on an iterator without a current "
                        "valid intersection");
        mark_as_end();
        return *this;
      }
      IntermediateTime current_intersection_start = intersection->start;

      bool s1_defined_current_segment_start =
          std::abs((current_intersection_start - s1_current.target_range.start)
                       .raw()) < EPSILON;
      bool s2_defined_current_segment_start =
          std::abs((current_intersection_start - s2_current.source_range.start)
                       .raw()) < EPSILON;

      // Decrement logic (similar to original, but with progress tracking)
      if (s1_defined_current_segment_start) {
        if (it1_ != comp_map_ptr_->map1_.begin()) {
          --it1_;
          it1_was_decremented_this_step = true;
        }
      }

      // If s2 also defined the start, or if s1 was at begin and couldn't
      // decrement.
      if (s2_defined_current_segment_start) {
        if (!it1_was_decremented_this_step ||
            s1_current.target_range.start ==
                current_intersection_start) { // If it1 didn't move or s1 still
                                              // defines the point
          if (it2_ != comp_map_ptr_->map2_.begin()) {
            --it2_;
          }
        } else if (it1_was_decremented_this_step &&
                   s1_defined_current_segment_start &&
                   s2_defined_current_segment_start) {
          // Both defined current start, it1 decremented. Decrement it2 also.
          if (it2_ != comp_map_ptr_->map2_.begin()) {
            --it2_;
          }
        }
      }

      // Check for actual progress to prevent infinite loops if underlying
      // iterators get stuck.
      bool it1_made_progress = (it1_ != prev_it1_for_check) ||
                               (it1_ == comp_map_ptr_->map1_.begin());
      bool it2_made_progress = (it2_ != prev_it2_for_check) ||
                               (it2_ == comp_map_ptr_->map2_.begin());

      if (!it1_made_progress && !it2_made_progress &&
          !(it1_ == comp_map_ptr_->map1_.begin() &&
            it2_ == comp_map_ptr_->map2_.begin())) {
        // Neither iterator moved, and we are not at the absolute begin(). This
        // implies a stuck state.
        assert(false && "operator-- stuck: No iterator made progress and not "
                        "at overall begin");
        mark_as_end();
        return *this;
      }

      // Re-synchronize it2_ based on it1_'s new state, similar to operator++.
      // This ensures ensure_valid_state works consistently.
      if (it1_was_decremented_this_step && it1_ != comp_map_ptr_->map1_.end()) {
        it2_ = comp_map_ptr_->map2_.getSegmentIteratorAt(
            (*it1_).target_range.start);
      } else if (it1_ == comp_map_ptr_->map1_.begin() &&
                 prev_it1_for_check != comp_map_ptr_->map1_.begin()) {
        // it1 moved to begin(), ensure it2 is synced to map1.begin's
        // target.start
        it2_ = comp_map_ptr_->map2_.getSegmentIteratorAt(
            (*it1_).target_range.start);
      }

      ensure_valid_state();
      return *this;
    }

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

    bool operator==(const iterator &other) const {
      // Both end iterators? (Check comp_map_ptr_ first)
      if (!comp_map_ptr_ && !other.comp_map_ptr_)
        return true;
      if (!comp_map_ptr_ || !other.comp_map_ptr_)
        return false;
      // Both valid, compare underlying iterators state
      return comp_map_ptr_ == other.comp_map_ptr_ && it1_ == other.it1_ &&
             it2_ == other.it2_;
    }

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

    friend class CompositionMap<Map1, Map2>;
  }; // End iterator class

private:
  const Map1 &map1_;
  const Map2 &map2_;

public:
  CompositionMap(const Map1 &map1, const Map2 &map2)
      : map1_(map1), map2_(map2) {}

  iterator begin() const {
    return iterator(this, map1_.begin(), map2_.begin());
  }
  iterator begin() { return iterator(this, map1_.begin(), map2_.begin()); }

  iterator end() const {
    return iterator(); // Default constructed iterator is end
  }
  iterator end() { return iterator(); }

  iterator getSegmentIteratorAt(const DomainTime &p) const {
    auto it1 = map1_.getSegmentIteratorAt(p);
    typename Map2::iterator it2;

    if (it1 == map1_.end()) {
      return end();
    } else {
      const auto &seg1 = *it1;
      auto p_intermediate_opt = seg1.map_point(p);

      if (!p_intermediate_opt) {
        // Mapping failed, likely p is exactly at end boundary or out of domain
        // Return end() as the segment is ambiguous or non-existent.
        return end();
      }
      it2 = map2_.getSegmentIteratorAt(*p_intermediate_opt);
    }

    return iterator(this, it1, it2);
  }
};

} // namespace time_transform