szczanieckiej/traffic/convert_departures.go

// SPDX-FileCopyrightText: Adam Evyčędo
//
// SPDX-License-Identifier: AGPL-3.0-or-later

package traffic

import (
	"bufio"
	"context"
	"database/sql"
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"git.sr.ht/~sircmpwn/go-bare"
	"github.com/lanrat/extsort"
	_ "github.com/mattn/go-sqlite3"
)

func readInputRoutesIndex(c feedConverter) (feedConverter, error) {
	index := map[string]int64{}

	path := c.TmpFeedPath
	err := forEachRow(filepath.Join(path, "routes.txt"), func(offset int64, fields map[string]int, record []string) error {
		routeID := record[fields["route_id"]]
		index[routeID] = offset
		return nil
	})

	c.routesInputIndex = index
	return c, err
}

func readInputStopsIndex(c feedConverter) (feedConverter, error) {
	index := map[string]int64{}

	path := c.TmpFeedPath
	forEachRow(filepath.Join(path, "stops.txt"), func(offset int64, fields map[string]int, record []string) error {
		stopID := record[fields["stop_id"]]
		index[stopID] = offset
		return nil
	})

	c.stopsInputIndex = index
	return c, nil
}

func readInputTripsIndex(c feedConverter) (feedConverter, error) {
	index := map[string]int64{}

	path := c.TmpFeedPath
	err := forEachRow(filepath.Join(path, "trips.txt"), func(offset int64, fields map[string]int, record []string) error {
		tripID := record[fields["trip_id"]]
		index[tripID] = offset
		return nil
	})

	c.tripsInputIndex = index
	return c, err
}

