#!/usr/bin/env python3

import click
import click_completion
import click_completion.core

from collections import defaultdict
import re
import sys
import os
import getopt
import requests
import subprocess
import datetime

DATE_FORMAT = '%d.%m.%Y'
PATH = os.path.dirname(os.path.abspath(__file__))
click_completion.init()

@click.group()
def completion():
    pass

@completion.command()
@click.option('--append/--overwrite', help="Append the completion code to the file", default=None)
@click.option('-i', '--case-insensitive/--no-case-insensitive', help="Case insensitive completion")
@click.argument('shell', required=False, type=click_completion.DocumentedChoice(click_completion.core.shells))
@click.argument('path', required=False)
def install(append, case_insensitive, shell, path):
    shell, path = click_completion.core.install(shell=shell, path=path, append=append)
    click.echo('%s completion installed in %s' % (shell, path))

class dotdict(dict):
    __getattr__ = dict.get
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__
 
def getStations():
    file = open(os.path.join(PATH, 'stations.in'), 'r')
    stationById = defaultdict(list)
    idByStation = defaultdict(list)
    all = []
    for string in file:
        l = string[:-1].rsplit(' ', 1)
        stationById[l[1]].append(l[0])
        idByStation[l[0]].append(l[1])
        all.append(l[0])
    file.close()
    return stationById, idByStation, all

def getMins(date):
    return int(date[0:2]) * 60 + int(date[3:5])

def getDate(mins):
    hours = mins // 60
    mins %= 60
    res = ''
    if hours != 0: res += str(hours) + ' ч'
    if mins != 0: res += str(mins) + ' м'
    return re.sub(r'ч(\d)', r'ч \1', res)

dayReplacers = [ ('по пятницам и выходным', 'ПСВ'), ('по пятницам и воскресеньям', 'ПВ'), ('по пятницам и субботам', 'ПС'), ('ежедневно', ''), ('по выходным', 'СВ'), ('по рабочим', 'КСВ'), ('по воскресеньям', 'В'), ('по субботам', 'С'), ('по пятницам', 'П'), ('кроме пятн\. и субб\.', 'КПС'), ('кроме пятниц и вых\.', 'КПСВ'), ('кроме четвергов и вых\.', 'КЧСВ'), ('кроме суббот', 'КС'), ('кроме воскресений', 'КВ') ]
placeReplacers = [ ('Красное Село', 'КрСело'), ('Лигово', ''), ('Калище', ''), ('Лебяжье', ''), ('Ораниенбаум-1', ''), ('Новый Петергоф', ''), ('Гатчина Балтийская', 'Г'), ('Гатчина Варшавская', 'Г'), ('Зеленогорск', 'Зел'), ('Кирилловское', 'Кир'), ('Каннельярви', 'Кан'), ('Рощино', 'Рощ'), ('Выборг', 'Выб'), ('Гаврилово', 'Гав'), ('Советский', 'Сов'), (r'Санкт-Петербург-.*?(\s|$)', r'\1') ]
daysSorter = { '': 0, 'К': 1, 'П': 2, 'С': 3, 'В': 4 }

def replacer(s, repl):
    for what, how in repl: s = re.sub(what, how, s)
    return s

def getTimetable(fro, to, date):
    for idFrom in fro:
        for idTo in to:
            try:
                x = requests.get('https://www.tutu.ru/spb/rasp.php?st1={}&st2={}&json&date={}'.format(idFrom, idTo, date)).json()
                if 'error' in x:
                    print(x['error'], file = sys.stderr)
                else:
                    return x
            except:
                print('No route from {} to {}'.format(idFrom, idTo))
    raise Exception('No route found:(')

