import requests
from datetime import datetime, timedelta
import xmltodict
from cachetools import cached, TTLCache, LRUCache

from rich.console import Console
from rich.table import Table
import fire

class Stop:
    line_colors = {
        "S1": "[turquoise2]",
        "S2": "[green3]",
        "S3": "[dark_magenta]",
        "S4": "[red1]",
        "S5": "[deep_sky_blue4]",
        "S6": "[green3]",
        "S7": "[dark_red]",
        "S8": "[gold1]",
        "S20": "[light_coral]"
    }

    def __init__(self, event):
        dp = event['dp']
        self.id = event['@id']
        try:
            self.line = f"{event['tl']['@c']}{dp['@l']}"
        except KeyError:
            self.line = '??'
        
        self.departure_time = datetime.strptime(dp['@pt'], "%y%m%d%H%M")
        self.old_departure_time = self.departure_time
        self.departure_delay = None
        self.departure_canceled = False
        self.departure_cancellation_time = None

        self.platform = dp['@pp']
        path = dp['@ppth']
        self.path = path.split('|')
        self.destination = self.path[-1]
        self.origin = None
        
        self.arrival_time = None
        if 'ar' in event:
            ar = event['ar']
            self.arrival_time = datetime.strptime(ar['@pt'], "%y%m%d%H%M")
            path = ar['@ppth']
            add_path:list = path.split('|')[:-1]
            add_path.extend(self.path)
            self.path = add_path
            self.origin = self.path[0]

        self.old_arrival_time = self.arrival_time
        self.arrival_delay = None
        self.arrival_canceled = False
        self.arrival_cancellation_time = None


    def lookup_line_color(self):
        if self.line in Stop.line_colors:
            return Stop.line_colors[self.line]
        else:
            return ""

    def __repr__(self):
        if not self.departure_canceled:
            return f"{self.line} to {self.destination}: {self.departure_time.strftime(r"%Y-%m-%d %H:%M")}{f" ({self.departure_delay:+d})" if self.departure_delay is not None else ""} @ platform {self.platform}"
        else:
            return f"{self.line} to {self.destination}: {self.departure_time.strftime(r"%Y-%m-%d %H:%M")} @ platform {self.platform} CANCELLED at {self.cancellation_time.strftime(r"%Y-%m-%d %H:%M")}"

