#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import logging
import coloredlogs
import requests
import os
import json
import zipfile
import datetime
import dateutil.parser
import csv
import shutil
import glob
import sys
import shutil

if not sys.version_info >= (3, 8):
    print("Python >= 3.8 is required")
    exit(1)

URL_API_GTFS = "https://data.explore.star.fr/api/records/1.0/search?apikey=3324de385f6b7b5bdfa1727d3e1a9f4bed5dd1accb288db6df5e8487&dataset=tco-busmetro-horaires-gtfs-versions-td&q=&sort=-debutvalidite"
URL_API_PICTO_BUS = "https://data.explore.star.fr/api/records/1.0/search/?apikey=3324de385f6b7b5bdfa1727d3e1a9f4bed5dd1accb288db6df5e8487&dataset=tco-bus-lignes-pictogrammes-dm&q=&rows=200&sort=idligne&refine.resolution=1%3A22"
URL_API_PICTO_METRO = "https://data.explore.star.fr/api/records/1.0/search/?apikey=3324de385f6b7b5bdfa1727d3e1a9f4bed5dd1accb288db6df5e8487&dataset=tco-metro-lignes-pictogrammes-dm&q=&rows=200&sort=idligne&refine.resolution=1%3A22"
URL_PICTO_PREFIX_BUS = (
    "https://data.explore.star.fr/explore/dataset/tco-bus-lignes-pictogrammes-dm/files/"
)
URL_PICTO_PREFIX_METRO = "https://data.explore.star.fr/explore/dataset/tco-metro-lignes-pictogrammes-dm/files/"


class Station(object):
    """
    Classe qui représente une station.
    Utilisé pour transformer un object en json
    """

    def __init__(self, sid, name=None, lon=None, lat=None):
        self.lat = lat
        self.lon = lon
        self.sid = sid
        self.name = name
        self.terminus = None
        self.limits = None
        self.lines = None
        self.parcours = None

    def __str__(self):
        o = {}
        o["type"] = "Feature"
        o["geometry"] = {}
        o["geometry"]["type"] = "Point"
        o["geometry"]["coordinates"] = [self.lon, self.lat]
        o["properties"] = {}
        o["properties"]["id"] = self.sid
        o["properties"]["name"] = self.name
        o["properties"]["terminus"] = list(self.terminus)
        o["properties"]["limits"] = self.limits
        o["properties"]["lines"] = list(self.lines)
        o["properties"]["parcours"] = list(self.parcours)
        return json.dumps(o)

    def __repr__(self):
        return self.__str__()


def get_gtfs_dir(today):
    """
    Récupère le repertoire GTFS le plus récent (en fonction de la publication)
    qui correspond à la date passé en paramètre.
    Si il n'y a pas de de données GTFS disponible retourne None.
    """
    gtfs_path = "./GTFS/"
    available_dir = []
    for filename in os.listdir(gtfs_path):
        logger.debug("Found directory: " + filename)
        gtfs_dir = gtfs_path + filename
        if os.path.isdir(gtfs_dir):
            start = datetime.datetime.strptime(
                filename.split("_")[0], "%Y-%m-%d"
            ).date()
            end = datetime.datetime.strptime(filename.split("_")[1], "%Y-%m-%d").date()
            if today >= start and today <= end:
                available_dir.append(gtfs_dir)
    available_dir_sorted = sorted(available_dir, key=lambda e: e.split("_")[2])
    if len(available_dir_sorted) > 0:
        gtfs_dir = available_dir_sorted[-1]
        logger.debug("Use directory: " + gtfs_dir)
        return gtfs_dir
    else:
        logger.debug("No GTFS data for date: " + str(today))
        return


#### Generateur de données ###
# Recodé en python depuis le générateur original en nodejs
# Ajout des parcours en plus pour chaque station


