path: root/src/main.rs

// See copyright information at the end of the file
pub mod stations;
use std::{process::exit, env::args, io};
use chrono::{DateTime, Local, Datelike, Timelike};
use serde::Deserialize;

#[derive(Deserialize, Debug)]
struct Station {
    station_ori_name: String,
}

#[derive(Deserialize, Debug)]
struct Train {
    train_category: Option<String>,
    train_operator: Option<String>,
    // delay: ...
    // status: Option<String>,
    train_id: Option<String>,
    average_crowding: Option<usize>,
    average_crowding_label: Option<String>,
}

#[derive(Deserialize, Debug)]
struct Stop {
    station: Station,
    #[serde(rename = "type")]
    stop_type: String, // "start", "pass" or "end"
    cancelled: bool,
    #[serde(rename = "dep_date_time")]
    departure: Option<DateTime<Local>>,
    #[serde(rename = "arr_date_time")]
    arrival: Option<DateTime<Local>>,
}

#[derive(Deserialize, Debug)]
struct Journey {
    train: Train,
    pass_list: Vec<Stop>,
}

#[derive(Deserialize, Debug)]
struct Ticket {
    arr_station: Station,
    cancelled: bool,
    // This is something like HH:MM:SS, I will need to parse it
    // arr_time: String,
    // date: String,
    // I still don't know what type this is...
    // delay: ,
    // delay_defined: bool,
    dep_station: Station,
    // See arr_time
    dep_time: String,
    // See dep_time
    duration: String,
    journey_list: Vec<Journey>,
}

const APIURL: &'static str = "https://www.trenord.it/mia/hafas";
const GREEN: &'static str = "\x1b[92m";
const YELLOW: &'static str = "\x1b[93m";
const RED: &'static str = "\x1b[31m";
const WHITE: &'static str = "\x1b[0m";