class Timetable:
    base_url = "https://iris.noncd.db.de/iris-tts/timetable"

    def __init__(self, eva=None, number_departures=10):
        self.eva = eva
        self.station = None
        self.get_changes()
        self.stops = []
        self.min_stop_count = number_departures
        self.timestamp = None
        self._timetable = None
        pass

    @cached(cache=TTLCache(maxsize=128, ttl=30))
    def get_changes(self):
        url = f"{Timetable.base_url}/fchg/{self.eva}"
        x = requests.get(url)
        changes = xmltodict.parse(x.text)
        if self.station is None:
            self.station = changes['timetable']['@station']
        return changes['timetable']['s']
    
    def fetch_actual(self, stop: Stop):
        changes = self.get_changes()
        for event in changes:
            if event['@id'] == stop.id:
                try:
                    changed_status = event['ar']['@cs']
                    if changed_status == 'c':
                        stop.arrival_canceled = True
                        stop.arrival_cancellation_time = datetime.strptime(event['ar']['@clt'], r"%y%m%d%H%M")
                except KeyError:
                    pass
                try:
                    changed_status = event['dp']['@cs']
                    if changed_status == 'c':
                        stop.departure_canceled = True
                        stop.departure_cancellation_time = datetime.strptime(event['dp']['@clt'], r"%y%m%d%H%M")
                except KeyError:
                    pass
                
                if stop.arrival_canceled or stop.departure_canceled: 
                    return

                try:
                    stop.arrival_time = datetime.strptime(event['ar']['@ct'], r"%y%m%d%H%M")
                    stop.arrival_delay = stop.arrival_time - stop.old_arrival_time
                    stop.arrival_delay = int(stop.arrival_delay.total_seconds()/60)
                    if stop.arrival_delay == 0:
                        stop.arrival_delay = None
                except KeyError:
                    pass
                try:
                    stop.departure_time = datetime.strptime(event['dp']['@ct'], r"%y%m%d%H%M")
                    stop.departure_delay = stop.departure_time - stop.old_departure_time
                    stop.departure_delay = int(stop.departure_delay.total_seconds()/60)
                    if stop.departure_delay == 0:
                        stop.departure_delay = None
                except KeyError:
                    pass

    @staticmethod
    # planned timetable doesn't change, so just cache a lot
    @cached(LRUCache(maxsize=1024))
    def _fetch_timetable_cached(daystamp, hourstamp, eva):
        url = f"{Timetable.base_url}/plan/{eva}/{daystamp}/{hourstamp}"
        x = requests.get(url)
        if x.status_code != 200:
            return None
        timetable = xmltodict.parse(x.text)['timetable']
        return timetable

    def set_timestamp(self, ts: datetime|None):
        if ts is None:
            ts = datetime.now()
        self.timestamp = ts

    def fetch_timetable(self, ts: datetime|None = None):
        if ts is None:
            ts = self.timestamp
        daystamp = ts.strftime(r"%y%m%d")
        hourstamp = ts.strftime(r"%H")

        self._timetable = self._fetch_timetable_cached(daystamp, hourstamp, self.eva)

    def get_stops(self, ts=None):
        self.set_timestamp(ts)
        # get departures for current hour
        self.fetch_timetable()
        self.extract_stops()

        # get departures for previous hour in case there is a long delay
        self.fetch_timetable(self.timestamp - timedelta(hours=1))
        self.extract_stops()

        # filter out departures earlier than now
        split_index = -1
        for di, dep in enumerate(self.stops):
            if dep.departure_time < self.timestamp:
                split_index = di
        self.stops = self.stops[split_index+1:]

        delta = timedelta(hours=1)
        while len(self.stops) < self.min_stop_count and delta.seconds < 86400:
            self.fetch_timetable(self.timestamp + delta)
            self.extract_stops()
            delta += timedelta(hours=1)

    def extract_stops(self):
        if self._timetable is None:
            return
        events = self._timetable['s']
        if not isinstance(events, list):
            events = [events]
        for event in events:
            if event['tl']['@f'] != 'S':
                # only s-bahnen are supported for now
                continue
            # check if there is a member 'dp' in event -> train does not end at the station
            if 'dp' not in event:
                continue
            stop = Stop(event)
            self.fetch_actual(stop)

            # check if ID is already in list
            # (inefficient, but there aren't that many trips available anyways as there aren't that many hours available)
            add = True
            for stp in self.stops:
                if stp.id == stop.id:
                    add = False
            if add:
                self.stops.append(stop)
        
        self.stops.sort(key=lambda stp: stp.departure_time)

    def print_stops(self):
        table = Table(title = f"{self.station} - {self.timestamp.strftime(r"%Y-%m-%d %H:%M:%S")}")
        
        table.add_column("Line", justify="center")
        # table.add_column("Origin", justify="left")
        table.add_column("Destination", justify="left")
        
        # table.add_column("Arrival", justify="center")
        # table.add_column("Arrives in", justify="right")
        # table.add_column("Arr. Delay", justify="center")
        # table.add_column("Orig. Arr.", justify="center")

        table.add_column("Departure", justify="center")
        table.add_column("Departs in", justify="right")
        table.add_column("Delay", justify="center")
        # table.add_column("Orig. Dep.", justify="center")
        
        # table.add_column("Pl", justify="left")

        for stop in self.stops[:min(self.min_stop_count, len(self.stops))]:
            stop: Stop = stop
            line_tag = stop.lookup_line_color()

            # arr_color_tag = "[yellow]" if stop.arrival_time is not None and int((stop.arrival_time - self.timestamp).total_seconds())/60 <= 10 else ""
            # arr_timestr = (f"{stop.arrival_time.strftime(r"%H:%M")}" if not stop.arrival_canceled else "[red]cancelled") if stop.arrival_time is not None else ""
            # arr_until =  f"{int((stop.arrival_time - self.timestamp).total_seconds()/60)} min" if stop.arrival_time is not None else ""

            dep_color_tag = ""
            mins_until_departure = int((stop.departure_time - self.timestamp).total_seconds()/60)
            if mins_until_departure <= 10:
                dep_color_tag = "[yellow]"
            if mins_until_departure <= 6:
                dep_color_tag = "[red]"

            dep_timestr = f"{stop.departure_time.strftime(r"%H:%M")}" if not stop.departure_canceled else "[red]cancelled"
            dep_until =  f"{int((stop.departure_time - self.timestamp).total_seconds()/60)} min"

            table.add_row(
                f"{line_tag}{stop.line}",
                # f"{self.station if stop.origin is None else stop.origin}",
                f"{stop.destination}",

                # f"{arr_timestr}",
                # f"{arr_color_tag}{arr_until}",
                # f"{f"{stop.arrival_delay:+d}" if stop.arrival_delay is not None else ""}",
                # f"{stop.old_arrival_time.strftime(r"%H:%M") if stop.arrival_delay is not None and stop.old_arrival_time is not None else ""}",

                f"{dep_timestr}",
                f"{dep_color_tag}{dep_until}",
                f"{f"{stop.departure_delay:+d}" if stop.departure_delay is not None else ""}",
                # f"{stop.old_departure_time.strftime(r"%H:%M") if stop.departure_delay is not None else ""}",

                # f"{stop.platform}"
            )
        
        console = Console()
        console.print(table)

def fetch_and_print_departures(eva):
    tt = Timetable(eva=eva)
    tt.get_stops()
    tt.print_stops()

if __name__ == "__main__":
    fire.Fire(fetch_and_print_departures)
