// SPDX-FileCopyrightText: 2025 Jonah Brüchert <jbb@kaidan.im
//
// SPDX-License-Identifier: MIT

use std::{collections::BTreeMap, error::Error, fmt::Display, fs::File, path::Path};

use csv::Writer;
use gtfs_structures::{
    Agency, Calendar, CalendarDate, FeedInfo, RawFrequency, RawGtfs, RawStopTime, RawTranslation,
    RawTrip, Route, Shape, Stop, TicketingDeepLink, TicketingIdentifier,
};
use zip::{write::SimpleFileOptions, write::ZipWriter};

#[cfg(feature = "osm")]
/// Features for automatically integrating data from OpenStreetMap
pub mod osm;

/// Error type for GTFS writing
#[derive(Debug)]
pub enum GtfsGeneratorError {
    /// An entry did not have an ID. This crate currently requires agency IDs to be set.
    NoId,
    /// The entry referred to by a foreign key could not be found.
    InvalidReference(String),
    /// An error occured while serializing CSV
    Csv(csv::Error),
    /// An error occured while writing the ZIP file
    Zip(zip::result::ZipError),
    /// A general IO error occured
    Io(std::io::Error),
    // Error in the data read from a preexisting GTFS-Feed
    ExistingData(gtfs_structures::Error),
}

impl Display for GtfsGeneratorError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NoId => write!(f, "The agency id can not be empty"),
            Self::InvalidReference(name) => write!(f, "The referenced {name:?} does not exist"),
            Self::Csv(e) => write!(f, "CSV error: {e}"),
            Self::Zip(e) => write!(f, "Zip error: {e}"),
            Self::Io(e) => write!(f, "IO error: {e}"),
            Self::ExistingData(e) => write!(f, "Existing data: {e}"),
        }
    }
}

impl Error for GtfsGeneratorError {}

impl From<csv::Error> for GtfsGeneratorError {
    fn from(value: csv::Error) -> Self {
        Self::Csv(value)
    }
}

impl From<zip::result::ZipError> for GtfsGeneratorError {
    fn from(value: zip::result::ZipError) -> Self {
        Self::Zip(value)
    }
}

impl From<std::io::Error> for GtfsGeneratorError {
    fn from(value: std::io::Error) -> Self {
        Self::Io(value)
    }
}

impl From<gtfs_structures::Error> for GtfsGeneratorError {
    fn from(value: gtfs_structures::Error) -> Self {
        Self::ExistingData(value)
    }
}

/// Result type used for all operations from this crate
pub type Result<T> = std::result::Result<T, GtfsGeneratorError>;

/// Incomplete GTFS feed
pub struct GtfsGenerator {
    pub agencies: BTreeMap<String, Agency>,
    pub stops: BTreeMap<String, Stop>,
    pub routes: BTreeMap<String, Route>,
    pub trips: BTreeMap<String, RawTrip>,
    pub stop_times: BTreeMap<(String, u32), RawStopTime>,
    pub shapes: BTreeMap<(String, usize), Shape>,
    pub services: BTreeMap<String, Calendar>,
    pub exceptions: BTreeMap<String, BTreeMap<chrono::NaiveDate, CalendarDate>>,
    pub frequencies: BTreeMap<(String, u32), RawFrequency>,
    pub ticketing_deep_links: BTreeMap<String, TicketingDeepLink>,
    pub ticketing_identifiers: BTreeMap<(String, String), TicketingIdentifier>,
    pub translations: BTreeMap<
        (
            String,
            String,
            String,
            Option<String>,
            Option<String>,
            String,
        ),
        RawTranslation,
    >,
    pub feed_info: Option<FeedInfo>,
}

impl Default for GtfsGenerator {
    fn default() -> Self {
        Self::new()
    }
}

impl GtfsGenerator {
    /// Create a new GTFS generator
    pub fn new() -> GtfsGenerator {
        GtfsGenerator {
            agencies: BTreeMap::new(),
            stops: BTreeMap::new(),
            routes: BTreeMap::new(),
            trips: BTreeMap::new(),
            stop_times: BTreeMap::new(),
            shapes: BTreeMap::new(),
            services: BTreeMap::new(),
            exceptions: BTreeMap::new(),
            frequencies: BTreeMap::new(),
            ticketing_deep_links: BTreeMap::new(),
            ticketing_identifiers: BTreeMap::new(),
            feed_info: None,
            translations: BTreeMap::new(),
        }
    }

    fn index_file<T, K: Ord, F: Fn(&T) -> K>(items: Vec<T>, key: F) -> Result<BTreeMap<K, T>> {
        let mut index = BTreeMap::new();

        for item in items {
            index.insert(key(&item), item);
        }

        return Ok(index);
    }