func convertDepartures(c feedConverter) (feedConverter, error) { // O(n:stop_times) ; (TmpFeedPath, tripsInputIndex, stopsInputIndex -- TripsOffsets:map[tripID]offset >> trips)
	path := c.TmpFeedPath
	cacheDir, err := os.UserCacheDir()
	if err != nil {
		return c, fmt.Errorf("while getting cache dir: %w", err)
	}
	db, err := sql.Open("sqlite3", filepath.Join(cacheDir, "turntable.db"))
	if err != nil {
		return c, fmt.Errorf("while opening db: %w", err)
	}
	_, err = db.Exec("create table change_options(stop_id text, line_name text, headsign text, primary key(stop_id, line_name, headsign))")
	if err != nil {
		return c, fmt.Errorf("while creating changeOptions table: %w", err)
	}
	tx, err := db.Begin()
	if err != nil {
		return c, fmt.Errorf("while beginning transaction: %w", err)
	}
	defer tx.Rollback()

	tripsOffsets := map[string]uint{}
	var outputOffset uint = 0
	previousTrip := ""
	firstDepartureTime := uint(0)
	trip := Trip{}

	result, err := os.Create(filepath.Join(path, "trips.bare"))
	if err != nil {
		return c, fmt.Errorf("while creating file: %w", err)
	}
	defer result.Close()

	tripsFile, err := os.Open(filepath.Join(path, "trips.txt"))
	if err != nil {
		return c, fmt.Errorf("while opening trips file: %w", err)
	}
	defer tripsFile.Close()

	// TODO unnecessary overhead, should parse single csv header line
	trips := csv.NewReader(tripsFile)
	tripsHeader, err := trips.Read()
	if err != nil {
		return c, fmt.Errorf("while reading trips header: %w", err)
	}

	tripsFields := map[string]int{}
	for i, headerField := range tripsHeader {
		tripsFields[headerField] = i
	}

	stopsFile, err := os.Open(filepath.Join(path, "stops.txt"))
	if err != nil {
		return c, fmt.Errorf("while opening stops file: %w", err)
	}
	defer stopsFile.Close()

	// TODO unnecessary overhead, should parse single csv header line
	stops := csv.NewReader(stopsFile)
	stopsHeader, err := stops.Read()
	if err != nil {
		return c, fmt.Errorf("while reading stops header: %w", err)
	}

	stopsFields := map[string]int{}
	for i, headerField := range stopsHeader {
		stopsFields[headerField] = i
	}

	routesFile, err := os.Open(filepath.Join(path, "routes.txt"))
	if err != nil {
		return c, fmt.Errorf("while opening routes file: %w", err)
	}
	defer routesFile.Close()

	// TODO unnecessary overhead, should parse single csv header line
	routes := csv.NewReader(routesFile)
	routesHeader, err := routes.Read()
	if err != nil {
		return c, fmt.Errorf("while reading routes header: %w", err)
	}

	routesFields := map[string]int{}
	for i, headerField := range routesHeader {
		routesFields[headerField] = i
	}

	tripsThroughStopFile, err := os.Create(filepath.Join(path, "tripsthroughstop.csv"))
	if err != nil {
		return c, fmt.Errorf("while creating tripsThroughStop file: %w", err)
	}
	defer tripsThroughStopFile.Close()
	defer tripsThroughStopFile.Sync()
	tripsThroughStop := csv.NewWriter(tripsThroughStopFile)
	defer tripsThroughStop.Flush()
	err = tripsThroughStop.Write([]string{"stop_id", "trip_id", "sequence"})
	if err != nil {
		return c, fmt.Errorf("while writing tripsThroughStop header: %w", err)
	}

	err = forEachRow(filepath.Join(path, "stop_times.txt"), func(offset int64, fields map[string]int, record []string) error {
		departure := Departure{}
		tripID := record[fields["trip_id"]]
		stopID := record[fields["stop_id"]]

		if previousTrip != tripID {
			if previousTrip != "" {
				tripsOffsets, outputOffset, err = finishTrip(trip, result, tripsOffsets, outputOffset)
				if err != nil {
					return fmt.Errorf("while finishing trip: %w", err)
				}
			}

			trip, err = beginTrip(tripsFile, c, tripID, len(tripsHeader), Trip{}, tripsFields)
			if err != nil {
				return fmt.Errorf("while beginning trip: %w", err)
			}

			firstDepartureTime = departure.Time
		}

		routeRecord, err := readCsvLine(routesFile, c.routesInputIndex[trip.LineID], len(routesHeader))
		if err != nil && err != io.EOF {
			return fmt.Errorf("while reading a routes record: %w", err)
		}

		lineName := c.Feed.Flags().LineName
		for _, template := range []string{"route_short_name", "route_long_name"} {
			lineName = strings.Replace(lineName, "{{"+template+"}}", routeRecord[routesFields[template]], -1)
		}
		/*headsign := translateFieldDefault(trip.Headsign, c.feedInfo.Language, c.defaultLanguage, c.translations)
		translatedHeadsigns := translateField(trip.Headsign, c.feedInfo.Language, c.defaultLanguage, c.translations)
		for _, translatedHeadsign := range translatedHeadsigns {
			changeOptionsRecord = append(changeOptionsRecord, translatedHeadsign.Language)
			changeOptionsRecord = append(changeOptionsRecord, translatedHeadsign.Value)
		}*/

		_, err = tx.Exec("insert into change_options values(?, ?, ?) on conflict(stop_id, line_name, headsign) do nothing", stopID, lineName, trip.Headsign)
		if err != nil {
			return fmt.Errorf("while writing changeOptions record: %w", err)
		}

		fmt.Sscanf(record[fields["stop_sequence"]], "%d", &departure.StopSequence)
		fmt.Sscanf(record[fields["pickup_type"]], "%d", &departure.Pickup)
		fmt.Sscanf(record[fields["drop_off_type"]], "%d", &departure.Dropoff)

		stopSequence := fmt.Sprintf("%d", departure.StopSequence)
		tripsThroughStopRecord := []string{stopID, tripID, stopSequence}
		err = tripsThroughStop.Write(tripsThroughStopRecord)
		if err != nil {
			return fmt.Errorf("while writing tripsThroughStop record: %w", err)
		}

		departureTime, err := parseDepartureTime(record[fields["arrival_time"]])
		if err != nil {
			return fmt.Errorf("while parsing arrival time: %w", err)
		}
		departure.Time = uint(departureTime)
		if _, ok := c.Headways[tripID]; ok {
			departure.Time -= firstDepartureTime
		}

		trip.Departures = append(trip.Departures, departure)
		previousTrip = tripID

		return nil
	})

	tripsOffsets, outputOffset, err = finishTrip(trip, result, tripsOffsets, outputOffset)
	if err != nil {
		return c, fmt.Errorf("while finishing trip: %w", err)
	}

	c.tripsOffsets = tripsOffsets
	tx.Commit()
	return c, err
}

