#!/usr/bin/python

# Copyright (C) 2022 Yuchen Pei.
#
# ptv.py is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.

# ptv.py is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public
# License along with ptv.py.  If not, see
# <https://www.gnu.org/licenses/>.

# PTV (Public Transport Victoria) timetable cli query tool

from hashlib import sha1
from urllib import request, error, parse
from datetime import datetime, timezone
from zoneinfo import ZoneInfo
import hmac
import json
import pprint
import itertools
import os

_ROUTE_TYPE = ["Train", "Tram", "Bus", "Vline", "Night Bus"]

def get_url(request):
    key = str.encode(os.environ.get('PTVKEY'))
    dev_id = os.environ.get('PTVID')
    request = parse.quote(request) + ('&' if ('?' in request) else '?')
    raw = request + f'devid={dev_id}'
    hashed = hmac.new(key, raw.encode('utf-8'), sha1)
    signature = hashed.hexdigest().upper()
    return f'http://timetableapi.ptv.vic.gov.au{raw}&signature={signature}'

def get_json(url):
    try:
        document = request.urlopen(url).read().decode()
    except error.HTTPError:
        print("http error!")
        return
    return json.loads(document)

def get_stops(search_result):
    return [{'routes': stop['routes'],
             'stop_name': stop['stop_name'],
             'stop_id': stop['stop_id']} for stop in search_result(stops)]

def get_stop_and_routes(search_result):
    return itertools.chain.from_iterable(
        [[(stop, route) for route in stop['routes']]
            for stop in search_result['stops']])

def search(keyword):
    return get_json(get_url(f'/v3/search/{keyword}'))

def format_stop_and_route_name(stop_and_route):
    stop, route = stop_and_route
    return 'Stop: {}, Route: {} {} {}'.format(
        stop['stop_name'], _ROUTE_TYPE[route['route_type']],
        route['route_number'], route['route_name'])

def format_stop_and_route_names(stop_and_routes):
    return '\n'.join(f'[{i}] {result}' for i, result in enumerate(
        map(format_stop_and_route_name, stop_and_routes)))
    return '\n'.join(f'[{i}] {result}' for i, result in enumerate(
        [format_stop_and_route_name(
            stop['stop_name'], _ROUTE_TYPE[route['route_type']],
            route['route_number'], route['route_name'])
         for stop, route in stop_and_routes]))

def get_departures(stop, route):
    return get_json(get_url(
        '/v3/departures/route_type/{}/stop/{}/route/{}'.format(
            route['route_type'], stop['stop_id'],
            route['route_id'])))

def parse_time(maybe_time):
    if maybe_time:
        return datetime.fromisoformat(maybe_time[:-1] + '+00:00')

def format_time(maybe_time):
    if maybe_time:
        return str(maybe_time.astimezone(
            ZoneInfo('Australia/Melbourne')))[:-6]

def filter_departures(departures):
    return [dep for dep in departures
            if dep['scheduled_departure_utc'] and
            parse_time(dep['scheduled_departure_utc']) >
            datetime.now().astimezone(timezone.utc)]

def get_directions(route_id):
    return get_json(get_url(f'/v3/directions/route/{route_id}'))

def get_direction_names(route_id):
    return {dir['direction_id']: dir['direction_name'] for dir in
            get_directions(route_id)['directions']}

def format_departures(departures, direction_names):
    return '\n'.join(['estimated: {}; scheduled: {}; direction: {}'.format(
        format_time(parse_time(dep['estimated_departure_utc'])),
        format_time(parse_time(dep['scheduled_departure_utc'])),
        direction_names[dep['direction_id']])
        for dep in filter_departures(departures['departures'])])

def main():
    stop_and_routes = list(get_stop_and_routes(search(input('Query: '))))
    if not stop_and_routes:
        print("No results")
        return
    print(format_stop_and_route_names(stop_and_routes))
    idx = int(input('Choose a number: '))
    print(format_stop_and_route_name(stop_and_routes[idx]))
    print(format_departures(
        get_departures(*stop_and_routes[idx]),
        get_direction_names(stop_and_routes[idx][1]['route_id'])))

if __name__ == "__main__":
    main()