use argh::FromArgs;
use bstr::ByteSlice;
use format_bytes::format_bytes;
use patchpal::common::{apply, commit_iter, diff_meta, git, io::WriteTo};
use patchpal::common::lines::{LineIndex, TypedLines};
use rayon::iter::{ParallelBridge, ParallelIterator};
use std::{cmp::min, ffi::OsString, fs, fs::File, io::Write, ops::Range, sync::mpsc};
use std::path::{Path, PathBuf};

#[derive(FromArgs)]
/// Extracts data from backports.
struct Args {
    /// the repository is a stable tree; if not specified, it is assumed to be a RHEL tree
    #[argh(switch, short = 's')]
    stable: bool,
    /// level of fuzziness to consider; 0: content conflicts only; 1: also context with fuzz;
    /// 2: everything including clean hunks
    #[argh(option, short = 'l', default = "0")]
    level: usize,
    /// the output directory; defaults to "output"
    #[argh(option, short = 'o', default = r#"String::from("output")"#)]
    output: String,
    /// show progress
    #[argh(switch, short = 'v')]
    verbose: bool,
    /// save duration of processing of each commit
    #[argh(switch)]
    duration: bool,
    /// path to the repository; must contain also upstream objects
    #[argh(positional)]
    repo: OsString,
    /// starting commit (exclusive)
    #[argh(positional)]
    start: String,
    /// ending commit (inclusive)
    #[argh(positional)]
    end: String,
}

// derive_more::AddAssign does not support fixed size arrays natively. Wrap them in a struct
// to be able to implement AddAssign manually.
#[derive(Clone, Copy, Default, derive_more::Index, derive_more::IndexMut,
         derive_more::IntoIterator)]
struct FuzzStats([usize; 3]);

#[derive(derive_more::AddAssign, Default)]
struct Stats {
    commits_too_long: usize,
    files_processed: usize,
    files_skipped: usize,
    files_missing: usize,
    hunks_processed: usize,
    hunks_too_long: usize,
    hunks_not_applied: usize,
    hunks_split: usize,
    pair_count: FuzzStats,
    pair_missing: FuzzStats,
    pair_okay: usize,
    pair_too_broad: usize,
    pair_cannot_encode: usize,
    pair_extra_code: usize,
    pair_chained: usize,
}

impl std::ops::AddAssign for FuzzStats {
    fn add_assign(&mut self, other: Self) {
        for i in 0 .. 3 {
            self.0[i] += other.0[i];
        }
    }
}

impl Stats {
    fn write_to(&self, f: &mut impl Write) {
        writeln!(f, "too long commits: {}", self.commits_too_long).unwrap();
        writeln!(f, "files total: {}, processed: {}, skipped: {}, not downstream: {}",
                 self.files_processed + self.files_skipped +
                 self.files_missing,
                 self.files_processed, self.files_skipped,
                 self.files_missing).unwrap();
        writeln!(f, "  hunks total: {}, processed: {}, not applied: {}, too long: {}, split: {}",
                 self.hunks_processed + self.hunks_not_applied + self.hunks_too_long,
                 self.hunks_processed, self.hunks_not_applied, self.hunks_too_long,
                 self.hunks_split).unwrap();
        writeln!(f, "  pairs total: {}   (content fuzz / context fuzz / no fuzz)",
                 self.pair_count.into_iter().sum::<usize>()).unwrap();
        writeln!(f, "    total: {} / {} / {}",
                 self.pair_count[0], self.pair_count[1], self.pair_count[2]).unwrap();
        writeln!(f, "    missing: {} / {} / {}",
                 self.pair_missing[0], self.pair_missing[1], self.pair_missing[2]).unwrap();
        writeln!(f, "    too broad: {}, can't encode: {}, extra code in backport: {}, chained: {}",
                 self.pair_too_broad, self.pair_cannot_encode, self.pair_extra_code,
                 self.pair_chained).unwrap();
        writeln!(f, "    successful: {}", self.pair_okay).unwrap();
    }
}

