# To communicate with the compiled userland program, a simple pipe is implemented from stdout
# of a second thread, running the python program to query the MTA API, to the stdin of the compiled
# userland program. The python program will use the requests library to send HTTP GET requests to
# the API endpoint. It will do basic parsing to make sure that only relevant data is passed to 
# the userspace program that will update the display with the appropriate information.

# INTERFACE:
# Receives: route number
#
# Responds: Array of Trains for that route
# Train: {from_station: <station_name>, to_station <station_name>, direction: <"northbound">/<"southbound">
# seconds_to_arrival: <seconds>} Note that if seconds_to_arrival == 0, we can consider a train "stopped"

import requests
import socket
from datetime import datetime
import json
import math
ROUTES_URL = "https://api.wheresthefuckingtrain.com/routes"
TRAINS_BY_ROUTE_URL = "https://api.wheresthefuckingtrain.com/by-route/"
DEFAULT_ROUTES = ["1","2","3","4","5","5X","6","6X","7","7X","A","B","C","D","E","F","FS","FX","G","H","J","L","M","N","Q","R","S","W"]
DEFAULT_TRAINS = [{'from_station': 0, 'to_station': 1, 'direction': 'northbound', 'seconds_to_arrival': 0}, {'from_station': 4, 'to_station': 3, 'direction': 'northbound', 'seconds_to_arrival': 38}, {'from_station': 7, 'to_station': 6, 'direction': 'northbound', 'seconds_to_arrival': 57}, {'from_station': 9, 'to_station': 8, 'direction': 'northbound', 'seconds_to_arrival': 52}, {'from_station': 0, 'to_station': 11, 'direction': 'northbound', 'seconds_to_arrival': 0}, {'from_station': 14, 'to_station': 13, 'direction': 'northbound', 'seconds_to_arrival': 78}, {'from_station': 17, 'to_station': 16, 'direction': 'northbound', 'seconds_to_arrival': 34}, {'from_station': 0, 'to_station': 18, 'direction': 'northbound', 'seconds_to_arrival': 0}, {'from_station': 24, 'to_station': 23, 'direction': 'northbound', 'seconds_to_arrival': 51}, {'from_station': 0, 'to_station': 24, 'direction': 'northbound', 'seconds_to_arrival': 0}, {'from_station': 27, 'to_station': 26, 'direction': 'northbound', 'seconds_to_arrival': 36}, {'from_station': 31, 'to_station': 30, 'direction': 'northbound', 'seconds_to_arrival': 14}, {'from_station': 0, 'to_station': 33, 'direction': 'northbound', 'seconds_to_arrival': 0}, {'from_station': 37, 'to_station': 36, 'direction': 'northbound', 'seconds_to_arrival': 37}, {'from_station': 36, 'to_station': 37, 'direction': 'southbound', 'seconds_to_arrival': 29}, {'from_station': 0, 'to_station': 35, 'direction': 'southbound', 'seconds_to_arrival': 0}, {'from_station': 0, 'to_station': 31, 'direction': 'southbound', 'seconds_to_arrival': 0}, {'from_station': 27, 'to_station': 28, 'direction': 'southbound', 'seconds_to_arrival': 16}, {'from_station': 23, 'to_station': 24, 'direction': 'southbound', 'seconds_to_arrival': 1}, {'from_station': 21, 'to_station': 22, 'direction': 'southbound', 'seconds_to_arrival': 42}, {'from_station': 18, 'to_station': 19, 'direction': 'southbound', 'seconds_to_arrival': 45}, {'from_station': 13, 'to_station': 14, 'direction': 'southbound', 'seconds_to_arrival': 112}, {'from_station': 9, 'to_station': 10, 'direction': 'southbound', 'seconds_to_arrival': 26}, {'from_station': 5, 'to_station': 6, 'direction': 'southbound', 'seconds_to_arrival': 16}, {'from_station': 2, 'to_station': 3, 'direction': 'southbound', 'seconds_to_arrival': 20}, {'from_station': 0, 'to_station': 1, 'direction': 'southbound', 'seconds_to_arrival': 37}]
HOST = "127.0.0.1"  # Standard loopback interface address 
PORT = 1989  # Port to listen on 
WORLD_TIME_API = "http://worldtimeapi.org/api/timezone/America/New_York"

