use std::cell::RefCell;
use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use thread_local::ThreadLocal;

use abstutil::{Timer, VecMap};
use geom::Duration;

use crate::pathfind::engine::CreateEngine;
use crate::pathfind::vehicles::VehiclePathfinder;
use crate::pathfind::walking::SidewalkPathfinder;
use crate::{
    BusRouteID, BusStopID, DirectedRoadID, Map, PathConstraints, PathRequest, PathV2, Position,
    RoutingParams,
};

#[derive(Serialize, Deserialize)]
pub struct Pathfinder {
    car_graph: VehiclePathfinder,
    bike_graph: VehiclePathfinder,
    bus_graph: VehiclePathfinder,
    train_graph: VehiclePathfinder,
    walking_graph: SidewalkPathfinder,
    walking_with_transit_graph: SidewalkPathfinder,

    // These params cover the main graphs
    params: RoutingParams,

    // Callers can opt into caching with pathfind_with_params
    // TODO VecMap is probably fast enough. RoutingParams is annoying to implement Hash.
    #[serde(skip_serializing, skip_deserializing)]
    cached_alternatives: ThreadLocal<RefCell<VecMap<(PathConstraints, RoutingParams), Pathfinder>>>,
}

// Implemented manually to deal with the ThreadLocal
impl Clone for Pathfinder {
    fn clone(&self) -> Self {
        Self {
            car_graph: self.car_graph.clone(),
            bike_graph: self.bike_graph.clone(),
            bus_graph: self.bus_graph.clone(),
            train_graph: self.train_graph.clone(),
            walking_graph: self.walking_graph.clone(),
            walking_with_transit_graph: self.walking_with_transit_graph.clone(),
            params: self.params.clone(),
            cached_alternatives: ThreadLocal::new(),
        }
    }
}

impl Pathfinder {
    /// Quickly create an invalid pathfinder, just to make borrow checking / initialization order
    /// work.
    pub fn empty() -> Pathfinder {
        Pathfinder {
            car_graph: VehiclePathfinder::empty(),
            bike_graph: VehiclePathfinder::empty(),
            bus_graph: VehiclePathfinder::empty(),
            train_graph: VehiclePathfinder::empty(),
            walking_graph: SidewalkPathfinder::empty(),
            walking_with_transit_graph: SidewalkPathfinder::empty(),
            params: RoutingParams::default(),
            cached_alternatives: ThreadLocal::new(),
        }
    }

    pub fn new(
        map: &Map,
        params: RoutingParams,
        engine: CreateEngine,
        timer: &mut Timer,
    ) -> Pathfinder {
        timer.start("prepare pathfinding for cars");
        let car_graph = VehiclePathfinder::new(map, PathConstraints::Car, &params, &engine);
        timer.stop("prepare pathfinding for cars");

        // The edge weights for bikes are so different from the driving graph that reusing the node
        // ordering actually hurts!
        timer.start("prepare pathfinding for bikes");
        let bike_graph = VehiclePathfinder::new(map, PathConstraints::Bike, &params, &engine);
        timer.stop("prepare pathfinding for bikes");

        timer.start("prepare pathfinding for buses");
        let bus_graph = VehiclePathfinder::new(
            map,
            PathConstraints::Bus,
            &params,
            &car_graph.engine.reuse_ordering(),
        );
        timer.stop("prepare pathfinding for buses");

        // Light rail networks are absolutely tiny; using a contraction hierarchy for them is
        // overkill. And in fact, it costs a bit of memory and file size, so don't do it!
        timer.start("prepare pathfinding for trains");
        let train_graph = VehiclePathfinder::new(
            map,
            PathConstraints::Train,
            &params,
            &CreateEngine::Dijkstra,
        );
        timer.stop("prepare pathfinding for trains");

        timer.start("prepare pathfinding for pedestrians");
        let walking_graph = SidewalkPathfinder::new(map, None, &engine);
        timer.stop("prepare pathfinding for pedestrians");

        timer.start("prepare pathfinding for pedestrians using transit");
        let walking_with_transit_graph =
            SidewalkPathfinder::new(map, Some((&bus_graph, &train_graph)), &engine);
        timer.stop("prepare pathfinding for pedestrians using transit");

        Pathfinder {
            car_graph,
            bike_graph,
            bus_graph,
            train_graph,
            walking_graph,
            walking_with_transit_graph,

            params,
            cached_alternatives: ThreadLocal::new(),
        }
    }

