package types

import (
	"iter"
	"math"
	"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) HaversineDistance(other *Stop) float64 {
	return haversine(s.StopLat, s.StopLon, other.StopLat, other.StopLon)
}

func haversine(lat1, lon1, lat2, lon2 float64) float64 {
	const R = 6371 // Earth radius in km
	dLat := (lat2 - lat1) * math.Pi / 180
	dLon := (lon2 - lon1) * math.Pi / 180
	a := math.Sin(dLat/2)*math.Sin(dLat/2) + math.Cos(lat1*math.Pi/180)*math.Cos(lat2*math.Pi/180)*math.Sin(dLon/2)*math.Sin(dLon/2)
	c := 2 * math.Atan2(math.Sqrt(a), math.Sqrt(1-a))
	return R * c
}

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
					}
				}
			}
		}

	}
}