def main(): 

    # listen for route requests. We assume only one connection at a time
    s= socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind((HOST, PORT))
    s.listen()
    while True:
        conn, addr = s.accept()
        data = conn.recv(1024)
        if not data:
            continue
        route = data.decode()
        trains = []
        print("received request for route ", route.rstrip()[0])
        try:
            trains = get_trains(route.rstrip()[0])
        except Exception as e:
            print("Error trying to fetch trains ", e)
            print("Sending default trains for Line 1")
            trains = DEFAULT_TRAINS # send default set of trains for the 1 line, so we display something.
        print(trains)
        train_bytes = get_trains_bytes(trains)
        conn.sendall(train_bytes)
        conn.close()

    # test the codez here
    
    route = "6"
    trains = get_trains(route)
    get_trains_bytes(trains)

# Protocol: first 2 bytes: number of trains, following by: 2 bytes station ID, 2 bytes station ID, 1 byte char, 4 bytes of seconds to arrival     

def get_trains_bytes(trains):
    train_bytes = bytearray()
    num_trains = 0
    for train in trains:
        train_bytes.extend(train['from_station'].to_bytes(length=2, byteorder='big'))
        train_bytes.extend(train['to_station'].to_bytes(length=2, byteorder='big'))
        train_bytes.extend(train['direction'][0].encode('ascii'))
        train_bytes.extend(train['seconds_to_arrival'].to_bytes(length=2, byteorder='big'))
        num_trains += 1
        # print(train_bytes)
    print(str(num_trains) + " trains")
    train_bytes[0:0] = num_trains.to_bytes(length=2, byteorder='big')
    return train_bytes


def get_trains(route):
    response = requests.get(TRAINS_BY_ROUTE_URL + route)
    stations = {}
    for station in response.json()["data"]:
        stop_id = list(station['stops'].keys())[0]
        if (stop_id[0] == route):
            stations[stop_id] = {"name": station["name"], "southbound": station['S'], "northbound": station['N']}

    # sort the stations by stop, so we can iterate over them easily
    sorted_stations_ids = sorted(stations)
    sorted_stations = []
    short_station_id = 0
    for id in sorted_stations_ids:
        sorted_stations.append(stations[id])
        sorted_stations[short_station_id]["id"] = short_station_id
        print(short_station_id)
        short_station_id += 1

    # parse MTA station data and find. those. trains!!!!
    current_time = get_current_time()
    trains = {}
    trains[route] = []
    print("Current Time:" + current_time.strftime("%H:%M:%S"))

    trains = []
    print("===================NORTHBOUND TRAINS=====================")
    trains = trains + find_trains("northbound", sorted_stations, route, current_time)
    print("===================SOUTHBOUND TRAINS=====================")
    trains = trains + find_trains("southbound", reverse(sorted_stations), route, current_time)

    return trains

def reverse(array):
    new_array = []
    for i in range(len(array)):
        new_array.append(array[len(array) - i - 1])
    return new_array

def get_current_time():
    response = requests.get(WORLD_TIME_API)
    if response.status_code != 200:
        print("Error fetching from the time API, status code=", response.status_code)
    time = datetime.strptime(response.json()['datetime'][0:19], '%Y-%m-%dT%H:%M:%S')
    return time

def find_trains(direction, sorted_stations, route, current_time):
    i = 0
    trains = []
    arrival_time = 0
    for station in sorted_stations:
        for arrival in station[direction]:
            if not arrival['route'] == route: # skip trains on different routes
                continue 
            
            arrival_time = datetime.strptime(arrival['time'], '%Y-%m-%dT%H:%M:%S-04:00')
            break
        
        if (arrival_time == 0):
            break
        # (likely incorrect) assumption: trains "stop" 20 seconds before their arrival time so we can map trains as "stopped"
        # we do this because the API does not let us always have access to "old" arrival times
        # we could keep state but that makes this unnecessarily complicated.

        if (current_time >= arrival_time and i > 0) or (i > 0 and arrival_time > prev_arrival_time 
            and current_time < prev_arrival_time):  # we have a TRAIN!
                seconds = math.ceil((prev_arrival_time - current_time).total_seconds())

                # subtract 20 seconds for simulating "stopped" status
                seconds = seconds - 20
                if (seconds < 0):
                    seconds = 0

                if (seconds == 0):
                    print(" TRAIN! STOPPED ")
                    trains.append({"from_station": station["id"], "to_station": sorted_stations[i-1]["id"],
                        "direction": direction, "seconds_to_arrival": seconds})
                else:
                    print("")
                    print("TRAIN! Arrives in " + str(seconds) + " seconds")
                    trains.append({"from_station": station["id"], "to_station": sorted_stations[i-1]["id"],
                        "direction": direction, "seconds_to_arrival": seconds})
        else: 
            print("")        
        print(station['name'][0:6] + '\t' + arrival_time.strftime("%H:%M:%S"), end = " ")

        prev_arrival_time = arrival_time
        i += 1
    
    print("")
    return trains


if __name__=="__main__":
   main()