fn search_station(stations: &[&'static str], station: &str) -> Option<&'static str> {
    for s in stations {
        if *s == station.to_uppercase() {
            return Some(s);
        }
    }
    eprintln!("No station named {station}");
    None
}

fn simple_url_encode(s: &'static str) -> String {
    s.replace(' ', "%20")
}

fn search_trains(agent: &ureq::Agent, stations: &[&'static str], from: &str, to: &str, at: &DateTime<Local>) -> Option<Vec<Ticket>> {
    // https://www.trenord.it/mia/hafas?orig=ROMANO&dest=MILANO%20LAMBRATE&departure_date=20240609&departure_hour=22:30&products=tickets&transfers=1
    // Url encoding
    let from = search_station(stations, from)?;
    let to = search_station(stations, to)?;
    match agent.get(APIURL).query("orig", &simple_url_encode(from)).query("dest", &simple_url_encode(to))
        .query("departure_date", &format!("{}", at.format("%Y%m%d")))
        .query("departure_hour", &format!("{}", at.format("%H:%M")))
        .query("products", "tickets")
        .query("transfers", "1").call() {
            Ok(resp) => {
		match resp.into_json::<Vec<Ticket>>() {
                    Ok(json) => {Some(json)},
                    Err(e) => {
			eprintln!("Error deserializing response from Trenord: {e:?}");
			None
                    },
		}
            },
            Err(e) => {
		eprintln!("Network error: {e:?}");
		None
            }
	}
}

fn to_title_case(s: &str) -> String {
    assert!(s.len() > 0);
    let mut result = Vec::from(s);
    let mut prev = s.bytes().nth(0).unwrap();
    for (i, ch) in result.iter_mut().enumerate() {
	if i == 0 || prev.is_ascii_whitespace() {
	    ch.make_ascii_uppercase();
	} else {
	    ch.make_ascii_lowercase();
	}
	prev = *ch
    }
    String::from_utf8(result).unwrap()
}

fn print_ticket(t: &Ticket, tstamp: &DateTime<Local>, verbose: bool) {
    let dep_name = to_title_case(&t.dep_station.station_ori_name);
    let arr_name = to_title_case(&t.arr_station.station_ori_name);
    print!("{} → {} {:02}/{:02}/{}@{YELLOW}{}{WHITE} ({})", dep_name, arr_name, tstamp.day(), tstamp.month(), tstamp.year(), t.dep_time, t.duration);
    if t.cancelled {
        println!(" {RED}CANCELLED{WHITE}");
    } else {
        println!("");
    }
    if verbose {
        for journey in &t.journey_list {
            for stop in &journey.pass_list {
                let station = to_title_case(&stop.station.station_ori_name);
                match stop.stop_type.as_str() {
                    "start" => {
                        match &journey.train.train_id {
                            Some(id) => {
                                match &journey.train.train_category {
                                    Some(c) => {
                                        print!("\t [{} {}] ", id, c);
                                        for _i in 0..(5+3-id.len() - c.len()) {
                                            print!(" "); // alignment
                                        }
                                    },
                                    None => match &journey.train.train_operator {
                                        Some(o) => {
                                            print!("\t [{} {}] ", id, o);
                                            for _i in 0..(5+3-id.len() - o.len()) {
                                                print!(" "); // alignment
                                            }
                                        },
                                        None => {},
                                    },
                                };
                            },
                            None => {},
                        };
                        if let Some(dep_time) = stop.departure {
                            print!("{YELLOW}{:02}:{:02}{WHITE} ", dep_time.hour(), dep_time.minute());
                        }
                        print!("{GREEN}{}{WHITE} → ", station);
                    }
                    "pass" => print!("{} → ", station),
                    "end" => {
                        print!("{GREEN}{}{WHITE} ", station);
                        if let Some(arr_time) = stop.arrival {
                            print!("{YELLOW}{:02}:{:02}{WHITE} ", arr_time.hour(), arr_time.minute());
                        }
                        if stop.cancelled {
                            print!("{RED}CANCELLED{WHITE} ");
                        }
                        if let (Some(label), Some(crowding)) = (&journey.train.average_crowding_label, journey.train.average_crowding) {
                            print!("(typically ");
                            if crowding >=20 && crowding <= 50 {
                                print!("{YELLOW}{}{WHITE})", label.to_lowercase());
                            } else if crowding > 50 {
                                print!("{RED}{}{WHITE})", label.to_lowercase());
                            } else {
                                print!("{GREEN}{}{WHITE})", label.to_lowercase());
                            }
                        }
                    },
                    _ => eprintln!("Unknown stop type: {}", stop.stop_type),
                };
            }
            println!("");
        }
    }
}

fn add_day(date: &DateTime<Local>, add: u32) -> Option<DateTime<Local>> {
    let day = date.day();
    let month = date.month();
    match month {
        11|4|6|9 if day + add > 30 => match date.with_day(1 + (day + add)) {
            Some(r) => Some(r.with_month(month+1).unwrap()),
            None => {
                eprintln!("Error setting date");
		None
            },
        },
        2 if day + add > 28 => match date.with_day(1 + (day + add) % 28) {
            Some(r) => Some(r.with_month(3).unwrap()),
            None => {
                eprintln!("Error setting date");
		None
            }
        },
        1|3|5|7|8|10|12 if day + add > 31 => match date.with_day(1 + (day + add) % 31) {
            Some(r) => Some(r.with_month((month + 1) % 12).unwrap()),
            None => {
                eprintln!("Error setting date");
                None
            },
        },
        _ => match date.with_day(day + add) {
            Some(r) => Some(r),
            None => {
                eprintln!("Error calculating date for tomorrow...?!");
                None
            },
        },
    }
}

fn add_hour(date: &DateTime<Local>, hour: u32) -> Option<DateTime<Local>> {
    if date.hour() + hour < 24 {
	date.with_hour(date.hour() + hour)
    } else {
	add_day(&date, 1)?.with_hour((date.hour() + hour) % 24)
    }
}

fn parse_timestring(tstamp: &str) -> Option<DateTime<Local>> {
    let mut result = Local::now();
    for tok in tstamp.split_whitespace() {
        if tok == "tomorrow" {
	    result = add_day(&result, 1)?;
        } else if tok.starts_with('+') {
	    if tok.ends_with('h') {
		let hour = match tok[1..tok.len()-1].parse::<u32>() {
		    Ok(min) => {min},
		    Err(e) => {
			eprintln!("Cannot parse {tok}: {e}");
			return None;
		    },
		};
		result = add_hour(&result, hour)?;
	    } else if tok.ends_with('d') {
		let add = match tok[1..tok.len()-1].parse::<u32>() {
		    Ok(day) => day,
		    Err(e) => {
			eprintln!("Cannot parse {tok}: {e}");
			return None;
		    },
		};
		result = add_day(&result, add)?;
	    } else if tok.ends_with('m') {
		let minute = match tok[1..tok.len()-1].parse::<u32>() {
		    Ok(min) => min,
		    Err(e) => {
			eprintln!("Cannot parse {tok}: {e}");
			return None;
		    }
		};
		if result.minute() + minute >= 60 {
		    result = add_hour(&result, 1)?;
		}
		result = result.with_minute((result.minute() + minute) % 60)?;
	    }
        } else if tok.contains(":") {
            let mut tok_colon_iter = tok.split(':');
            let hour = match match tok_colon_iter.next() {
                Some(h) => h,
                None => {
                    eprintln!("Cannot parse HH:MM timestamp {tok}");
                    return None;
                },
            }.parse::<i32>() {
                Ok(h) => h,
                Err(e) => {
                    eprintln!("Error parsing HH:MM timestamp {tok}: {e}");
                    exit(1);
		},
            };
            let minute = match match tok_colon_iter.next() {
                Some(m) => m,
                None => {
                    eprintln!("Cannot parse HH:MM timestamp {tok}");
                    return None;
                }
            }.parse::<i32>() {
                Ok(m) => m,
                Err(e) => {
                    eprintln!("Error parsing HH::MM timestamp {tok}: {e}");
                    exit(1);
                },
            };
            if hour < 0 || hour >= 24 {
                eprintln!("Hour indication {hour} should be in the range [0,24)");
                return None;
            } else if minute < 0 || minute >= 60 {
                eprintln!("Minute indication {minute} should be in the range [0,60)");
                return None;
            }
            result = result.with_minute(minute as u32).unwrap();
            result = result.with_hour(hour as u32).unwrap();
        } else {
            eprintln!("Cannot parse time indication {tok}");
            // TODO(mario): Explain the available format
            return None;
        }
    }
    return Some(result);
}

fn usage() {
    eprintln!("USAGE: trenord [-n #trains] [-@ time] [-v] FROM TO");
}

const NUMBER_SOLUTIONS: usize = 5;

fn main() -> Result<(), io::Error> {
    let mut number_solutions = NUMBER_SOLUTIONS;
    let mut verbose = false;
    let mut args = args();
    let mut from = String::new();
    let mut to = String::new();
    let mut timestamp =  Local::now();
    if args.len() < 3 {
        usage();
        exit(1);
    }
    args.next();
    while let Some(arg) = args.next() {
        if arg.chars().nth(0).unwrap() != '-' {
            from = arg;
            match args.next() {
                Some(a) => to = a,
                None => {
                    usage();
                    exit(1);
                },
            }
        } else if arg == "-n" {
            match args.next() {
                Some(sols) => number_solutions = sols.parse().map_err(|_| {
                    eprintln!("Argument to '-n' must be a number");
                    exit(1);
                }).unwrap(),
                None => {
                    eprintln!("Argument to '-n' must be a number");
                    exit(1);
                }
            }
        } else if arg == "-v" {
            verbose = true;
        } else if arg == "-@" {
            match args.next() {
                Some(tstamp) => timestamp = parse_timestring(&tstamp).unwrap_or_else(|| {
                    eprintln!("Argument to '-@' must be a timestring");
                    exit(1);
                }),
                None => {
                    eprintln!("Argument to '-@' must be a timestring");
                    exit(1);
                }
            }
        }
    }
    if from == "" {
        eprintln!("Departing station not found");
        usage();
        exit(1);
    }
    if to == "" {
        eprintln!("Arriving station not found");
        usage();
        exit(1);
    }
    let agent = ureq::AgentBuilder::new().build();
    match search_trains(&agent, &stations::STATIONS, &from, &to, &timestamp) {
	Some(tickets) => {
	    for ticket in tickets.iter().take(number_solutions) {
		print_ticket(ticket, &timestamp, verbose);
	    }
	},
	None => exit(1),
    };
    Ok(())
}

/*
 * Copyright ©️ 20224 Mario Forzanini <mf@marioforzanini.com>
 *
 * This file is part of trenord.
 *
 * This file 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.
 *
 * This file 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 this file. If not, see <https://www.gnu.org/licenses/>.
 *
 */