// 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, ValueEnum};
use dsi_progress_logger::{ProgressLog, concurrent_progress_logger};
use either::Either;
use mimalloc::MiMalloc;
use rusqlite::Connection;
use swh_contents::blame::blame::Blamer;
use swh_contents::diff::tokenizers::{Tokenizer, TrimmedAsciiTokenizer};
use swh_contents::retriever::rocksdb::{ReadonlyRocksdbContentCache, RocksdbContentCache};
use swh_contents::retriever::{CachedContentRetriever, ContentRetriever};
use swh_graph::graph::{
    SwhBidirectionalGraph, SwhGraphWithProperties, SwhLabeledBackwardGraph, SwhLabeledForwardGraph,
};
use swh_graph::mph::DynMphf;
use swh_graph::properties;

use swh_osv::szz::strategies::NaiveStrategyFactory;
use swh_osv::szz::tokenizers::StrippedAsciiWhitespaceTokenizer;
use swh_osv::{db, szz};

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

#[derive(ValueEnum, Debug, Clone)]
pub enum TokenizerArg {
    /// Line-level diff
    Line,
    /// Line-level diff which ignore changes to leading and trailing whitespace
    TrimmedLine,
    /// Line-level diff which ignore all ASCII whitespace changes
    WhitespaceStrippedLine,
}

#[derive(Parser, Debug)]
/// Runs the SZZ algorithm on all "fixed" commits and writes the results.
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)]
    /// Ignore fix revisions that touch more files than this
    max_tree_diff_items: Option<usize>,
    #[arg(long)]
    /// How many steps of "blame" to run on a single hunk before giving up on the hunk
    max_blame_steps_per_hunk: Option<u64>,
    #[arg(long, default_value = "line")]
    /// How to compute diffs between revisions
    tokenizer: TokenizerArg,
    #[arg(long)]
    /// How many hunks on a revision's diff to consider before giving up on the revision
    max_hunks: Option<u64>,
    #[arg(long)]
    /// How many modified files on a revision's diff to consider before giving up on the revision
    max_modified_files: Option<u64>,
    #[arg(long)]
    /// If given and < 1., applies fuzzy-matching to lines while looking for an introducing commit
    ///
    /// If a line would be attributed to a revision while blaming but there is a similar line
    /// (levenshtein distance < this_ratio * line_len) in a parent revision, that line is
    /// assumed to be an earlier version of it.
    ///
    /// See [the V-SZZ algorithm](https://baolingfeng.github.io/papers/ICSE2022VSZZ.pdf#page=4)
    /// for details:
    ///
    /// > We compute the line similarity based on the edit distance
    /// > and set the similarity threshold is equal to 0.75. But this might
    /// > result in some incorrect mappings, which would be discussed in
    /// > the evaluation (see Section 5). If the line mapping algorithm can
    /// > identify a corresponding line in the previous commit or descen-
    /// > dants commits, we continue going back by running the git blame
    /// > command for the line. This process stops until the line mapping
    /// > algorithm identifies the line as a new line, the source code file of
    /// > the line is firstly added, or the commit is the initial commit of the
    /// > repository.
    min_line_similarity_ratio: Option<f64>,
    #[arg(long, default_value_t = szz::DEFAULT_MAX_LINE_SIMILARITY_DISTANCE)]
    /// Upper bound on the edit distance of two similar lines.
    ///
    /// This complements `--min-line-similarity-ratio` for large lines.
    max_line_similarity_distance: u32,
    #[arg(long, default_value_t = szz::DEFAULT_MAX_SIMILAR_LINE_OFFSET)]
    /// Defines the size of the window of old lines where to look for similar lines to a new line
    ///
    /// For example, if `--max-similar-line-offset 1000` (the default), only the 1k old lines
    /// before and after a new line's position will be checked for similarity.
    ///
    /// Larger values linearly impact performance on large files.
    max_similar_line_offset: u32,
    #[arg(long)]
    middle_revisions_out_dir: Option<PathBuf>,
    #[arg(long)]
    /// Path to the directory to write as ndjson files with a line for each vulnerability
    vuln_out_dir: Option<PathBuf>,
    #[arg(long)]
    /// Path to the directory to write as ndjson files with a line for each vulnerable line
    line_details_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_strings())
        .context("Could not load strings")?
        .load_properties(|properties| properties.load_timestamps())
        .context("Could not load timestamps")?
        .load_properties(|properties| properties.load_label_names())
        .context("Could not load label names")?;
    let graph = Arc::new(graph);

    let content_retriever = 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")?;
            Either::Left(
                CachedContentRetriever::new(
                    cache,
                    download_args
                        .build_retriever(Arc::clone(&graph))
                        .context("Could not initialize downloader")?,
                )
                .context("Could not join cache and downloader")?,
            )
        }
        None => {
            tracing::info!("Opening cache...");
            Either::Right(
                ReadonlyRocksdbContentCache::new(graph.clone(), &args.content_cache)
                    .context("Could not load Rocksdb cache")?,
            )
        }
    };
    let content_retriever = Arc::new(content_retriever);

    let blamer = Blamer::new(graph.clone(), Arc::clone(&content_retriever));
    match args.tokenizer {
        TokenizerArg::Line => main2(args, graph, content_retriever, blamer).await,
        TokenizerArg::TrimmedLine => {
            let blamer = blamer.with_tokenizer(TrimmedAsciiTokenizer);
            main2(args, graph, content_retriever, blamer).await
        }
        TokenizerArg::WhitespaceStrippedLine => {
            let blamer = blamer.with_tokenizer(StrippedAsciiWhitespaceTokenizer);
            main2(args, graph, content_retriever, blamer).await
        }
    }
}