def get_serviceid_today(gtfs_dir, today):
    services = []
    csv_calendar = list(csv.DictReader(open(gtfs_dir + "/calendar.txt")))
    weekdays = list(csv_calendar[0].keys())[1:8]
    for i in csv_calendar:
        start_date = datetime.datetime.strptime(i["start_date"], "%Y%m%d").date()
        end_date = datetime.datetime.strptime(i["end_date"], "%Y%m%d").date()
        weekday = weekdays[today.weekday()]
        if i[weekday] == "1" and today >= start_date and today <= end_date:
            services.append(i["service_id"])

    csv_calendar_dates = list(csv.DictReader(open(gtfs_dir + "/calendar_dates.txt")))
    for i in csv_calendar_dates:
        c_date = datetime.datetime.strptime(i["date"], "%Y%m%d").date()
        if today == c_date:
            services.append(i["service_id"])

    return services


def get_trips_from_services(gtfs_dir, services_today):
    trips = {}
    csv_trips = list(csv.DictReader(open(gtfs_dir + "/trips.txt")))
    for i in csv_trips:
        if i["service_id"] in services_today:
            info = {}
            info["route_id"] = i["route_id"]
            info["shape_id"] = i["shape_id"]
            info["sens"] = i["direction_id"]
            info["last_sequence"] = "1"
            trips[i["trip_id"]] = info

    return trips


def generate_stop_map(gtfs_dir):
    all_stops = {}

    csv_stops = list(csv.DictReader(open(gtfs_dir + "/stops.txt")))
    for i in csv_stops:
        stop_detail = {}
        stop_detail["id"] = i["stop_id"]
        stop_detail["name"] = i["stop_name"]
        stop_detail["lat"] = i["stop_lat"]
        stop_detail["lon"] = i["stop_lon"]
        stop_detail["trips"] = {}
        all_stops[i["stop_id"]] = stop_detail

    return all_stops


def associate_trips(gtfs_dir, all_stops, trips):

    csv_stop_times = list(csv.DictReader(open(gtfs_dir + "/stop_times.txt")))
    for i in csv_stop_times:
        if i["trip_id"] in trips:
            stop_info = {}
            stop_info["arrival_time"] = i["arrival_time"]
            stop_info["departure_time"] = i["departure_time"]
            stop_info["route_id"] = trips[i["trip_id"]]["route_id"]
            stop_info["shape_id"] = trips[i["trip_id"]]["shape_id"]
            stop_info["sens"] = trips[i["trip_id"]]["sens"]
            trips[i["trip_id"]]["last_sequence"] = str(
                max(int(trips[i["trip_id"]]["last_sequence"]), int(i["stop_sequence"]))
            )
            stop_info["stop_sequence"] = i["stop_sequence"]
            all_stops[i["stop_id"]]["trips"][i["trip_id"]] = stop_info

    return all_stops, trips


def filter_no_trips(all_stops):
    newDict = {}
    for (k, v) in all_stops.items():
        if len(all_stops[k]["trips"]) != 0:
            newDict[k] = v
    return newDict


def finalize_stops(all_stops, trips):
    all_stops = filter_no_trips(all_stops)
    all_stations = []
    for stop_key, stop_value in all_stops.items():
        station = Station(stop_key)
        station.lat = stop_value["lat"]
        station.lon = stop_value["lon"]
        station.name = stop_value["name"]
        terminus_dict = {}
        station.terminus = set()
        first_trip = next(iter((stop_value["trips"].items())))
        first_arrival_time = first_trip[1]["arrival_time"]
        last_departure_time = first_trip[1]["departure_time"]
        station.lines = set()
        station.parcours = set()
        for (k, v) in stop_value["trips"].items():
            if v["stop_sequence"] == trips[k]["last_sequence"]:
                terminus_dict[k] = v
            arrival_time = v["arrival_time"]
            departure_time = v["departure_time"]
            first_arrival_time = min(first_arrival_time, arrival_time)
            last_departure_time = max(last_departure_time, departure_time)
            station.lines.add(v["route_id"] + "," + v["sens"])
            if len(v["shape_id"]):
                station.parcours.add(v["shape_id"])
        for (k, v) in terminus_dict.items():
            station.terminus.add(v["route_id"] + "," + v["sens"])
        station.limits = [first_arrival_time, last_departure_time]
        all_stations.append(station)
    return all_stations


### Fin du générateur


