package types

import (
	"iter"
	"time"
)

type Stop struct {
	Trips        map[string]*Trip `json:"-" csv:"-"`
	StopId       string           `json:"stop_id" csv:"stop_id"`
	StopName     string           `json:"stop_name" csv:"stop_name"`
	StopLat      float64          `json:"stop_lat" csv:"stop_lat"`
	StopLon      float64          `json:"stop_lon" csv:"stop_lon"`
	LocationType int              `json:"location_type" csv:"location_type"`
	Transfers    []*Transfer      `json:"-" csv:"transfers"`
}

func (s *Stop) AddTransfer(transfer *Transfer) {
	if s.Transfers == nil {
		s.Transfers = make([]*Transfer, 0)
	}
	s.Transfers = append(s.Transfers, transfer)
}

func (s *Stop) AddTrip(trip *Trip) {

	s.Trips[trip.TripId] = trip
}

type TripWithDepartureTime struct {
	*Trip
	DepartureTime SecondsAfterMidnight
}

func (s *Stop) GetTripsAfter(when time.Time) iter.Seq[*TripWithDepartureTime] {
	startAfterMidnight := AsSecondsAfterMidnight(when)
	return func(yield func(*TripWithDepartureTime) bool) {
		for _, trip := range s.Trips {

			for _, stop := range trip.Stops {
				if stop.StopId == s.StopId && stop.ArrivalTime >= startAfterMidnight {

					if !yield(&TripWithDepartureTime{Trip: trip, DepartureTime: stop.DepartureTime}) {
						return
					}
					break
				}
			}

		}
	}
}

func (s *Stop) GetUpcomingStops(start *StopTime) iter.Seq[*StopTime] {
	return func(yield func(*StopTime) bool) {
		found := false
		for _, trip := range s.Trips {
			for _, stop := range trip.Stops {
				if !found {
					if stop.StopId == start.StopId && stop.DepartureTime >= start.ArrivalTime {
						found = true
					}
				} else {
					if !yield(stop) {
						return
					}
				}
			}
		}
	}
}

func (s *Stop) GetStopsAfter(when time.Time) iter.Seq2[*StopTime, *StopTime] {
	startAfterMidnight := AsSecondsAfterMidnight(when)
	return func(yield func(start, stop *StopTime) bool) {
		for _, trip := range s.Trips {
			found := false
			var start *StopTime
			for _, stop := range trip.Stops {
				if stop.StopId == s.StopId && stop.ArrivalTime >= startAfterMidnight {
					found = true
					start = stop
				}
				if found {
					if !yield(start, stop) {
						return
					}
				}
			}
		}

	}
}