func beginTrip(tripsFile *os.File, c feedConverter, tripID string, tripsHeaderLen int, trip Trip, tripsFields map[string]int) (Trip, error) {
	tripRecord, err := readCsvLine(tripsFile, c.tripsInputIndex[tripID], tripsHeaderLen)
	if err != nil && err != io.EOF {
		return trip, fmt.Errorf("while reading a trips record: %w", err)
	}

	trip.Id = tripID
	switch c.Feed.Flags().Headsign {
	case HeadsignTripHeadsing:
		trip.Headsign = tripRecord[tripsFields["trip_headsign"]]
	case HeadsignTripLastStop:
		// TODO test this case
		/*
			stopRecord, err := readCsvLine(stopsFile, c.stopsInputIndex[stopID], len(stopsHeader))
			if err != nil && err != io.EOF {
				return fmt.Errorf("while reading a stops record: %w", err)
			}

			trip.Headsign = stopRecord[stopsFields["stop_name"]]
		*/
	}
	// TODO translated headsign(-s)

	if h, ok := c.Headways[trip.Id]; ok {
		trip.Headways = h
	}

	trip.ScheduleID = tripRecord[tripsFields["service_id"]]
	trip.LineID = tripRecord[tripsFields["route_id"]]
	fmt.Sscanf(tripRecord[tripsFields["direction_id"]], "%d", &trip.Direction)

	return trip, nil
}

func finishTrip(trip Trip, result io.Writer, tripsOffsets map[string]uint, outputOffset uint) (map[string]uint, uint, error) {

	trip.Departures[0].Ordinality = ORIGIN
	trip.Departures[len(trip.Departures)-1].Ordinality = TERMINUS

	bytes, err := bare.Marshal(&trip)
	if err != nil {
		return tripsOffsets, outputOffset, fmt.Errorf("while marshalling: %w", err)
	}
	b, err := result.Write(bytes)
	if err != nil {
		return tripsOffsets, outputOffset, fmt.Errorf("while writing: %w", err)
	}
	tripsOffsets[trip.Id] = outputOffset
	outputOffset += uint(b)

	return tripsOffsets, outputOffset, nil
}

func sortOutIndex(fileName string) error {
	file, err := os.Open(fileName)
	if err != nil {
		return fmt.Errorf("while opening file: %w", err)
	}
	defer file.Close()

	result, err := os.Create(fileName + "2")
	if err != nil {
		return fmt.Errorf("while creating file: %w", err)
	}
	defer result.Close()

	scanner := bufio.NewScanner(file)
	scanner.Scan()
	headerLine := scanner.Text()

	if err := scanner.Err(); err != nil {
		return fmt.Errorf("while scanning: %w", err)
	}
	inputChan := make(chan string)
	go func() {
		for scanner.Scan() {
			inputChan <- scanner.Text()
		}
		close(inputChan)
	}()

	sorter, outputChan, errChan := extsort.Strings(inputChan, nil)
	sorter.Sort(context.Background())

	result.WriteString(headerLine + "\n")
	for data := range outputChan {
		result.WriteString(data + "\n")
	}

	if err := <-errChan; err != nil {
		return fmt.Errorf("while sorting: %w", err)
	}

	result.Sync()
	result.Close()
	os.Rename(fileName, fileName+"_xxx")
	err = os.Rename(fileName+"2", fileName)
	if err != nil {
		return fmt.Errorf("while replacing file: %w", err)
	}
	return nil

}

func sortTripsThroughStop(c feedConverter) error {
	return sortOutIndex(filepath.Join(c.TmpFeedPath, "tripsthroughstop.csv"))
}

// TODO out to separate file
func readCsvLine(r *os.File, offset int64, fields int) ([]string, error) {
	if offset != -1 {
		r.Seek(offset, io.SeekStart)
	}
	line := []byte{}
	for {
		b := make([]byte, 1)
		_, err := r.Read(b)
		if err != nil {
			if err == io.EOF {
				break
			}
			return []string{}, fmt.Errorf("while reading byte: %w", err)
		}
		if b[0] == '\n' {
			break
		}
		line = append(line, b[0])
	}

	if string(line) == "" {
		return []string{}, io.EOF
	}
	// TODO unnecessary overhead, should parse single csv line, expecting $fields fields
	csvReader := csv.NewReader(strings.NewReader(string(line)))
	csvReader.FieldsPerRecord = fields
	record, err := csvReader.Read()
	if err != nil {
		log.Printf("fields: %d\nline: %s\nrecord:%v\n", fields, string(line), record)
		return record, fmt.Errorf("while reading record: %w", err)
	}

	return record, nil
}