def stations():
    """
    Génère toutes les données geojson possible depuis aujourd'hui
    jusqu'à la fin des données GTFS disponible.
    Si les données des stations ont déjà été generées et qu'elles proviennt de données
    GTFS qui ont la même date de publication on ne les génère pas à nouveau.
    """

    os.makedirs("./GTFS/", exist_ok=True)
    os.makedirs("./data/", exist_ok=True)
    today = datetime.datetime.now().date()
    while gtfs_dir := get_gtfs_dir(today):
        logger.debug("Generate data for " + str(today))
        publication_gtfs = gtfs_dir.split("_")[-1]
        publication_gtfs_date = datetime.datetime.fromtimestamp(
            int(publication_gtfs)
        ).date()
        logger.debug("Use %s (%s)", gtfs_dir, str(publication_gtfs_date))
        filename = "./data/" + today.strftime("stations_%Y_%m_%d.geojson")
        if os.path.exists(filename):
            logger.debug(filename + " already exists")
            with open(filename + ".info", "r") as f:
                present_file_info = json.load(f)
            present_file_publication = present_file_info["publication"]
            logger.debug(
                "File is generated from data published on " + present_file_publication
            )
            present_file_publication_date = datetime.datetime.fromtimestamp(
                int(present_file_publication)
            ).date()
            if present_file_publication_date == publication_gtfs_date:
                logger.debug("stations file is up to date")
                today += datetime.timedelta(days=1)
                continue
        services_today = get_serviceid_today(gtfs_dir, today)
        trips = get_trips_from_services(gtfs_dir, services_today)
        all_stops = generate_stop_map(gtfs_dir)
        all_stops, trips = associate_trips(gtfs_dir, all_stops, trips)
        all_stations = finalize_stops(all_stops, trips)
        with open(filename, "w") as f:
            f.write(str(all_stations))
        with open(filename + ".info", "w") as f:
            info = {"publication": publication_gtfs}
            json.dump(info, f)
        today += datetime.timedelta(days=1)
    # import ipdb; ipdb.set_trace()


def download_file(url, output_filename):
    """
    Helper: télécharge un fichier
    """
    logger.info("Download: %s", url)
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(output_filename, "wb") as f:
            for chunk in r.iter_content(chunk_size=8192):
                # If you have chunk encoded response uncomment if
                # and set chunk_size parameter to None.
                # if chunk:
                f.write(chunk)


def picto():
    """
    Télécharge tous les picto disponibles.
    Si ils sont déjà présent on ne fait rien.
    """

    for param in [
        (URL_API_PICTO_BUS, URL_PICTO_PREFIX_BUS),
        (URL_API_PICTO_METRO, URL_PICTO_PREFIX_METRO),
    ]:
        r = requests.get(param[0])
        data = r.json()
        records = data["records"]
        logger.debug(json.dumps(records, indent=4))
        for entry in records:
            filename = "./images/picto_" + entry["fields"]["idligne"] + ".png"
            if not os.path.exists(filename):
                url = param[1] + entry["fields"]["image"]["id"] + "/download/"
                download_file(url, filename)
            assert os.path.exists(filename)


def download():
    """
    Télécharge les fichier zip des données GTFS puis les extrait.
    Si le fichier est déjà téléchargé on ne fait rien.
    """

    os.makedirs("./downloads/", exist_ok=True)
    os.makedirs("./GTFS/", exist_ok=True)
    r = requests.get(URL_API_GTFS)
    data = r.json()
    records = data["records"]
    logger.debug(json.dumps(records, indent=4))
    for entry in records:
        # Le dataset contient aussi des flux NETEX (fichiers XML) qui n'ont pas
        # la même structure que le GTFS: on les ignore.
        typeflux = entry["fields"].get("typeflux", "GTFS")
        if typeflux != "GTFS":
            logger.debug("Skip %s record: %s", typeflux, entry["fields"]["id"])
            continue
        filename = "./downloads/" + entry["fields"]["fichier"]["filename"]
        filename_info = filename + ".json"
        with open(filename_info, "w") as f:
            f.write(json.dumps(entry, indent=4))
        size = entry["fields"]["tailleoctets"]
        if not os.path.exists(filename) or os.path.getsize(filename) != size:
            url = entry["fields"]["url"]
            output_filename = "./downloads/" + url.split("/")[-1]
            download_file(url, output_filename)
        assert os.path.exists(filename)
        assert os.path.getsize(filename) == size
        publication_timestamp = dateutil.parser.parse(
            entry["fields"]["publication"]
        ).timestamp()
        path_extract = (
            "./GTFS/"
            + entry["fields"]["debutvalidite"]
            + "_"
            + entry["fields"]["finvalidite"]
            + "_"
            + str(int(publication_timestamp))
        )
        if not os.path.exists(path_extract):
            os.mkdir(path_extract)
        if not os.listdir(path_extract):
            with zipfile.ZipFile(filename) as zf:
                zf.extractall(path_extract)


