"""
Query
-----

Module to handle database IO while abstracting the specific SQL and django model filtering.

Functions
---------
get_field_names(model: django.db.models.base.ModelBase) -> list[str]:
    Return all the model's field names as a list of strings
get_pks_from_get(req_get: dict[str,str]) -> dict[django.db.models.base.ModelBase, list[str]]:
    For given GET data return the primary keys requested mapped to the model classes
get_obj_by_pk(mdl: django.db.models.base.ModelBase, pks: list[str]) -> list[mdl]:
    For a given model class and list of primary keys of that class, return the corresponding model objects.
obj_from_get(req_get: django dict-like request.GET) -> dict[django.db.base.models.ModelBase, list[models.Model]):
    For given GET data, return a dict mapping the model classes to the requested instances.
get_trips(routes: list[pt_map.models.Route]) -> dict[str, list[pt_map.models.Trip]]:
    For a list of Routes, return a dict mapping a list of all Trips associated with an individual Route, mapped to its route_id.
get_stop_sequences(routes: list[pt_map.models.Route], [trips: dict[str, list[pt_map.models.Trip]]]) -> dict[str, list[pt_map.models.Stop]]:
    For a list of routes, return a dict mapping an ordered sequence of stops to the Routes route_id. If a dict of lists trips mapped to the respective route_ids is given, the first elements of each of those lists is used, otherwise it is calculated. If the trips argument is given, it is expected to hold a value for all the Routes in routes.
get_timetable(r: pt_map.models.Route, [trips_r: list[pt_map.models.Trip]], [stop_sequence: list[pt_map.models.Stop]]) -> :
    Calculate a timetable for a Route r. If a list of trips for that route and/or a stop sequence is given, their first elements are used, otherwise all are calculated and the first elements are used.
get_all_stops() -> dict[str, dict[str, str]]:
    Return a dict mapping all stop_ids in the database to a dict of relevant values.
get_all_routes() -> list[dict[str,str]]:
    Return a list of dicts containing route_ids and other relevant values of all Routes in the database.
"""

# This file is part of transport-accessibility.
# Copyright (C) 2024 Janek Kiljanski, Johannes Randerath
# 
# transport-accessibility is free software: you can redistribute it and/or modify it under the terms of the 
# GNU General Public License as published by the Free Software Foundation, either version 3 
# of the License, or (at your option) any later version.
# 
# transport-accessibility 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 General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along with transport-accessibility. 
# If not, see <https://www.gnu.org/licenses/>. 

import django
from pt_map.models import *
from django.db import models
from .class_names import *
from django.http import HttpRequest
from datetime import datetime, date

def get_field_names(model: models.Model) -> list[str]:
    """
    Given a model, returns a list of the name strings of all the model's fields.
    """
    return [field.name for field in model._meta.fields]

def get_pks_from_get(req_get: dict[str, str]) -> dict[models.base.ModelBase, list[str]]:
    """
    Extract primary keys from a request.GET dict-like and find the corresponding classes.

    Parameters
    ----------
    req_get : dict[str, str]
        dict-like object, a HTTP Requests GET data.

    Returns
    -------
    : dict[django.db.models.base.ModelBase, list[str]]
        dict mapping a model to a list of id fields passed with GET
    """
    result = {}
    for k in req_get.keys():
        if k in classes_by_primary_keys.keys():
            result[classes_by_primary_keys[k]] = req_get.getlist(k)
    if not result:
        raise ValueError("No pks found.")
    return result

def get_obj_by_pk(mdl: models.base.ModelBase, pks: list[str]) -> list[models.Model]:
    """
    Given a model, and a list of corresponding primary keys, return a list of objects of the given model identified by the given primary keys.

    Parameters
    ----------
    mdl: django.db.models.base.ModelBase
        Model class to look for
    pks: list[str]
        primary keys of the objects to return

    Returns
    -------
    : list[mdl]:
        Objects corresponding to primary keys in pk.

    Raises
    ------
    Exception: mdl.DoesNotExist:
        If at least one object from the list of pks could not be found.
    """
    return [obj for obj in [mdl.objects.get(**{primary_keys[mdl]: pk}) for pk in pks] if obj]
        
def obj_from_get(req_get: dict[str,str]) -> dict[django.db.models.base.ModelBase, list[models.Model]]:
    """
    Given the GET data of a HTTP Request, return a dict with the requested model classes as keys and lists of the requested model objects as values.
    """
    return {mdl: get_obj_by_pk(mdl, keys) for mdl, keys in get_pks_from_get(req_get).items()}

