<commit_before>package partners

import (
	"fmt"
	"log"
	"net/url"
	"time"

	"github.com/brnstz/bus/internal/conf"
	"github.com/brnstz/bus/internal/etc"
	"github.com/brnstz/bus/internal/models"

	"github.com/brnstz/bus/internal/partners/nyct_subway"
	"github.com/brnstz/bus/internal/partners/transit_realtime"

	"github.com/golang/protobuf/proto"
)

var (
	esiURL = "http://datamine.mta.info/mta_esi.php"

	mtaSubwayRouteToFeed = map[string]string{
		"1":  "1",
		"2":  "1",
		"3":  "1",
		"4":  "1",
		"5":  "1",
		"6":  "1",
		"6X": "1",
		"S":  "1",
		"GS": "1",
		"L":  "2",
		"SI": "11",
	}
)

type mtaNYCSubway struct{}

// getURL returns the URL for getting this routeID's feed. Second return
// value is false if there is no feed to get.
func (p mtaNYCSubway) getURL(routeID string) (string, bool) {
	var u string

	// Get the feed for this route, if there is one. Otherwise, nothing
	// to return.
	feed, exists := mtaSubwayRouteToFeed[routeID]
	if !exists {
		return "", false
	}

	// Construct URL and call external API, possibly getting cached
	// value.
	q := url.Values{}
	q.Set("key", conf.Partner.DatamineAPIKey)
	q.Set("feed_id", feed)
	u = fmt.Sprint(esiURL, "?", q.Encode())

	return u, true
}

func (p mtaNYCSubway) Precache(agencyID, routeID string, directionID int) error {
	k := fmt.Sprintf("%v|%v|%v", agencyID, routeID, directionID)

	u, exists := p.getURL(routeID)
	if !exists {
		return nil
	}

	// Since the URL we call is the same no matter which direction, arbirarily
	// decide to ignore one of the directions.
	if directionID == 1 {
		return nil
	}

	_, err := etc.RedisCacheURL(u)
	if err != nil {
		log.Println("can't cache live subway response", err)
		return err
	}

	// attempt to parse response to ensure it is valid
	_, _, err = p.Live(agencyID, routeID, "", directionID)
	if err != nil {
		log.Println("can't parse response", err)
		return err
	}

	log.Println("succesfully saved", k)

	return nil
}