    fn get_optional_gtfs_subfile<T, K: Ord, F: Fn(&T) -> K>(
        result: Option<std::result::Result<Vec<T>, gtfs_structures::Error>>,
        key: F,
    ) -> Result<Option<BTreeMap<K, T>>> {
        match result {
            Some(result) => Ok(Some(Self::index_file(
                result.map_err(GtfsGeneratorError::ExistingData)?,
                key,
            )?)),
            None => Ok(None),
        }
    }

    pub fn from_existing_feed(gtfs: RawGtfs) -> Result<GtfsGenerator> {
        let feed_info: Option<FeedInfo> = if let Some(feed_info_result) = gtfs.feed_info {
            feed_info_result
                .map_err(GtfsGeneratorError::ExistingData)?
                .first()
                .map(|info| info.clone())
        } else {
            None
        };

        Ok(GtfsGenerator {
            agencies: Self::index_file(gtfs.agencies?, |agency: &Agency| {
                agency.id.clone().unwrap_or("default_agency".to_string())
            })?,
            stops: Self::index_file(gtfs.stops?, |stop: &Stop| stop.id.clone())?,
            routes: Self::index_file(gtfs.routes?, |route: &Route| route.id.clone())?,
            trips: Self::index_file(gtfs.trips?, |trip: &RawTrip| trip.id.clone())?,
            stop_times: Self::index_file(gtfs.stop_times?, |st: &RawStopTime| {
                (st.trip_id.to_string(), st.stop_sequence)
            })?,
            shapes: Self::get_optional_gtfs_subfile(gtfs.shapes, |shape: &Shape| {
                (shape.id.clone(), shape.sequence)
            })?
            .unwrap_or_default(),
            services: Self::get_optional_gtfs_subfile(gtfs.calendar, |cal: &Calendar| {
                cal.id.clone()
            })?
            .unwrap_or_default(),
            exceptions: {
                let mut exceptions = BTreeMap::new();

                if let Some(dates) = gtfs.calendar_dates {
                    let dates = dates?;

                    for date in dates {
                        exceptions
                            .entry(date.service_id.clone())
                            .or_insert(BTreeMap::new())
                            .insert(date.date, date);
                    }
                }

                exceptions
            },
            feed_info,
            frequencies: Self::get_optional_gtfs_subfile(gtfs.frequencies, |f: &RawFrequency| {
                (f.trip_id.clone(), f.start_time)
            })?
            .unwrap_or_default(),
            translations: Self::get_optional_gtfs_subfile(gtfs.translations, |translation| {
                (
                    translation.table_name.clone(),
                    translation.field_name.clone(),
                    translation.language.clone(),
                    translation.record_id.clone(),
                    translation.record_sub_id.clone(),
                    translation.field_name.clone(),
                )
            })?
            .unwrap_or_default(),
            ticketing_deep_links: Self::get_optional_gtfs_subfile(
                gtfs.ticketing_deep_links,
                |deep_link| deep_link.ticketing_deep_link_id.clone(),
            )?
            .unwrap_or_default(),
            ticketing_identifiers: Self::get_optional_gtfs_subfile(
                gtfs.ticketing_identifiers,
                |ticketing_identifier| {
                    (
                        ticketing_identifier.agency_id.clone(),
                        ticketing_identifier.stop_id.clone(),
                    )
                },
            )?
            .unwrap_or_default(),
        })
    }

    /// Add an agency
    ///
    /// See <https://gtfs.org/schedule/reference/#agencytxt> for rules of the individual fields
    ///
    /// Note: This crate requires agency IDs to be set, even if there's only a single agency.
    pub fn add_agency(&mut self, agency: Agency) -> Result<()> {
        self.agencies.insert(
            agency.id.as_ref().ok_or(GtfsGeneratorError::NoId)?.clone(),
            agency,
        );
        Ok(())
    }

    /// Add a stop
    ///
    /// See <https://gtfs.org/schedule/reference/#stopstxt> for rules of the individual fields
    pub fn add_stop(&mut self, stop: Stop) -> Result<()> {
        self.stops.insert(stop.id.clone(), stop);
        Ok(())
    }

    /// Add a route
    ///
    /// See <https://gtfs.org/schedule/reference/#routestxt> for rules of the individual fields
    ///
    /// Note: This crate requires agency IDs to be set, even if there's only a single agency.
    pub fn add_route(&mut self, route: Route) -> Result<()> {
        if !self
            .agencies
            .contains_key(route.agency_id.as_ref().ok_or(GtfsGeneratorError::NoId)?)
        {
            return Err(GtfsGeneratorError::InvalidReference(
                "agency_id".to_string(),
            ));
        }
        self.routes.insert(route.id.clone(), route);
        Ok(())
    }

