// Copyright (C) 2025  The Software Heritage developers
// See the AUTHORS file at the top-level directory of this distribution
// License: GNU General Public License version 3, or any later version
// See top-level LICENSE file for more information

use std::num::NonZeroUsize;
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result};
use clap::Parser;
use mimalloc::MiMalloc;
use swh_contents::diff::differ::RevDiffer;
use swh_contents::retriever::rocksdb::{ReadonlyRocksdbContentCache, RocksdbContentCache};
use swh_contents::retriever::{CachedContentRetriever, ContentRetriever};
use swh_graph::graph::{SwhBidirectionalGraph, SwhGraphWithProperties, SwhLabeledForwardGraph};
use swh_graph::mph::DynMphf;
use swh_graph::properties;

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

#[derive(Parser, Debug)]
/// Given the path to a directory of `.ndjson` files written by the `szz`, computes the diffs of
/// all revisions referenced by it, and writes them to a new directory.
struct Args {
    #[arg(long)]
    /// Path to swh-graph's files
    graph: PathBuf,
    #[arg(long)]
    /// Number of tasks to run in parallel. Defaults to the number of CPUs * 10
    concurrency: Option<NonZeroUsize>,
    #[arg(long)]
    /// Path to swh-osv's sqlite database
    db: PathBuf,
    #[clap(flatten)]
    download_args: Option<swh_contents::retriever::cli::DownloadArgs>,
    #[arg(long)]
    content_cache: PathBuf,
    #[arg(long)]
    /// Path to the directory of CSV files written by V-SZZ's --similar-lines-out-dir
    similar_lines: Option<PathBuf>,
    #[arg(long)]
    /// Path to the directory of ndjson files written by SZZ
    szz_revs: PathBuf,
    #[arg(long)]
    /// Path to the directory where to write diffs
    out_dir: Option<PathBuf>,
}

#[tokio::main(flavor = "multi_thread")]
pub async fn main() -> Result<()> {
    let mut args = Args::parse();

    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive(tracing::Level::INFO.into()),
        )
        .init();

    tracing::info!("Loading graph...");
    let graph = SwhBidirectionalGraph::new(&args.graph)
        .with_context(|| format!("Could not load graph from {}", args.graph.display()))?
        .load_labels()
        .context("Could not load graph labels")?
        .init_properties()
        .load_properties(|properties| properties.load_maps::<DynMphf>())
        .context("Could not load maps")?
        .load_properties(|properties| properties.load_label_names())
        .context("Could not load label names")?;
    let graph = Arc::new(graph);

    match args.download_args.take() {
        Some(download_args) => {
            tracing::info!("Opening cache...");
            let cache = RocksdbContentCache::new(graph.clone(), &args.content_cache)
                .context("Could not load Rocksdb cache")?;
            let caching_retriever = CachedContentRetriever::new(
                cache,
                download_args
                    .build_retriever(Arc::clone(&graph))
                    .context("Could not initialize downloader")?,
            )
            .context("Could not join cache and downloader")?;
            main2(args, graph, caching_retriever).await
        }
        None => {
            tracing::info!("Opening cache...");
            let cache = ReadonlyRocksdbContentCache::new(graph.clone(), &args.content_cache)
                .context("Could not load Rocksdb cache")?;
            main2(args, graph, cache).await
        }
    }
}

async fn main2<G>(
    args: Args,
    graph: Arc<G>,
    content_retriever: impl ContentRetriever + Send,
) -> Result<()>
where
    G: SwhLabeledForwardGraph
        + SwhGraphWithProperties<Maps: properties::Maps, LabelNames: properties::LabelNames>
        + Sync
        + Send,
{
    let content_retriever = Arc::new(content_retriever);

    let revs = swh_osv::szz::list_revs_from_szz_output(
        &graph,
        &args.szz_revs,
        args.similar_lines.as_deref(),
    )
    .with_context(|| format!("Could not list revs in {}", args.szz_revs.display()))?;

    // filter out merges because we don't support computing diffs of merges yet
    let revs = swh_osv::git::filter_non_merge_revisions(&*graph, revs);

    let differ = RevDiffer::new(graph.clone(), content_retriever);

    let num_revs = revs.len();
    let mut batch_differ = swh_contents::diff::batch::BatchDiffer::new(graph, differ)
        .with_out_dir(args.out_dir.as_deref())
        .with_expected_num_revs(Some(num_revs));
    if let Some(concurrency) = args.concurrency {
        batch_differ = batch_differ.with_concurrency(concurrency);
    }

    batch_differ
        .diff_revisions(futures::stream::iter(revs))
        .await
}