async fn main2<G, CR>(
    args: Args,
    graph: Arc<G>,
    content_retriever: Arc<CR>,
    blamer: Blamer<Arc<G>, Arc<CR>, impl Tokenizer + Clone + Sync + Send>,
) -> Result<()>
where
    G: SwhLabeledForwardGraph
        + SwhLabeledBackwardGraph
        + SwhGraphWithProperties<
            Maps: properties::Maps,
            LabelNames: properties::LabelNames,
            Strings: properties::Strings,
            Timestamps: properties::Timestamps,
        > + Sync
        + Send,
    CR: ContentRetriever + Sync + Send,
{
    tracing::debug!("Opening sqlite database...");
    let conn = Connection::open(&args.db)
        .with_context(|| format!("Could not open {}", args.db.display()))?;

    tracing::info!("Counting vulnerabilities...");
    let num_vulnerabilities =
        db::count_vulnerabilities(&conn).context("Could not count vulnerabilities")?;

    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "file",
        local_speed = true,
        expected_updates = Some(num_vulnerabilities),
    );
    pl.threshold(100);
    pl.start(format!(
        "Processing {num_vulnerabilities} vulnerabilities..."
    ));

    let vuln_dataset_writer = args
        .vuln_out_dir
        .map(|out_dir| {
            dataset_writer::ParallelDatasetWriter::<dataset_writer::PlainZstTableWriter>::new(
                out_dir.clone(),
            )
            .with_context(|| format!("Could not create {}", out_dir.display()))
        })
        .transpose()?;
    let line_details_dataset_writer = args
        .line_details_out_dir
        .map(|out_dir| {
            dataset_writer::ParallelDatasetWriter::<dataset_writer::PlainZstTableWriter>::new(
                out_dir.clone(),
            )
            .with_context(|| format!("Could not create {}", out_dir.display()))
        })
        .transpose()?;

    let strategy_factory = NaiveStrategyFactory::new(graph.clone());
    let mut processor = szz::SzzProcessor::new(graph, content_retriever, strategy_factory, blamer)
        .with_max_tree_diff_items(args.max_tree_diff_items)
        .with_max_blame_steps_per_hunk(args.max_blame_steps_per_hunk)
        .with_max_hunks(args.max_hunks)
        .with_max_modified_files(args.max_modified_files)
        .with_min_line_similarity_ratio(args.min_line_similarity_ratio)
        .with_max_similar_line_offset(args.max_similar_line_offset);

    if let Some(middle_revisions_out_dir) = args.middle_revisions_out_dir {
        processor = processor.with_middle_revisions_out_dir(&middle_revisions_out_dir)?;
    }

    let runtime_handle = tokio::runtime::Handle::current();
    tokio::task::block_in_place(|| {
        db::for_each_git_range(
            &conn,
            &pl,
            Some(
                args.concurrency
                    .unwrap_or(NonZeroUsize::try_from(num_cpus::get() * 10).expect("0 CPUs")),
            ),
            || {
                (
                    vuln_dataset_writer.as_ref().map(|vuln_dataset_writer| {
                        vuln_dataset_writer.get_thread_writer().unwrap()
                    }),
                    line_details_dataset_writer
                        .as_ref()
                        .map(|line_details_dataset_writer| {
                            line_details_dataset_writer.get_thread_writer().unwrap()
                        }),
                )
            },
            |(vuln_writer, line_details_writer), _rowid, filename, _doc, range| {
                runtime_handle.block_on(processor.process_range(
                    vuln_writer,
                    line_details_writer,
                    filename,
                    range,
                ))
            },
        )?;

        processor.flush()
    })?;

    let stats = processor.stats();

    pl.done();

    println!("{stats:#?}");

    Ok(())
}