    /// Add a trip
    ///
    /// See <https://gtfs.org/schedule/reference/#tripstxt> for rules of the individual fields
    ///
    /// Note: All entries that the trip refers to already need to exist
    pub fn add_trip(&mut self, trip: RawTrip) -> Result<()> {
        if !self.routes.contains_key(&trip.route_id) {
            return Err(GtfsGeneratorError::InvalidReference("route_id".to_string()));
        }
        if !self.services.contains_key(&trip.service_id)
            && !self.exceptions.contains_key(&trip.service_id)
        {
            return Err(GtfsGeneratorError::InvalidReference(
                "service_id".to_string(),
            ));
        }
        self.trips.insert(trip.id.clone(), trip);
        Ok(())
    }

    /// Add a stop time
    ///
    /// See <https://gtfs.org/schedule/reference/#stoptimestxt> for rules of the individual fields
    ///
    /// Note: All entries that the trip refers to already need to exist
    pub fn add_stop_time(&mut self, stop_time: RawStopTime) -> Result<()> {
        if !self.trips.contains_key(&stop_time.trip_id) {
            return Err(GtfsGeneratorError::InvalidReference("trip_id".to_string()));
        }
        self.stop_times.insert(
            (stop_time.trip_id.clone(), stop_time.stop_sequence),
            stop_time,
        );
        Ok(())
    }

    /// Add a shape point
    ///
    /// See <https://gtfs.org/schedule/reference/#shapestxt> for rules of the individual fields
    pub fn add_shape(&mut self, shape: Shape) -> Result<()> {
        self.shapes
            .insert((shape.id.clone(), shape.sequence), shape);
        Ok(())
    }

    /// Add a calendar entry
    ///
    /// See <https://gtfs.org/schedule/reference/#calendartxt> for rules of the individual fields
    pub fn add_service(&mut self, calendar: Calendar) -> Result<()> {
        self.services.insert(calendar.id.to_string(), calendar);
        Ok(())
    }

    /// Add calendar exception
    ///
    /// See <https://gtfs.org/schedule/reference/#calendardates> for rules of the individual fields
    pub fn add_service_exception(&mut self, calendar_date: CalendarDate) -> Result<()> {
        self.exceptions
            .entry(calendar_date.service_id.clone())
            .or_insert(BTreeMap::new())
            .insert(calendar_date.date, calendar_date);
        Ok(())
    }

    /// Add frequency entry
    ///
    /// See <https://gtfs.org/schedule/reference/#frequenciestxt> for rules of the individual fields
    ///
    /// Note: The referenced trip must already exist
    pub fn add_frequency(&mut self, frequency: RawFrequency) -> Result<()> {
        if !self.trips.contains_key(&frequency.trip_id) {
            return Err(GtfsGeneratorError::InvalidReference("trip_id".to_string()));
        }
        self.frequencies
            .insert((frequency.trip_id.clone(), frequency.start_time), frequency);
        Ok(())
    }

    /// Add ticketing deep link
    ///
    /// See <https://developers.google.com/transit/gtfs/reference/google-transit-ticketing-extension?hl=de#field-definitions> for rules of the individual fields
    pub fn add_ticketing_deep_link(&mut self, link: TicketingDeepLink) -> Result<()> {
        self.ticketing_deep_links
            .insert(link.ticketing_deep_link_id.clone(), link);
        Ok(())
    }

    /// Add ticketing deep link
    ///
    /// See <https://developers.google.com/transit/gtfs/reference/google-transit-ticketing-extension?hl=de#field-definitions> for rules of the individual fields
    pub fn add_ticketing_identifier(&mut self, link: TicketingIdentifier) -> Result<()> {
        self.ticketing_identifiers
            .insert((link.agency_id.to_string(), link.stop_id.to_string()), link);
        Ok(())
    }

    /// Add translation
    ///
    /// See <https://gtfs.org/schedule/reference/#translationstxt> for rules of the individual fields
    pub fn add_translation(&mut self, translation: RawTranslation) -> Result<()> {
        let translation_id = (
            translation.table_name.to_string(),
            translation.field_name.to_string(),
            translation.language.to_string(),
            translation.record_id.clone(),
            translation.record_sub_id.clone(),
            translation.field_name.to_string(),
        );
        self.translations.insert(translation_id, translation);
        Ok(())
    }

    /// Set the feed info metadata
    ///
    /// See <https://gtfs.org/schedule/reference/#feedinfotxt> for rules of the individual fields
    pub fn set_feed_info(&mut self, feed_info: FeedInfo) -> Result<()> {
        self.feed_info = Some(feed_info);
        Ok(())
    }

