use crate::{
    algorithm::search::{Direction, EdgeTraversal, SearchError, SearchTree},
    model::{
        constraint::ConstraintModel,
        cost::CostModel,
        label::{label_model::LabelModel, Label},
        map::MapModel,
        network::{EdgeId, EdgeListId, Graph},
        state::StateModel,
        termination::TerminationModel,
        traversal::{EdgeFrontierContext, TraversalModel},
    },
};
use std::sync::Arc;

/// A `SearchInstance` represents the collection of read-only models and data required
/// to execute a search query. It encapsulates the graph, constraints, traversal logic,
/// cost calculations, and termination criteria, providing a unified interface for
/// search algorithms to interact with the underlying network and models.
pub struct SearchInstance {
    pub graph: Arc<Graph>,
    pub constraint_models: Vec<Arc<dyn ConstraintModel>>,
    pub traversal_models: Vec<Arc<dyn TraversalModel>>,
    pub map_model: Arc<MapModel>,
    pub state_model: Arc<StateModel>,
    pub cost_model: Arc<CostModel>,
    pub termination_model: Arc<TerminationModel>,
    pub label_model: Arc<dyn LabelModel>,
    pub default_edge_list: Option<usize>,
}

impl SearchInstance {
    /// Retrieves the traversal model that should be used for traversal estimation
    /// when no specific edge is being traversed (e.g., at the start of a search or
    /// when calculating heuristics). It falls back to the default edge list's model.
    pub fn get_traversal_estimation_model(&self) -> Arc<dyn TraversalModel> {
        self.traversal_models[self.default_edge_list.unwrap_or_default()].clone()
    }

    /// Retrieves the constraint model associated with a specific edge list.
    ///
    /// # Arguments
    ///
    /// * `edge_list_id` - The ID of the edge list to get the constraint model for.
    ///
    /// # Returns
    ///
    /// A reference to the constraint model, or an error if the ID is not found.
    pub fn get_constraint_model(
        &self,
        edge_list_id: &EdgeListId,
    ) -> Result<Arc<dyn ConstraintModel>, SearchError> {
        self.constraint_models
            .get(edge_list_id.0)
            .ok_or_else(|| SearchError::InternalError(format!("during search, attempting to retrieve constraint models for edge list {edge_list_id} that does not exist")))
            .cloned()
    }

    /// Retrieves the traversal model associated with a specific edge list.
    ///
    /// # Arguments
    ///
    /// * `edge_list_id` - The ID of the edge list to get the traversal model for.
    ///
    /// # Returns
    ///
    /// A reference to the traversal model, or an error if the ID is not found.
    pub fn get_traversal_model(
        &self,
        edge_list_id: &EdgeListId,
    ) -> Result<Arc<dyn TraversalModel>, SearchError> {
        self.traversal_models
            .get(edge_list_id.0)
            .ok_or_else(|| SearchError::InternalError(format!("during search, attempting to retrieve traversal models for edge list {edge_list_id} that does not exist")))
            .cloned()
    }

    /// Computes the sequence of `EdgeTraversal` objects for a given path of edge IDs.
    ///
    /// This method is essential for reconstructing the full state (costs, state transitions)
    /// along a path that was generated by an external process, such as map matching,
    /// which typically only returns a sequence of edge identifiers.
    ///
    /// # Implementation Note: Vector Index Tracking
    ///
    /// When building the internal `SearchTree` for the path, this method uses "indexed labels"
    /// (`Label::VertexWithIntState`) where the state is the current index in the path.
    ///
    /// We must track the vector index of each insertion for the following reasons:
    ///
    /// 1. **Path Cycles**: A path may visit the same vertex multiple times (e.g., in complex
    ///    intersections or specific routing constraints). If we used plain `Label::Vertex`,
    ///    subsequent visits to the same vertex would overwrite previous entries in the
    ///    `SearchTree`'s internal map, breaking the tree structure.
    /// 2. **Backtracking**: The `SearchTree` depends on unique labels to correctly backtrack
    ///    from a destination to the root. By indexing each step, we ensure that every node
    ///    in the path is a unique entity in the tree, even if they share the same physical vertex.
    /// 3. **State Consistency**: It ensures that the state at each step of the path is
    ///    linearly associated with its predecessor, avoiding any ambiguity that could arise
    ///    if multiple paths through the same vertex were present in the tree.
    pub fn compute_path(
        &self,
        path: &[(EdgeListId, EdgeId)],
    ) -> Result<Vec<EdgeTraversal>, SearchError> {
        let mut edge_traversals = Vec::with_capacity(path.len());
        let mut current_state = self.state_model.initial_state(None)?;
        let mut tree = SearchTree::new_stateful(Direction::Forward);

        let mut prev_label = if let Some((edge_list_id, edge_id)) = path.first() {
            let (src, _, _) = self.graph.edge_triplet(edge_list_id, edge_id)?;
            let root_label = Label::VertexWithIntState {
                vertex_id: src.vertex_id,
                state: 0,
            };
            tree.set_root(root_label.clone())?;
            root_label
        } else {
            return Ok(Vec::new());
        };

        for (i, (edge_list_id, edge_id)) in path.iter().enumerate() {
            let (src, edge, dst) = self.graph.edge_triplet(edge_list_id, edge_id)?;
            let tm = self.get_traversal_model(edge_list_id)?;
            let ctx = EdgeFrontierContext::new(&prev_label, src, edge, dst, &tree);
            let traversal = EdgeTraversal::new_local(
                &ctx,
                &current_state,
                &self.state_model,
                tm.as_ref(),
                &self.cost_model,
            )?;

            // Use indexed labels to handle cycles in the path correctly. i+1 is used as the root will have state 0.
            let child_label = Label::VertexWithIntState {
                vertex_id: dst.vertex_id,
                state: i + 1,
            };
            tree.insert_trajectory(prev_label, traversal.clone(), child_label.clone())
                .map_err(|e| SearchError::InternalError(e.to_string()))?;

            prev_label = child_label;
            current_state = traversal.result_state.clone();
            edge_traversals.push(traversal);
        }

        Ok(edge_traversals)
    }
}