fn get_function_range(
    lines: &TypedLines,
    code_range: Range<LineIndex>,
) -> Range<LineIndex> {
    let tag_end = lines.tag_end().unwrap();
    let (code_start, code_end) = if code_range.is_empty() {
        (
            code_range.start.saturating_sub(1),
            min(code_range.end + 1, tag_end)
        )
    } else {
        (code_range.start, code_range.end)
    };
    let tag = lines.tag(code_start).unwrap();
    let start = match lines.tag_slice(0.into() .. code_start)
        .unwrap()
        .iter()
        .rev()
        .position(|&t| t != tag)
    {
        Some(idx) => (code_start - 1) - idx + 1,
        None => 0.into(),
    };
    let tag = lines.tag(code_end - 1).unwrap();
    let end = match lines.tag_slice(code_end .. tag_end)
        .unwrap()
        .iter()
        .position(|&t| t != tag)
    {
        Some(idx) => code_end + idx - 1,
        None => tag_end,
    };
    start .. end
}

fn process(
    fileno: usize,
    tx_stats: mpsc::Sender<Stats>,
    (downstream, upstream): (Vec<u8>, Vec<u8>),
    repo: git::Git,
    level: usize,
    output_dir: PathBuf,
    duration: bool,
) {
    let mut stats = Stats::default();
    let file_changes = repo.raw_diff(&upstream)
        .unwrap_or_else(|err|
            panic!("cannot get patch {}: {}", upstream.as_bstr(), err.to_string())
        );
    if file_changes.len() == 0 {
        // skip backports of merges (a merge commit produces an empty diff)
        return;
    }
    let mut downstream_changes = repo.raw_diff(&downstream)
        .unwrap_or_else(|err|
            panic!("cannot get patch {}: {}", downstream.as_bstr(), err.to_string())
        );
    let start_time = std::time::SystemTime::now();
    if duration {
        let fname = output_dir.join(&format!("{fileno:06}")).with_extension("time");
        let mut f = File::create(fname).unwrap();
        f.write(&downstream).unwrap();
        f.write(b" / ").unwrap();
        f.write(&upstream).unwrap();
        f.write(b"\n").unwrap();
        let start_secs = start_time
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs();
        f.write(&format_bytes!(b"start {}\n", start_secs)).unwrap();
    }
    let mut files = Vec::with_capacity(file_changes.len());
    for change in file_changes {
        match &change.metadata.op {
            diff_meta::Op::Modified { .. } => (),
            _ => {
                stats.files_skipped += 1;
                continue;
            }
        }
        let old_name = change.metadata.name.clone();
        let hunks = change.data.diff();
        files.push((hunks, old_name));
    }
    if files.iter().map(|(hunks, _)| hunks.len()).sum::<usize>() > 350 {
        // Skip commits with too many hunks. Their processing takes too long and
        // none of those patches contribute anything to the samples.
        let stats = Stats {
            commits_too_long: 1,
            ..Default::default()
        };
        tx_stats.send(stats).unwrap();
        return;
    }
    for (hunks, old_name) in files {
        // Find the corresponding downstream file.
        let change = match downstream_changes.iter()
            .position(|ch| ch.metadata.name == old_name)
        {
            Some(pos) => downstream_changes.remove(pos),
            None => {
                stats.files_missing += 1;
                continue;
            }
        };
        // ignore file that got deleted in the backported tree (e.g.
        // RHEL commit af405037d992)
        match change.metadata.op {
            diff_meta::Op::Modified { .. } => (),
            _ => {
                stats.files_missing += 1;
                continue;
            }
        }
        stats.files_processed += 1;

        let resolved_hunks = change.data.diff();
        // Reuse the "before" downstream data. It's exactly the file we're patching.
        let mut lines = change.data.src;
        // Annotate the data first. We'll need the annotation later when extracting the
        // enclosing function.
        lines.annotate();
        let mut extract = apply::Extract::new(lines);
        extract.apply.add_hunks(hunks);
        extract.add_resolved_hunks(&resolved_hunks);

        stats.hunks_processed += extract.apply.collected_hunks();
        stats.hunks_too_long += extract.apply.too_long;
        stats.hunks_not_applied += extract.apply.missed;
        stats.hunks_split += extract.apply.split;

        let mut it = extract.iter_pairs();
        for (repl_applied, repl_extracted) in it.by_ref() {
            let fuzz_level = if repl_applied.fuzz_content {
                0
            } else if
                repl_applied.fuzz_top || repl_applied.fuzz_bottom || repl_applied.fuzz_middle
            {
                1
            } else {
                2
            };
            stats.pair_count[fuzz_level] += 1;
            let repl_extracted = match repl_extracted {
                Some(x) => x,
                None => {
                    stats.pair_missing[fuzz_level] += 1;
                    continue;
                }
            };
            if fuzz_level > level {
                continue;
            }
            let context_range = repl_applied.range_with_context();
            if repl_extracted.range.start < context_range.start ||
                repl_extracted.range.end > context_range.end
            {
                // the resolution is broader than the context
                stats.pair_too_broad += 1;
                continue;
            }

            stats.pair_okay += 1;
            let base_fname = output_dir.join(&format!("{fileno:06}-{:02}", stats.pair_okay));
            let mut f = File::create(base_fname.with_extension("desc")).unwrap();
            f.write(&downstream).unwrap();
            f.write(b" / ").unwrap();
            f.write(&upstream).unwrap();
            f.write(b"\n").unwrap();
            f.write(&old_name).unwrap();
            f.write(b"\n").unwrap();
            let mut f_a = File::create(base_fname.with_extension("patch")).unwrap();
            let mut f_i = File::create(base_fname.with_extension("raw_i")).unwrap();
            let mut f_o = File::create(base_fname.with_extension("raw_o")).unwrap();
            extract.write_pair_to(
                repl_applied, repl_extracted, &mut f_a, &mut f_i, &mut f_o
            ).unwrap();

            // the context
            let range = get_function_range(&extract.apply.lines, repl_applied.range.clone());
            let mut f_c = File::create(base_fname.with_extension("raw_c")).unwrap();
            extract.apply.lines.line_slice(range).write_to(&mut f_c).unwrap();
        }
        stats.pair_extra_code += it.extracted_skipped;
        stats.pair_chained += it.chained;
    }
    tx_stats.send(stats).unwrap();
    if duration {
        let end_time = std::time::SystemTime::now();
        let fname = output_dir.join(&format!("{fileno:06}")).with_extension("time");
        let mut f = std::fs::OpenOptions::new().append(true).open(fname).unwrap();
        let end_secs = end_time
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs();
        let dur_secs = end_time
            .duration_since(start_time)
            .unwrap()
            .as_secs_f32();
        f.write(&format_bytes!(b"end {}\n", end_secs)).unwrap();
        f.write(&format_bytes!(b"duration {}\n", dur_secs)).unwrap();
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Args = argh::from_env();
    let repo = git::Git::open(&args.repo)
        .map_err(|e| format!("Invalid git repository path: {}", e))?;
    let output_dir = Path::new(&args.output);
    fs::create_dir_all(&output_dir)
        .map_err(|e| format!("cannot create directory {}: {}", output_dir.display(), e))?;
    let parser: Box<dyn commit_iter::CommitParser + Send> =
        if args.stable {
            Box::new(commit_iter::KernelStableCommitParser::new())
        } else {
            Box::new(commit_iter::DownstreamCommitParser::new())
        };

    // the stats writing thread
    let (tx_stats, rx_stats) = mpsc::channel();
    let verbose = args.verbose;
    let output_dir2 = output_dir.to_path_buf();
    let stat_thread = std::thread::spawn(move || {
        let mut stats = Stats::default();
        for substats in rx_stats.iter() {
            let prev_files_processed = stats.files_processed;
            stats += substats;
            if verbose && prev_files_processed / 1000 != stats.files_processed / 1000 {
                stats.write_to(&mut std::io::stdout());
            }
        }
        let mut f = File::create(output_dir2.join("stats")).unwrap();
        stats.write_to(&mut f);
        stats.write_to(&mut std::io::stdout());
    });

    // the output data writing threads
    commit_iter::ResolvedCommitsIterator::new_with_parser(
        parser,
        &repo.walk(args.end.as_bytes()).base(Some(args.start.as_bytes())),
    )
        .ignore_reverts()
        .verify_in(repo.clone())
        .map_while(|v| v.ok())
        .filter_map(|(d_sha, r)| r.map(|r_sha| (d_sha, r_sha)))
        .enumerate()
        .par_bridge()
        .for_each(|(fileno, sha_pair)| process(
            fileno,
            tx_stats.clone(),
            sha_pair,
            repo.clone(),
            args.level,
            output_dir.to_path_buf(),
            args.duration,
        ));
    drop(tx_stats);
    stat_thread.join().unwrap();
    Ok(())
}