    /// Set the start and end date of the feed based on the calendar
    ///
    /// Only does something if feed info and calendar have been set before.
    fn set_feed_timeframe(&mut self) {
        if !self.services.is_empty() || !self.exceptions.is_empty() {
            let calendar_start_date = self
                .services
                .values()
                .min_by_key(|cal| cal.start_date)
                .map(|cal| cal.start_date);
            let calendar_end_date = self
                .services
                .values()
                .max_by_key(|cal| cal.end_date)
                .map(|cal| cal.end_date);

            let exceptions_start_date = self
                .exceptions
                .values()
                .flat_map(|v| v.values())
                .min_by_key(|exc| exc.date)
                .map(|exc| exc.date);

            let exceptions_end_date = self
                .exceptions
                .values()
                .flat_map(|v| v.values())
                .max_by_key(|exc| exc.date)
                .map(|exc| exc.date);

            let timeframe = if let (Some(cal_start), Some(cal_end)) =
                (calendar_start_date, calendar_end_date)
            {
                if let (Some(exc_start), Some(exc_end)) =
                    (exceptions_start_date, exceptions_end_date)
                {
                    let start = if exc_start < cal_start {
                        exc_start
                    } else {
                        cal_start
                    };

                    let end = if exc_end > cal_end { exc_end } else { cal_end };

                    Some((start, end))
                } else {
                    Some((cal_start, cal_end))
                }
            } else {
                if let (Some(exc_start), Some(exc_end)) =
                    (exceptions_start_date, exceptions_end_date)
                {
                    Some((exc_start, exc_end))
                } else {
                    None
                }
            };

            if let Some(ref mut info) = self.feed_info {
                if let Some((start_date, end_date)) = timeframe {
                    info.start_date = Some(start_date);
                    info.end_date = Some(end_date);
                }
            }
        }
    }

    fn write_file<T: serde::Serialize, U: Iterator<Item = T>>(
        zip: &mut ZipWriter<File>,
        name: &str,
        data: U,
    ) -> Result<()> {
        let options =
            SimpleFileOptions::default().compression_method(zip::CompressionMethod::Deflated);

        zip.start_file(name, options)?;
        {
            let mut writer = Writer::from_writer(zip);
            for item in data {
                writer.serialize(item)?;
            }
        }
        Ok(())
    }

    /// Write the resulting feed to a ZIP file
    ///
    /// Before writing, all fields that can be set automatically are set. Currently that is:
    /// - feed_start_date (Set to the first service day in the calendar)
    /// - feed_end_date (Set to the last service day in the calendar)
    ///
    /// If the assumptions to not match for your purpose, for example because you include days
    /// in the calendar for which the schedule is not yet certain, please set the fields to
    /// the correct value before writing.
    pub fn write_to<T: AsRef<Path>>(mut self, path: T) -> Result<()> {
        self.set_feed_timeframe();

        let mut zip = ZipWriter::new(File::create(path)?);

        Self::write_file(&mut zip, "agency.txt", self.agencies.values())?;
        Self::write_file(&mut zip, "stops.txt", self.stops.values())?;
        Self::write_file(&mut zip, "routes.txt", self.routes.values())?;
        Self::write_file(&mut zip, "trips.txt", self.trips.values())?;
        Self::write_file(&mut zip, "stop_times.txt", self.stop_times.values())?;

        if !self.shapes.is_empty() {
            Self::write_file(&mut zip, "shapes.txt", self.shapes.values())?;
        }

        if !self.services.is_empty() {
            Self::write_file(&mut zip, "calendar.txt", self.services.values())?;
        }

        if !self.exceptions.is_empty() {
            Self::write_file(
                &mut zip,
                "calendar_dates.txt",
                self.exceptions.values().flat_map(|v| v.values()),
            )?;
        }

        if !self.frequencies.is_empty() {
            Self::write_file(&mut zip, "frequencies.txt", self.frequencies.values())?;
        }

        if !self.translations.is_empty() {
            Self::write_file(&mut zip, "translations.txt", self.translations.values())?;
        }

        if !self.ticketing_deep_links.is_empty() {
            Self::write_file(
                &mut zip,
                "ticketing_deep_links.txt",
                self.ticketing_deep_links.values(),
            )?;
        }

        if !self.ticketing_identifiers.is_empty() {
            Self::write_file(
                &mut zip,
                "ticketing_identifiers.txt",
                self.ticketing_identifiers.values(),
            )?;
        }

        if let Some(feed_info) = &self.feed_info {
            Self::write_file(&mut zip, "feed_info.txt", [feed_info].iter())?;
        }

        Ok(())
    }
}
