// Copyright 2023 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

use std::{
    collections::{HashMap, HashSet},
    fmt,
    sync::Arc,
};

use petgraph::{
    dot::{self, Dot},
    graph::{Neighbors, NodeIndex},
    Directed, Direction, Graph,
};

use crate::target::{Provider, Target};
use crate::{Error, Result};

/// A node in the dependency graph.
#[derive(Debug)]
pub enum Dependency {
    /// A resolved dependency.
    Resolved(Arc<Target>),

    /// An unresolved dependency.  This occurs when a target declares a
    /// dependency for a target that has not been seen.
    Unresolved(String),
}

impl fmt::Display for Dependency {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Resolved(target) => write!(f, "{}", target.full_name()),
            Self::Unresolved(name) => write!(f, "{name}"),
        }
    }
}

/// An edge in the dependency graph.
/// We define an empty struct for edges instead of using the unit type
/// to allow us to implement `fmt::Display` for it.  This is needed for
/// graphviz output.
#[derive(Debug)]
pub struct Edge {}

impl fmt::Display for Edge {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "")
    }
}

/// A database of packages known to `qg`.
#[derive(Debug)]
pub struct Registry {
    /// Mapping of package slug to package. Slugs are globally unique.
    targets: HashMap<String, Arc<Target>>,

    /// All known package providers by name.
    providers: HashMap<String, Provider>,

    /// Unknown target dependencies.
    unresolved_dependencies: HashSet<String>,

    /// Dependency Graph of targets.
    ///
    /// Nodes that have a `None` weight are unresolved.
    dependency_graph: Graph<Dependency, Edge, Directed, u32>,

    /// Map from target full name to dependency graph node IDe.
    node_id_map: HashMap<String, NodeIndex<u32>>,
}

impl Registry {
    /// Initializes an empty package registry.
    #[must_use]
    pub fn new() -> Self {
        Self {
            targets: HashMap::new(),
            providers: HashMap::new(),
            unresolved_dependencies: HashSet::new(),
            dependency_graph: Graph::new(),
            node_id_map: HashMap::new(),
        }
    }

    /// Returns the number of registered targets.
    #[must_use]
    pub fn target_count(&self) -> usize {
        self.targets.len()
    }

    /// Returns an iterator over all known targets, in arbitrary order.
    pub fn targets(&self) -> impl Iterator<Item = &Target> {
        self.targets.values().map(Arc::as_ref)
    }

    /// Registers a new provider.
    pub(crate) fn add_provider(&mut self, provider: Provider) -> Result<()> {
        if self.providers.contains_key(&provider.name) {
            // TODO(frolv): Provider names must be unique.
            return Err(Error::GenericErrorPlaceholder);
        }
        self.providers.insert(provider.name.clone(), provider);
        Ok(())
    }

    /// Returns the dependency graph in graphviz dot format.
    #[must_use]
    pub fn dot_dependency_graph(&self) -> String {
        format!(
            "{:}",
            Dot::with_config(&self.dependency_graph, &[dot::Config::EdgeNoLabel])
        )
    }

    /// Inserts a target into the registry.
    pub(crate) fn add_target(&mut self, target: Target) -> Result<()> {
        if !self.providers.contains_key(target.provider()) {
            // TODO(frolv): Target must have a known provider.
            return Err(Error::GenericErrorPlaceholder);
        };

        let target_name = target.full_name();
        for dep in target
            .dependencies()
            .filter(|&d| !self.targets.contains_key(d))
        {
            self.unresolved_dependencies.insert(dep.into());
        }

        self.unresolved_dependencies.remove(&target_name);

        let target = Arc::new(target);
        self.update_dependency_graph_for_new_target(&target);
        self.targets.insert(target_name.clone(), target);

        Ok(())
    }

    fn update_dependency_graph_for_new_target(&mut self, target: &Arc<Target>) {
        let target_name = target.full_name();

        let node_id = if let Some(id) = self.node_id_map.get(&target_name) {
            // If the target is already in the graph, it is because a previous
            // target as taken an unresolved dependency on it.  Update the node
            // with the resolved target.
            *self
                .dependency_graph
                .node_weight_mut(*id)
                .expect("node should exist in graph") = Dependency::Resolved(target.clone());
            *id
        } else {
            let id = self
                .dependency_graph
                .add_node(Dependency::Resolved(target.clone()));
            self.node_id_map.insert(target_name, id);
            id
        };

        for dep in target.dependencies() {
            // Get the node id of the dependency.  Insert an unresolved
            // dependency if the node does not already exist.
            let dep_node_id = *self.node_id_map.entry(dep.to_string()).or_insert_with(|| {
                self.dependency_graph
                    .add_node(Dependency::Unresolved(dep.to_string()))
            });
            self.dependency_graph
                .add_edge(node_id, dep_node_id, Edge {});
        }
    }

    /// Returns the node ID for a node.
    ///
    /// # Errors
    /// Returns an error if the given node name does not exist.
    pub fn get_node_id(&self, name: &str) -> Result<NodeIndex<u32>> {
        self.node_id_map
            .get(name)
            .ok_or(
                // Node not found.
                Error::GenericErrorPlaceholder,
            )
            .map(|id| *id)
    }

    /// Returns an iterator of dependencies for the node.
    #[must_use]
    pub fn node_deps(&self, node: NodeIndex<u32>) -> Neighbors<Edge, u32> {
        self.dependency_graph
            .neighbors_directed(node, Direction::Outgoing)
    }
}

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