    /// Create a new Pathfinder with custom routing params that can only serve some modes.
    pub fn new_limited(
        map: &Map,
        params: RoutingParams,
        engine: CreateEngine,
        modes: Vec<PathConstraints>,
        timer: &mut Timer,
    ) -> Pathfinder {
        let mut p = Pathfinder::empty();
        for constraints in modes {
            timer.start(format!("prepare pathfinding for just {:?}", constraints));
            match constraints {
                PathConstraints::Pedestrian => {
                    p.walking_graph = SidewalkPathfinder::new(map, None, &engine);
                }
                PathConstraints::Car => {
                    p.car_graph = VehiclePathfinder::new(map, constraints, &params, &engine);
                }
                PathConstraints::Bike => {
                    p.bike_graph = VehiclePathfinder::new(map, constraints, &params, &engine);
                }
                PathConstraints::Bus => {
                    p.bus_graph = VehiclePathfinder::new(map, constraints, &params, &engine);
                }
                PathConstraints::Train => {
                    p.train_graph = VehiclePathfinder::new(map, constraints, &params, &engine);
                }
            }
            timer.stop(format!("prepare pathfinding for just {:?}", constraints));
        }
        p.params = params;
        p
    }

    /// Finds a path from a start to an end for a certain type of agent.
    pub fn pathfind(&self, req: PathRequest, map: &Map) -> Option<PathV2> {
        self.pathfind_with_params(req, map.routing_params(), false, map)
    }

    /// Finds a path from a start to an end for a certain type of agent. May use custom routing
    /// parameters. If caching is requested and custom routing parameters are used, then the
    /// intermediate graph is saved to speed up future calls with the same routing parameters.
    pub fn pathfind_with_params(
        &self,
        req: PathRequest,
        params: &RoutingParams,
        cache_custom: bool,
        map: &Map,
    ) -> Option<PathV2> {
        let constraints = req.constraints;
        if params == &self.params {
            return match constraints {
                PathConstraints::Pedestrian => self.walking_graph.pathfind(req, map),
                PathConstraints::Car => self.car_graph.pathfind(req, map),
                PathConstraints::Bike => self.bike_graph.pathfind(req, map),
                PathConstraints::Bus => self.bus_graph.pathfind(req, map),
                PathConstraints::Train => self.train_graph.pathfind(req, map),
            };
        }

        // If the params differ from the ones baked into the map, the CHs won't match. Do we have a
        // cached alternative?
        if let Some(alt) = self
            .cached_alternatives
            .get_or(|| RefCell::new(VecMap::new()))
            .borrow()
            .get(&(constraints, params.clone()))
        {
            return alt.pathfind_with_params(req, params, false, map);
        }

        // If somebody's repeatedly calling this without caching, log very obnoxiously.
        let mut timer = Timer::new(format!("Pathfinding slowly for {} with custom params", req));
        let tmp_pathfinder = Pathfinder::new_limited(
            map,
            params.clone(),
            CreateEngine::Dijkstra,
            vec![constraints],
            &mut timer,
        );
        let result = tmp_pathfinder.pathfind_with_params(req, params, false, map);
        if cache_custom {
            self.cached_alternatives
                .get_or(|| RefCell::new(VecMap::new()))
                .borrow_mut()
                .push((constraints, params.clone()), tmp_pathfinder);
        }
        result
    }

    pub fn all_costs_from(
        &self,
        req: PathRequest,
        map: &Map,
    ) -> Option<(Duration, HashMap<DirectedRoadID, Duration>)> {
        let req_cost = self.pathfind(req.clone(), map)?.get_cost();
        let all_costs = match req.constraints {
            PathConstraints::Pedestrian => self.walking_graph.all_costs_from(req.start, map),
            PathConstraints::Car => self.car_graph.all_costs_from(req.start, map),
            PathConstraints::Bike => self.bike_graph.all_costs_from(req.start, map),
            PathConstraints::Bus | PathConstraints::Train => unreachable!(),
        };
        Some((req_cost, all_costs))
    }

    // TODO Consider returning the walking-only path in the failure case, to avoid wasting work
    pub fn should_use_transit(
        &self,
        map: &Map,
        start: Position,
        end: Position,
    ) -> Option<(BusStopID, Option<BusStopID>, BusRouteID)> {
        self.walking_with_transit_graph
            .should_use_transit(map, start, end)
    }

    pub fn apply_edits(&mut self, map: &Map, timer: &mut Timer) {
        timer.start("apply edits to car pathfinding");
        self.car_graph.apply_edits(map);
        timer.stop("apply edits to car pathfinding");

        timer.start("apply edits to bike pathfinding");
        self.bike_graph.apply_edits(map);
        timer.stop("apply edits to bike pathfinding");

        timer.start("apply edits to bus pathfinding");
        self.bus_graph.apply_edits(map);
        timer.stop("apply edits to bus pathfinding");

        timer.start("apply edits to train pathfinding");
        self.train_graph.apply_edits(map);
        timer.stop("apply edits to train pathfinding");

        timer.start("apply edits to pedestrian pathfinding");
        self.walking_graph.apply_edits(map, None);
        timer.stop("apply edits to pedestrian pathfinding");

        timer.start("apply edits to pedestrian using transit pathfinding");
        self.walking_with_transit_graph
            .apply_edits(map, Some((&self.bus_graph, &self.train_graph)));
        timer.stop("apply edits to pedestrian using transit pathfinding");
    }
}