def clean():
    """
    Supprime toutes les données produites par ce programme.
    """

    shutil.rmtree("./downloads")
    shutil.rmtree("./GTFS")
    for picto in glob.glob("./images/picto_*"):
        os.remove(picto)
    for stations in glob.glob("./data/stations_*"):
        os.remove(stations)


def prune():
    """
    Supprime les fichier zip, les données GTFS et les stations si elles ne sont
    plus valide (avec un delta de deux jours).
    """

    today = (datetime.datetime.now() - datetime.timedelta(days=2)).date()
    for zipfile in glob.glob("./downloads/*.zip"):
        zipfileinfo = zipfile + ".json"
        with open(zipfileinfo, "r") as f:
            info = json.load(f)
        fin = dateutil.parser.parse(info["fields"]["finvalidite"]).date()
        if today > fin:
            os.remove(zipfile) if os.path.exists(zipfile) else None
            os.remove(zipfileinfo) if os.path.exists(zipfileinfo) else None
            logger.debug(zipfile + " removed")
            logger.debug(zipfileinfo + " removed")
    for gtfsdir in os.listdir("./GTFS/"):
        fin = dateutil.parser.parse(gtfsdir.split("_")[1]).date()
        if today > fin:
            dir_to_remove = "./GTFS/" + gtfsdir
            shutil.rmtree(dir_to_remove)
            logger.debug(dir_to_remove + " removed")
    for stationsfile in glob.glob("./data/stations_*"):
        filename = os.path.basename(stationsfile)
        datefromfile = "-".join(filename.split(".")[:1][0].split("_")[1:])
        date = dateutil.parser.parse(datefromfile).date()
        if today > date:
            os.remove(stationsfile) if os.path.exists(stationsfile) else None
            os.remove(stationsfile + ".info") if os.path.exists(
                stationsfile + ".info"
            ) else None
            logger.debug(stationsfile + "removed")


def auto():
    """
    Execute toutes les tâches nécéssaire au fonctionnement du service.
    Commande à executer chaque jour avec tâche cron.
    """

    picto()
    download()
    stations()
    prune()


if __name__ == "__main__":
    # fmt: off
    description_str = ('ouestmonbus.com GTFS downloader \n' +
                       'and data generator\n')
    # fmt: on
    usage_str = "%(prog)s action"
    epilog_str = "v 0.1"
    frm = argparse.RawTextHelpFormatter
    parser = argparse.ArgumentParser(
        description=description_str,
        epilog=epilog_str,
        usage=usage_str,
        formatter_class=frm,
    )

    parser.add_argument(
        "--log",
        help="log level",
        default="INFO",
    )
    cmd_parser = parser.add_subparsers(dest="command", help="command parser")
    subparser = cmd_parser.add_parser("download", help="download GTFS")
    subparser = cmd_parser.add_parser("stations", help="generate stations geojson")
    subparser = cmd_parser.add_parser("picto", help="download picto")
    subparser = cmd_parser.add_parser("clean", help="delete all")
    subparser = cmd_parser.add_parser("prune", help="delete expired data")
    subparser = cmd_parser.add_parser(
        "auto", help="execute all needed task to use with cron"
    )
    subparser.add_argument(
        "--days",
        type=int,
        default=1,
        help="nb days to generate",
    )

    args = parser.parse_args()
    coloredlogs.install(level=args.log)
    logger = logging.getLogger(__name__)
    os.chdir(os.path.dirname(__file__))
    locals()[args.command]()