def get_trips(routes: list[pt_map.models.Route]) -> dict[str, list[pt_map.models.Trip]]:
    """
    Return a list of all Trips associated with a Route in the argument.
    
    Parameters
    ----------
    routes: list[str]:
        List of primary keys for the Routes to search the trips for.

    Returns
    -------
    : dict[str, list[pt_map.models.Trip]]:
        Keys: route_ids from parameter.
        Values: lists of corresponding trip objects.
    """
    return {r["route_id"]: [t for t in Trip.objects.filter(route_id_id=r["route_id"])] for r in routes} 

def get_stop_sequences(routes: list[pt_map.models.Route], trips: dict[str,list[pt_map.models.Trip]]=None) -> dict[str, list[str]]:
    """
    For all given routes, return a list of stops in the order of appearance along the route. The first trip in the list of trips is used to define the sequence.

    Parameters
    ----------
    routes: list[pt_map.models.Route]
        List of pt_map.models.Route to find stop sequences for.
    trips: dict[str,list[pt_map.models.Trip]]
        List of at least one trip for each Route in routes. If none, all are calculated and the first used for the sequence.

    Returns
    -------
    : dict[str, list[str]]
        Keys: 
            route_ids.
        Values: 
            Lists of stop_ids in the order of appearance in the first Trip in the routes' trips list given.
    """
    if not trips:
        trips = get_trips(routes)
    stop_sequences = {}
    for r in routes:
        seq = []
        t = trips[r["route_id"]]
        for s in StopTime.objects.filter(trip_id_id__exact=t[0].trip_id):
            seq.append(s)
        stop_sequences[r["route_id"]] = [s.stop_id.stop_id for s in sorted(seq, key=lambda st : st.stop_sequence)]
    return stop_sequences

def get_timetable(r: pt_map.models.Route, trips_r: list[pt_map.models.Trip]=None, stop_sequence: list[str]=None):
    """ 
    Given a pt_map.models.Route, calculate the timetable for all its stops.

    Parameters
    ----------
    r: pt_map.models.Route
        Route the timetable should be calculated for.
    trips_r: list[pt_map.models.Trip]
        Optional.
        List of trips travelling on the Route r.
    stop_sequence: list[str]
        Optional.
        List of stop_ids the Route r serves.
        Currently the first trip is taken as reference for stops and sequence.

    Returns
    -------
    : dict{"stop_sequence": list(str), "stop_times": dict[str, list(str)]}
        Dict containing two elements:
            "stop_sequence" : list(str):
                list of stop_ids the route serves.
            "stop_times" : dict(str, list(str)):
                dict mapping stop_ids from stop_sequence to time strings the route is serving the stop at.
    """
    if not trips_r:
        trips_r = get_trips([r])[r["route_id"]]
    if not stop_sequence:
        stop_sequence = get_stop_sequences([r])[r["route_id"]]

    timetable = {"stop_sequence": stop_sequence}
    sts = {}
    for stop in stop_sequence:
        times = []
        for t in trips_r:
            for st in StopTime.objects.filter(trip_id=t.trip_id):
                times.append(st.departure_time.strftime("%H:%M")) 
        sts[stop] = times
    timetable["stop_times"] =  sts
    return timetable

def get_all_stops() -> dict[str, dict[str,str]]:
    """
    Return all Stop object stored in the database.

    Representation of the result:
    dict:

    .. highlight:: python
    .. code-block:: python

        {
            stop_id (str): {
                    'stop_name': pt_map.models.Stop.stop_name,
                    'stop_lat': pt_map.models.Stop.stop_lat,
                    'stop_lon': pt_map.models.Stop.stop_lon,
                }
        }
    """
    return {s.stop_id: {name: getattr(s, name) for name in ['stop_name', 'stop_lat', 'stop_lon']} for s in Stop.objects.all()} 

def get_all_routes() -> list[dict[str, str]]:
    """
    Return a list of all Route objects found in the database.

    Representation of the result:
    list:

    .. highlight:: python
    .. code-block:: python

        [
                {
                    'route_id': pt_map.models.Route.route_id, 
                    'route_type': pt_map.models.Route.route_type,
                    'route_name': pt_map.models.Route.route_short_name if set else pt_map.models.Route.route_long_name,
                    'agency_id': pt_map.models.Route.agency_id.agency_id,
                }
        ]
    """
    route_name = lambda r : r.route_short_name if r.route_short_name else r.route_long_name
    return [{"route_id": r.route_id, "route_type": r.route_type, "route_name": route_name(r), "agency_id": r.agency_id.agency_id} for r in Route.objects.all()]