package main import ( "database/sql" "encoding/json" "fmt" "io" "log" "math" "net/http" "os" "strconv" "strings" pb "github.com/cjdenio/uta-trax-api/proto" "google.golang.org/protobuf/proto" _ "github.com/mattn/go-sqlite3" ) var scheduleDb *sql.DB var headsignMappings map[string]string = map[string]string{ "to medical": "To University Medical", "to airport": "To Airport", "to draper": "To Draper", "to salt lake": "To Salt Lake Central", "to meadowbrook": "To Meadowbrook", "to ballpark": "To Ballpark", } func distance(lat1 float64, lng1 float64, lat2 float64, lng2 float64) float64 { radlat1 := float64(math.Pi * lat1 / 180) radlat2 := float64(math.Pi * lat2 / 180) theta := float64(lng1 - lng2) radtheta := float64(math.Pi * theta / 180) dist := math.Sin(radlat1)*math.Sin(radlat2) + math.Cos(radlat1)*math.Cos(radlat2)*math.Cos(radtheta) if dist > 1 { dist = 1 } dist = math.Acos(dist) dist = dist * 180 / math.Pi dist = dist * 60 * 1.1515 return dist } func getStationForVehicle(vehicle *pb.VehiclePosition) *pb.VehicleFeed_Station { rows, err := scheduleDb.Query(`SELECT stop_lat, stop_lon, stops.stop_id, stop_name FROM stop_times INNER JOIN stops ON stops.stop_id = stop_times.stop_id WHERE stop_times.trip_id = ? AND stop_times.pickup_type != 1 AND stop_times.drop_off_type != 1;`, vehicle.Trip.TripId) if err != nil { fmt.Println(err) return nil } nearestStop := new(pb.VehicleFeed_Station) nearestStopDistance := math.Inf(1) for rows.Next() { var stop_lat float32 var stop_lon float32 var stop_id string var stop_name string rows.Scan(&stop_lat, &stop_lon, &stop_id, &stop_name) distance := distance(float64(stop_lat), float64(stop_lon), float64(*vehicle.Position.Latitude), float64(*vehicle.Position.Longitude)) if distance < nearestStopDistance { nearestStop.Id = stop_id nearestStop.Name = stop_name nearestStop.Lat = stop_lat nearestStop.Lon = stop_lon nearestStopDistance = distance } } // sometimes trains aren't where they should be, ignore nearest stop if more than 2 miles away if nearestStopDistance > 2 { return nil } return nearestStop } func fetchVehicleFeed() ([]byte, error) { if os.Getenv("PRETEND_FEED") == "true" { return os.ReadFile("pretend-feed.bin") } else { resp, err := http.Get("https://apps.rideuta.com/tms/gtfs/Vehicle") if err != nil { return nil, err } bytes, err := io.ReadAll(resp.Body) if err != nil { return nil, err } return bytes, nil } } func getVehicles() ([]*pb.VehiclePosition, *pb.FeedHeader, error) { bytes, err := fetchVehicleFeed() if err != nil { return nil, nil, err } feed := pb.FeedMessage{} err = proto.Unmarshal(bytes, &feed) if err != nil { return nil, nil, err } log.Printf("Found %d vehicles...", len(feed.Entity)) vehicles := make([]*pb.VehiclePosition, 0, len(feed.Entity)) for _, v := range feed.Entity { if v.Vehicle != nil { vehicles = append(vehicles, v.Vehicle) } } return vehicles, feed.Header, nil } func feedifyVehicles(vehicles []*pb.VehiclePosition, header *pb.FeedHeader, routeTypes map[int]bool) pb.VehicleFeed { vehicle_feed := make([]*pb.VehicleFeed_Vehicle, 0, len(vehicles)) for _, vehicle := range vehicles { rows, err := scheduleDb.Query(` SELECT trips.route_id, trips.trip_headsign, routes.route_type, routes.route_color, routes.route_short_name, routes.route_long_name FROM trips INNER JOIN routes ON routes.route_id = trips.route_id WHERE trips.trip_id = ? LIMIT 1; `, *vehicle.Trip.TripId) if err != nil { fmt.Println(err) continue } var route_id string var trip_headsign sql.NullString var route_type int32 var route_color sql.NullString var route_short_name sql.NullString var route_long_name sql.NullString if !rows.Next() { fmt.Println(rows.Err()) continue } err = rows.Scan(&route_id, &trip_headsign, &route_type, &route_color, &route_short_name, &route_long_name) if err != nil { fmt.Println(err) continue } // sometimes headsigns are weird. this lets me manually change them if trip_headsign.Valid { if mapping, ok := headsignMappings[strings.ToLower(trip_headsign.String)]; ok { trip_headsign.String = mapping } } rows.Close() if routeTypes != nil && !routeTypes[int(route_type)] { continue } vehicle_feed = append(vehicle_feed, &pb.VehicleFeed_Vehicle{ Lat: *vehicle.Position.Latitude, Lon: *vehicle.Position.Longitude, Bearing: *vehicle.Position.Bearing, // Line: trip.Line, Id: *vehicle.Vehicle.Id, // Direction: trip.Direction, NearestStation: getStationForVehicle(vehicle), Headsign: trip_headsign.String, Route: &pb.VehicleFeed_Route{ Id: route_id, Type: pb.VehicleFeed_Route_RouteType(route_type + 1), // lol Color: route_color.String, ShortName: route_short_name.String, LongName: route_long_name.String, }, }) } return pb.VehicleFeed{ Vehicles: vehicle_feed, Info: &pb.VehicleFeed_FeedInfo{ LastUpdate: *header.Timestamp, }, } } func main() { fmt.Println("Opening database...") _db, err := sql.Open("sqlite3", "uta-gtfs.db") if err != nil { log.Fatal(err) } scheduleDb = _db http.Handle("/", http.FileServer(http.Dir("./static"))) http.HandleFunc("/schema.proto", func(w http.ResponseWriter, r *http.Request) { http.ServeFile(w, r, "./proto/schema.proto") }) http.HandleFunc("/api", func(w http.ResponseWriter, r *http.Request) { vehicles, header, err := getVehicles() if err != nil { log.Println(err) w.WriteHeader(http.StatusInternalServerError) return } routeTypes := r.URL.Query().Get("route_types") var routeTypeFilter map[int]bool if routeTypes != "" { routeTypeFilter = make(map[int]bool) for t := range strings.SplitSeq(routeTypes, ",") { routeTypeNumber, err := strconv.Atoi(t) if err == nil { routeTypeFilter[routeTypeNumber] = true } } } if r.Header.Get("Accept") == "application/json" { feed := feedifyVehicles(vehicles, header, routeTypeFilter) response, err := json.Marshal(&feed) if err != nil { log.Println(err) w.WriteHeader(http.StatusInternalServerError) return } w.Header().Add("Content-Type", "application/json") w.Header().Add("Content-Length", strconv.Itoa(len(response))) w.Write(response) } else { // send a protobuf feed := feedifyVehicles(vehicles, header, routeTypeFilter) b, err := proto.Marshal(&feed) if err != nil { log.Println(err) w.WriteHeader(http.StatusInternalServerError) return } w.Header().Add("Content-Type", "application/protobuf") w.Header().Add("Content-Length", strconv.Itoa(len(b))) w.Write(b) } }) port := "3000" if portEnv, ok := os.LookupEnv("PORT"); ok { port = portEnv } fmt.Printf("Started on port %s\n", port) err = http.ListenAndServe(":"+port, nil) if err != nil { log.Fatal(err) } }