// 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::ffi::OsStr;
use std::fs::File;
use std::io::Cursor;
use std::io::Write;
use std::num::NonZeroUsize;
use std::os::unix::ffi::OsStrExt; // FIXME: Unix-only
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result};
use clap::Parser;
use dataset_writer::{ParallelDatasetWriter, TableWriter};
use dsi_progress_logger::{ProgressLog, concurrent_progress_logger};
use mimalloc::MiMalloc;
use rapidhash::RapidHashSet;
use rayon::prelude::*;
use sux13::traits::IndexedDict;
use swh_contents::retriever::ContentRetriever;
use swh_contents::retriever::rocksdb::{ReadonlyRocksdbContentCache, load_content_ids};
use swh_graph::graph::{SwhBidirectionalGraph, SwhGraphWithProperties, SwhLabeledForwardGraph};
use swh_graph::labels::DirEntry;
use swh_graph::mph::DynMphf;
use swh_graph::{NodeType, SWHID};

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

#[derive(Parser, Debug)]
/// Given the path to a directory of `.ndjson` files written by the `szz`, produces a partial
/// checkout of all revisions referenced by it, and writes them to a new directory of tarballs.
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,
    #[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 checkouts
    out_dir: PathBuf,
}

pub fn main() -> Result<()> {
    let 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);

    tracing::info!("Opening cache...");
    let content_retriever = ReadonlyRocksdbContentCache::new(&graph, &args.content_cache)
        .context("Could not load Rocksdb cache")?;
    let cache_cache = load_content_ids(&graph, &content_retriever)
        .context("Could not load set of already cached contents")?;
    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()))?;

    let num_revs = revs.len();

    let dataset_writer = ParallelDatasetWriter::<TarbleWriter>::new(args.out_dir)
        .context("Could not create directory")?;
    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "revision",
        local_speed = true,
        expected_updates = Some(num_revs),
    );
    pl.start("Listing files in each revision");

    let files: RapidHashSet<_> = revs
        .into_par_iter()
        .map_init(
            || (pl.clone(), dataset_writer.get_thread_writer().unwrap()),
            |(pl, writer), rev| {
                let Some(dir) = swh_graph_stdlib::find_root_dir(&graph, rev)
                    .expect("Could not lookup revision root dir")
                else {
                    pl.light_update();
                    return vec![];
                };
                let mut todo_stack = vec![dir];
                let mut path_stack = swh_graph_stdlib::collections::PathStack::new();
                path_stack.push([]);

                let mut files = Vec::new();
                while let Some(node) = todo_stack.pop() {
                    let reversed_path: Vec<_> = path_stack // root leftmost in the stack -> last in 'path'
                        .pop()
                        .expect("path_stack is shorter than todo_stack")
                        .collect();

                    for (succ, labels) in graph.untyped_labeled_successors(node) {
                        for label in labels {
                            let label = DirEntry::from(label);
                            todo_stack.push(succ);
                            path_stack.push(reversed_path.iter().copied()); // root last in 'path' -> first to be inserted -> leftmost in stack
                            path_stack.push_filename(label.label_name_id());
                        }
                    }
                    if graph.properties().node_type(node) == NodeType::Content {
                        let cnt = node;
                        if !cache_cache.contains(cnt) {
                            // if we never downloaded the content, it means neither it nor its next
                            // version is affected by a vulnerability or SZZ would have downloaded it
                            // already.
                            continue;
                        }
                        files.push(cnt);

                        let mut header = tar::Header::new_ustar();
                        let mut path = PathBuf::from("revisions");
                        path.push(swhid_to_path(&graph.properties().swhid(rev)));
                        let mut target_path = PathBuf::new();
                        for &path_part in reversed_path.iter().rev() {
                            let name = graph.properties().label_name(path_part);

                            // FIXME: Unix-only, but 'tar' forces us to by making the Path encoding
                            // platform-dependent instead of accepting bytestrings as it should
                            let name = OsStr::from_bytes(&name);

                            path.push(name);
                            target_path.push("..");
                        }
                        target_path.push("..");
                        target_path.push("contents");
                        target_path.push(swhid_to_path(&graph.properties().swhid(cnt)));

                        header.set_mode(0o644);
                        header.set_size(0);
                        header.set_entry_type(tar::EntryType::Symlink);
                        writer
                            .builder
                            .append_link(&mut header, path, target_path)
                            .expect("Could not write tar entry");
                    }
                }

                pl.light_update();
                files
            },
        )
        .flatten()
        .collect();
    pl.done();
    drop(pl);

    let mut files: Vec<_> = files.into_iter().collect();
    files.par_sort_unstable(); // Keep consecutive ids close for better compression

    let rt = tokio::runtime::Runtime::new().context("Could not initialize tokio runtime")?;

    let mut pl = concurrent_progress_logger!(
        display_memory = true,
        item_name = "file",
        local_speed = true,
        expected_updates = Some(files.len()),
    );
    pl.start("Writing files to tarballs");
    files.par_iter().try_for_each_init(
        || (pl.clone(), dataset_writer.get_thread_writer().unwrap()),
        |(pl, writer), cnt| -> Result<()> {
            let &cnt = cnt;
            let file_content = rt
                .block_on(content_retriever.get_pinned(cnt))
                .context("Could not get content")?
                .context("Content disappeared from rocksdb database")?;

            let mut header = tar::Header::new_ustar();
            header.set_size(
                u64::try_from(file_content.as_ref().len()).context("File size overflowed u64")?,
            );
            header.set_mode(0o644);
            let mut path = PathBuf::from("contents");
            path.push(swhid_to_path(&graph.properties().swhid(cnt)));
            writer
                .builder
                .append_data(&mut header, path, Cursor::new(file_content))
                .context("Could not write tar entry")?;
            pl.light_update();
            Ok(())
        },
    )?;
    pl.done();

    Ok(())
}

struct TarbleWriter {
    builder: tar::Builder<zstd::stream::write::Encoder<'static, File>>,
}

impl TableWriter for TarbleWriter {
    type Schema = ();
    type CloseResult = ();
    type Config = ();

    // Required methods
    fn new(path: PathBuf, _schema: (), _config: ()) -> Result<Self>
    where
        Self: Sized,
    {
        let path = path.with_extension("tar.zst");
        let file =
            File::create(&path).with_context(|| format!("Could not create {}", path.display()))?;

        // level 3  -> compression ratio 4.5%
        // level 6  -> compression ratio 4.2%
        // level 9  -> compression ratio 3.5%
        // level 19 -> compression ratio 3.2%
        let compression_level = 9;

        let encoder = zstd::stream::write::Encoder::new(file, compression_level)
            .context("Could not initialize zstd encoder")?;
        let builder = tar::Builder::new(encoder);
        Ok(Self { builder })
    }

    fn flush(&mut self) -> Result<()> {
        self.builder
            .get_mut()
            .flush()
            .context("Could not flush zstd encoder")
    }

    fn close(self) -> Result<Self::CloseResult> {
        self.builder
            .into_inner()
            .context("Could not flush tar builder")?
            .finish()
            .context("Could not flush zstd encoder")?
            .flush()
            .context("Could not flush zstd file")?;
        Ok(())
    }
}

fn swhid_to_path(swhid: &SWHID) -> String {
    // Windows and macOS don't allow ':' in file names
    swhid.to_string().replace(':', "_")
}