def parse(directory, fro, to, date, full, collapse, remove, st = getStations()):
    stationById, idByStation, _ = st
    x = getTimetable(idByStation[fro], idByStation[to], date)

    filename = directory + stationById[x['dep-st']][0] + '—' + stationById[x['arr-st']][0]
    if not (x['dat'] is None):
        filename += ', ' + x['dat']
    filename = re.sub(r'(Санкт-Петербург)-.*?\.', r'\1', filename)
    sys.stdout = open(filename + '.out', 'w')

    trains = []
    for y in x['tra-list']:
        train = dotdict(
            departure = y['tra']['dep']['tim'],
            arrival   = y['tra']['arr']['tim'],
            fro       = stationById[y['tra']['dep']['st']][0],
            to        = stationById[y['tra']['arr']['st']][0],
            schedule  = replacer(y['tra']['sch'], dayReplacers),
            type      = '★' if y['tra']['typ'] == 'Ласточка' else '',
            minutes   = getMins(y['tra']['tr-tim']),
            change    = y['tra']['cha'].replace('. Уточните дату поездки', ''),
            skip      = False
        )
        if train.schedule.find('отменен') != -1:
            continue
        trains.append(train)

    trains = sorted(trains, key = lambda x: (x.departure, daysSorter[x.schedule[:1].upper()]))
    averageTime = sum(map(lambda train: train.minutes, trains)) / len(trains)
    for train in trains:
        if full:
            train.time = getDate(train.minutes)
        else:
            if len(trains) > 10 and train.minutes + 4 < averageTime:
                train.time = '' if remove else getDate(train.minutes).replace(' ', '')
            elif train.minutes - 10 > averageTime:
                train.time = getDate(train.minutes)
                train.skip = True
            else:
                train.time = ''
                train.skip = train.type == '★'

    if remove: trains = list(filter(lambda train: not train.skip, trains))

    sameDep, sameDst = True, True
    if collapse:
        trains2 = []
        lastTime, lastDep = 0, 0
        for train in trains:
            if train.departure == lastDep and abs(train.minutes - lastTime) <= 2 and collapse:
                trains2[-1].fro      += '/' + train.fro
                trains2[-1].to       += '/' + train.to
                trains2[-1].schedule += '/' + train.schedule
                trains2[-1].type     += '/' + train.type
            else:
                trains2.append(train)
                lastDep, lastTime = train.departure, train.minutes
        trains = trains2

        for train in trains:
            if not full:
                train.time = train.type + ' ' + train.time
                train.time = re.sub(r'^\s*(.*?)\s*$', r'\1', train.time)
                train.fro  = replacer(train.fro, placeReplacers)
                train.to   = replacer(train.to, placeReplacers)

            train.fro  = re.sub(r'^(\w+)(/\1)+$', r'\1', train.fro)
            train.to   = re.sub(r'^(\w+)(/\1)+$', r'\1', train.to)
            train.time = re.sub(r'^\s*/*\s*$', '', train.time)
            train.fro  = re.sub(r'^\s*/*\s*$', '', train.fro)
            train.to   = re.sub(r'^\s*/*\s*$', '', train.to)
            if train.fro == '' and train.to == '' and (train.schedule == 'КСВ/СВ' or train.schedule == 'СВ/КСВ'): train.schedule = ''
            train.time = re.sub('^/', '❤/', train.time)
            train.time = re.sub('/$', '/❤', train.time)
            train.time = re.sub('//', '/❤/', train.time)
        for i in range(1, len(trains)):
            sameDep = sameDep and (trains[i].fro == trains[i - 1].fro)
            sameDst = sameDst and (trains[i].to  == trains[i - 1].to )

    if full:
        fields = ['departure', 'arrival', 'schedule', 'type', 'time', 'fro', 'to', 'change']
    else:
        fields = ['departure', 'schedule', 'time']
        if not sameDep:
            fields.append('fro')
        if not sameDst:
            fields.append('to')

    resultedFields = []
    for f in fields:
        isNeeded = False
        for train in trains:
            isNeeded = isNeeded or train[f] != ''
        if isNeeded:
            resultedFields.append(f)

    for train in trains:
        print('\t'.join(map(lambda f: train[f], resultedFields)))

    sys.stdout.close()
    subprocess.Popen(['mousepad', filename + '.out'])

@completion.command()
@click.option('-d', '--date',      type=click.DateTime(formats=[DATE_FORMAT]), help='load timetable for date=DATE')
@click.option('-o', '--directory', type=click.Path(file_okay=False),           help='output directory for timetable files',            default='~/.cache/timetable')
@click.option('-t', '--today',     is_flag=True,                               help='load timetable for today',                        default=False)
@click.option('-m', '--tomorrow',  is_flag=True,                               help='load timetable for tomorrow',                     default=False)
@click.option('-f', '--full',      is_flag=True,                               help='do not minify',                                   default=False)
@click.option('-c', '--collapse',  is_flag=True,                               help='do not collapse trains with same departure time', default=True)
@click.option('-r', '--remove',    is_flag=True,                               help='remove slow trains and normal lastochkas',        default=True)
@click.argument('stations', type=click.Choice(getStations()[2], case_sensitive=False), required=True, nargs=-1)
def load(stations, directory, today, tomorrow, date, full, collapse, remove):
    if directory[-1] != os.sep:
        directory += os.sep
    if directory[0] == '~':
        directory = os.path.expanduser('~') + directory[1:]
    date = 'all' if date is None else date.strftime(DATE_FORMAT)
    if today:
        date = datetime.date.today().strftime(DATE_FORMAT)
    if tomorrow:
        date = (datetime.date.today() + datetime.timedelta(days = 1)).strftime(DATE_FORMAT)
    click.echo('Получаю расписание на ' + date)
    if full:
        collapse = False
    if len(stations) % 2 == 1:
        raise click.ClickException('Count of stations must be even')
    os.makedirs(directory, exist_ok=True)
    for i in range(0, len(stations), 2):
        parse(directory, stations[i], stations[i + 1], date, full, collapse, remove)

if __name__ == "__main__":
    completion()
