#!/bin/sh

print_usage() {
  cat <<EOF
Usage: $(basename "$0") [OPTIONS] COMMAND [COMMAND_OPTIONS]

Get information about Trenitalia trains

Options:
  -h --help                             print this message

Commands:

  search FROM TO [OPTIONS]              search for trains between two stations
   -i --interactive                     allows selecting a train for more information
   -d --date DATE                       use DATE instead of now

  status TRAIN_ID [--board STATION]     get realtime information about a train
   -b --board STATION                   prints the platform number when boarding
                                        then train from a given station
EOF
  exit 1
}

# Prints realtime information about a train
train_status() {
  train="$1"
  test "$2" = "-b" -o "$2" = "--board" && board_station="$3"

  # Resolve the station name, if necessary
  if test "${board_station#id:}" = "$board_station"; then
    board_station=$(station_search "$board_station" | sed 's/^8300/S/g')
  else
    board_station=${board_station#id:}
  fi

  url='http://www.viaggiatreno.it/infomobilita/resteasy/viaggiatreno'
  today="$(date -d "$(date +%F)" +%s)000"
  station=$(curl -s "$url/cercaNumeroTrenoTrenoAutocomplete/$1" |
            grep "$today" | cut -d\| -f2 | cut -d- -f2)
  curl -s "$url/andamentoTreno/$station/$train/$today" |
  jq -r --arg board_station "$board_station" '
    # Functions to color text
    def red($x):    "\u001B[1;31m\($x)\u001B[0m";
    def green($x):  "\u001B[1;32m\($x)\u001B[0m";
    def yellow($x): "\u001B[1;36m\($x)\u001B[0m";

      (if   .tipoTreno == "PG" then green("regolare")
       elif .tipoTreno == "ST" then red("cancellato")
       elif .tipoTreno == "DV" then yellow("deviato")
       elif .tipoTreno == "PP" or
            .tipoTreno == "SI" or
            .tipoTreno == "SF" then yellow(.subTitle)
       else "sconosciuto"
       end) as $status
    | (if .nonPartito
        then []
        else ["ultima posizione: " + .stazioneUltimoRilevamento
                                   + "(\(.compOraUltimoRilevamento))"]
       end) as $position
    | (if .ritardo <= 0
         then green(.compRitardo[0])
         else red(.compRitardo[0])
       end) as $late
    | (if $board_station
        then ["binario arrivo: " +
          (.fermate[] | select(.id == $board_station)
                      | .binarioProgrammatoArrivoDescrizione)]
        else []
      end) as $platform
    | ["nome: \(.compNumeroTreno)",
       "stato: \($status)",
       "ritardo: \($late)"
      ] + $platform + $position
    | join("\n")'
}


# Searches for a station by name and returns its ID
station_search() {
  url='https://www.lefrecce.it/Channels.Website.BFF.WEB/website'
  # URL-encode the search string
  encoded=$(printf '%s' "$1" | jq -sRr '@uri')
  curl -s "$url/locations/search?name=$encoded&limit=1" | jq '.[0].id'
}


# Searches for a route between two stations
route_search() {
  url='https://www.lefrecce.it/Channels.Website.BFF.WEB/website'
  query=$(cat <<EOF
{
  "departureLocationId": "$1",
  "arrivalLocationId": "$2",
  "departureTime": "$(date -d "$3" --iso-8601=s)",
  "adults": 1,
  "children": 0,
  "criteria": {
    "frecceOnly": false,
    "regionalOnly": false,
    "noChanges": false,
    "order": "DEPARTURE_DATE",
    "limit": 6,
    "offset": 0
  },
  "advancedSearchRequest": { "bestFare": false }
}
EOF
  )

  printf 'data: \033[36m%s\033[0m\n\n' "$(date -d "$3" +'%F %R')"
  curl -s --json "$query" "$url/ticket/solutions" |
  jq -r '
  "%FT%T.%G%z" as $timefmt |

  # Functions to color text
  def red($x):    "\u001B[1;31m\($x)\u001B[0m";
  def green($x):  "\u001B[1;32m\($x)\u001B[0m";
  def blue($x):   "\u001B[1;34m\($x)\u001B[0m";

  def time:
    # Extracts "HH:MM" from a timestamp
    match("T([0-9]+:[0-9]+)"; null) | .captures[0].string;

  def delta($stop; $start):
    # Computes the time elapsed between two dates
    [$stop, $start]
    | map(strptime($timefmt) | mktime) # seconds
    | .[0] - .[1]
    | ((./60) % 60) as $m
    | (./3600 | floor) as $h
    | (if $h > 0 then "\($h)h " else "" end) + "\($m)min";

  def print_trip:
    # Prints a trip "station(time) → station(time)"
    (.duration // delta(.arrivalTime; .departureTime)) as $delta |
    "\(.origin)(\(green(.departureTime | time)))" +
    " → \(.destination)(\(.arrivalTime | time)) " +
    blue($delta);

  def print_train:
    # Prints the train name and its trip
    "\(.train.description // .train.trainCategory): " + print_trip;

  def print_summary:
    (.trains | length) as $ntrain |
    (if $ntrain > 2 then (", " + red("\($ntrain-1) cambi"))
     elif $ntrain == 2 then (", " + red("1 cambio"))
     else "" end) as $switch |
       "- " + print_trip + $switch + "\n  "
            + (.nodes | map(print_train) | join("\n  "));

  .solutions | map(.solution | print_summary) | join("\n\n")'
}

handle_search() {
  # Get station IDs
  start_id=$(station_search "$1")
  stop_id=$(station_search "$2")
  shift 2

  while :; do
    case "$1" in
      -d|--date) date=$2; shift ;;
      -i|--interactive) interactive=1; shift ;;
      -?*) print_usage ;;
        *) break
    esac
  done

  results=$(route_search "$start_id" "$stop_id" "${date:-now}")
  printf '%s\n\n' "$results"

  if test -n "$interactive"; then
    # Ask user to select a train
    info=$(printf '%s' "$results" | fzf --ansi --no-sort --disabled | cut -d: -f1)
    # Print the train status
    train_id=$(printf '%s' "$info" | cut -d: -f1 | rev | cut -d' ' -f 1 | rev)
    # convert between different API conventions
    station_id=$(printf '%s' "$start_id" | sed 's/^8300/S/g')
    train_status "$train_id" --board "id:$station_id"
  fi
}

test $# -eq 0 && print_usage

cmd="$1"; shift
case "$cmd" in
  search) handle_search "$@" ;;
  status) train_status "$@" ;;
       *) print_usage
esac