// Copyright 2020 The Jujutsu 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.

#![expect(missing_docs)]

use std::collections::BTreeMap;
use std::sync::Arc;

use futures::TryStreamExt as _;
use futures::future::BoxFuture;
use futures::stream::FuturesUnordered;

use crate::backend;
use crate::backend::BackendResult;
use crate::backend::TreeId;
use crate::backend::TreeValue;
use crate::repo_path::RepoPathBuf;
use crate::repo_path::RepoPathComponentBuf;
use crate::repo_path::RepoPathTree;
use crate::store::Store;
use crate::tree::Tree;

#[derive(Debug)]
enum Override {
    Tombstone,
    Replace(TreeValue),
}

#[derive(Debug)]
pub struct TreeBuilder {
    store: Arc<Store>,
    base_tree_id: TreeId,
    overrides: BTreeMap<RepoPathBuf, Override>,
}

impl TreeBuilder {
    pub fn new(store: Arc<Store>, base_tree_id: TreeId) -> Self {
        let overrides = BTreeMap::new();
        Self {
            store,
            base_tree_id,
            overrides,
        }
    }

    pub fn store(&self) -> &Store {
        self.store.as_ref()
    }

    pub fn set(&mut self, path: RepoPathBuf, value: TreeValue) {
        assert!(!path.is_root());
        self.overrides.insert(path, Override::Replace(value));
    }

    pub fn remove(&mut self, path: RepoPathBuf) {
        assert!(!path.is_root());
        self.overrides.insert(path, Override::Tombstone);
    }

    pub fn set_or_remove(&mut self, path: RepoPathBuf, value: Option<TreeValue>) {
        assert!(!path.is_root());
        if let Some(value) = value {
            self.overrides.insert(path, Override::Replace(value));
        } else {
            self.overrides.insert(path, Override::Tombstone);
        }
    }

    pub async fn write_tree(self) -> BackendResult<TreeId> {
        if self.overrides.is_empty() {
            return Ok(self.base_tree_id);
        }

        let mut trees_to_write = self.get_base_trees().await?;

        // Update entries in parent trees for file overrides
        for (path, file_override) in self.overrides {
            let (dir, basename) = path.split().unwrap();
            let tree_entries = trees_to_write.get_mut(dir).unwrap();
            match file_override {
                Override::Replace(value) => {
                    tree_entries.insert(basename.to_owned(), value);
                }
                Override::Tombstone => {
                    tree_entries.remove(basename);
                }
            }
        }

        // Write trees in reverse lexicographical order, starting with trees without
        // children.
        // TODO: Writing trees concurrently should help on high-latency backends
        let store = &self.store;
        while let Some((dir, cur_entries)) = trees_to_write.pop_last() {
            if let Some((parent, basename)) = dir.split() {
                let parent_entries = trees_to_write.get_mut(parent).unwrap();
                if cur_entries.is_empty() {
                    if let Some(TreeValue::Tree(_)) = parent_entries.get(basename) {
                        parent_entries.remove(basename);
                    } else {
                        // Entry would have been replaced with file (see above)
                    }
                } else {
                    let data =
                        backend::Tree::from_sorted_entries(cur_entries.into_iter().collect());
                    let tree = store.write_tree(&dir, data).await?;
                    parent_entries.insert(basename.to_owned(), TreeValue::Tree(tree.id().clone()));
                }
            } else {
                // We're writing the root tree. Write it even if empty. Return its id.
                assert!(trees_to_write.is_empty());
                let data = backend::Tree::from_sorted_entries(cur_entries.into_iter().collect());
                let written_tree = store.write_tree(&dir, data).await?;
                return Ok(written_tree.id().clone());
            }
        }

        unreachable!("trees_to_write must contain the root tree");
    }

    async fn get_base_trees(
        &self,
    ) -> BackendResult<BTreeMap<RepoPathBuf, BTreeMap<RepoPathComponentBuf, TreeValue>>> {
        // All base trees we need
        let mut needed_dirs: RepoPathTree<()> = RepoPathTree::default();
        for path in self.overrides.keys() {
            if let Some(dir) = path.parent() {
                needed_dirs.add(dir);
            }
        }

        let mut tree_reads: FuturesUnordered<BoxFuture<'_, BackendResult<(RepoPathBuf, Tree)>>> =
            FuturesUnordered::new();

        // Schedule reading the root tree
        tree_reads.push(Box::pin(async move {
            let root_dir = RepoPathBuf::root();
            let tree = self
                .store
                .get_tree(root_dir.clone(), &self.base_tree_id)
                .await?;
            Ok((root_dir, tree))
        }));

        let mut base_trees = BTreeMap::new();
        while let Some((dir, tree)) = tree_reads.try_next().await? {
            if let Some(node) = needed_dirs.get(&dir) {
                for (basename, _child) in node.children() {
                    let basename = basename.to_owned();
                    let sub_dir = dir.join(&basename);
                    let tree = tree.clone();
                    tree_reads.push(Box::pin(async move {
                        let sub_tree = tree
                            .sub_tree(&basename)
                            .await?
                            .unwrap_or_else(|| Tree::empty(self.store.clone(), sub_dir.clone()));
                        Ok((sub_dir, sub_tree))
                    }));
                }
            }
            base_trees.insert(dir, tree);
        }

        Ok(base_trees
            .into_iter()
            .map(|(dir, tree)| {
                let entries = tree
                    .data()
                    .entries()
                    .map(|entry| (entry.name().to_owned(), entry.value().clone()))
                    .collect();
                (dir, entries)
            })
            .collect())
    }
}