func (p mtaNYCSubway) Live(agencyID, routeID, stopID string, directionID int) (d []*models.Departure, v []models.Vehicle, err error) {
	now := time.Now()

	u, exists := p.getURL(routeID)
	if !exists {
		return
	}

	b, err := etc.RedisGet(u)
	if err != nil {
		log.Println("can't get live subways", err)
		return
	}

	// Load the protobuf struct
	tr := &transit_realtime.FeedMessage{}
	err = proto.Unmarshal(b, tr)
	if err != nil {
		log.Println("can't unmarshal", err)
		return
	}

	// Look at each message in the feed
	for _, e := range tr.Entity {

		// Get some updates
		tripUpdate := e.GetTripUpdate()
		trip := tripUpdate.GetTrip()
		stopTimeUpdates := tripUpdate.GetStopTimeUpdate()

		// If we have at least one stopTimeUpdate and the trip is non-nil,
		// we can get the NYCT extensions.
		if len(stopTimeUpdates) > 0 && trip != nil {
			var event interface{}

			// Get the NYC extension so we can see if the Trip is "assigned"
			// yet. If it's assigned, we'll put the vehicle on the map.
			event, err = proto.GetExtension(
				trip, nyct_subway.E_NyctTripDescriptor,
			)
			if err != nil {
				log.Println("can't get extension", err)
				return
			}
			nycTrip, ok := event.(*nyct_subway.NyctTripDescriptor)
			if !ok {
				log.Println("can't coerce to nyct_subway.NyctTripDescriptor")
				return
			}

			// The first update in an entity is the stop where the train will
			// next be. Include only "assigned" trips, which are those that
			// are about to start.
			if nycTrip.GetIsAssigned() {
				var vehicle models.Vehicle<commit_msg>match routes and don't recache same info<commit_after>package partners

import (
	"fmt"
	"log"
	"net/url"
	"time"

	"github.com/brnstz/bus/internal/conf"
	"github.com/brnstz/bus/internal/etc"
	"github.com/brnstz/bus/internal/models"

	"github.com/brnstz/bus/internal/partners/nyct_subway"
	"github.com/brnstz/bus/internal/partners/transit_realtime"

	"github.com/golang/protobuf/proto"
)

var (
	esiURL = "http://datamine.mta.info/mta_esi.php"

	mtaSubwayRouteToFeed = map[string]string{
		"1":  "1",
		"2":  "1",
		"3":  "1",
		"4":  "1",
		"5":  "1",
		"6":  "1",
		"6X": "1",
		"S":  "1",
		"GS": "1",
		"L":  "2",
		"SI": "11",
	}

	// mapping from feed routeIDs to actual routeIDs
	expressMatch = map[string]string{
		// The feed uses 6 to mean 6X. We need to dig deeper to
		// find out that it's express.
		"6": "6X",
	}
)

type mtaNYCSubway struct{}

// getURL returns the URL for getting this routeID's feed. Second return
// value is false if there is no feed to get.
func (p mtaNYCSubway) getURL(routeID string) (string, bool) {
	var u string

	// Get the feed for this route, if there is one. Otherwise, nothing
	// to return.
	feed, exists := mtaSubwayRouteToFeed[routeID]
	if !exists {
		return "", false
	}

	// Construct URL and call external API, possibly getting cached
	// value.
	q := url.Values{}
	q.Set("key", conf.Partner.DatamineAPIKey)
	q.Set("feed_id", feed)
	u = fmt.Sprint(esiURL, "?", q.Encode())

	return u, true
}

func (p mtaNYCSubway) Precache(agencyID, routeID string, directionID int) error {
	k := fmt.Sprintf("%v|%v|%v", agencyID, routeID, directionID)

	u, exists := p.getURL(routeID)
	if !exists {
		return nil
	}

	// Since the URL we call is the same no matter which direction, arbirarily
	// decide to ignore one of the directions.
	if directionID == 1 {
		return nil
	}

	// Also feed ID "1" applies to multiple routes. We only need to cache one.
	// Arbitrarily choose the "1" route.
	if mtaSubwayRouteToFeed[routeID] == "1" && routeID != "1" {
		return nil
	}

	_, err := etc.RedisCacheURL(u)
	if err != nil {
		log.Println("can't cache live subway response", err)
		return err
	}

	// attempt to parse response to ensure it is valid
	_, _, err = p.Live(agencyID, routeID, "", directionID)
	if err != nil {
		log.Println("can't parse response", err)
		return err
	}

	log.Println("succesfully saved", k)

	return nil
}

// matchingRoutes determines if a route from the feed is equivalent to our
// routeID
func (p mtaNYCSubway) matchingRoute(feedRouteID, routeID string) bool {
	// The obvious case
	if feedRouteID == routeID {
		return true
	}

	if expressMatch[feedRouteID] == routeID {
		return true
	}

	return false
}

func (p mtaNYCSubway) Live(agencyID, routeID, stopID string, directionID int) (d []*models.Departure, v []models.Vehicle, err error) {
	now := time.Now()

	u, exists := p.getURL(routeID)
	if !exists {
		return
	}

	b, err := etc.RedisGet(u)
	if err != nil {
		log.Println("can't get live subways", err)
		return
	}

	// Load the protobuf struct
	tr := &transit_realtime.FeedMessage{}
	err = proto.Unmarshal(b, tr)
	if err != nil {
		log.Println("can't unmarshal", err)
		return
	}

	// Look at each message in the feed
	for _, e := range tr.Entity {

		// Get some updates
		tripUpdate := e.GetTripUpdate()
		trip := tripUpdate.GetTrip()
		stopTimeUpdates := tripUpdate.GetStopTimeUpdate()

		if !p.matchingRoute(trip.GetRouteId(), routeID) {
			continue
		}

		// If we have at least one stopTimeUpdate and the trip is non-nil,
		// we can get the NYCT extensions.
		if len(stopTimeUpdates) > 0 && trip != nil {
			var event interface{}

			// Get the NYC extension so we can see if the Trip is "assigned"
			// yet. If it's assigned, we'll put the vehicle on the map.
			event, err = proto.GetExtension(
				trip, nyct_subway.E_NyctTripDescriptor,
			)
			if err != nil {
				log.Println("can't get extension", err)
				return
			}
			nycTrip, ok := event.(*nyct_subway.NyctTripDescriptor)
			if !ok {
				log.Println("can't coerce to nyct_subway.NyctTripDescriptor")
				return
			}

			// The first update in an entity is the stop where the train will
			// next be. Include only "assigned" trips, which are those that
			// are about to start.
			if nycTrip.GetIsAssigned() {
				var vehicle models.Vehicle
{'fraction_non_alphanumeric': 0.13440514469453377, 'fraction_numerical': 0.006430868167202572, 'mean_word_length': 3.3080332409972297, 'pattern_counts': {'":': 23, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7469697', 'n_tokens_mistral': 3108, 'n_tokens_neox': 2836, 'n_words': 857}