pileuphi_lib 0.9.3

use crate::bamio::OutputDataDest;
use crate::engine::BUFWRITER_CAP;
use crate::errors::Error;
use crate::position_queue::GenomeInterval;
use crate::utils::temp_fname;
use crate::utils::OutputWriter;
use crossbeam::channel::{unbounded, Receiver, Sender};
use log::{info, warn};
use std::collections::VecDeque;
use std::fs::File;
use std::io::BufReader;
use std::sync::{Arc, Mutex};

pub static FILE_MERGE_SINGLETON: Mutex<Vec<OutputDataDest>> = Mutex::new(vec![]);

/// A job a worker thread can accept. Contains temp output file handle, the interval to process, and a "done" state.
pub struct IntervalJobInner {
    pub out: OutputDataDest,
    pub interval: GenomeInterval,
    pub done: Mutex<bool>,
}

impl IntervalJobInner {
    fn new(interval: &GenomeInterval) -> Self {
        Self {
            out: OutputDataDest::from_string(&temp_fname(
                &format!("{}:{}-{}", interval.name, interval.start, interval.end),
                "",
                ".temp",
            )),
            done: Mutex::new(false),
            interval: interval.clone(),
        }
    }
}

pub type IntervalJob = Arc<IntervalJobInner>;

/// Data structure storing tasks to send to threads. Each interval job (chunk) is mapped to the larger interval it was fragmented from. Once all chunks are completed, their respective temp files are copied into the main output file and deleted. Importantly, larger intervals are ORDERED.
pub struct IntervalJobs {
    map: VecDeque<(GenomeInterval, Vec<IntervalJob>)>,
    pub queue: VecDeque<IntervalJob>,
    handle: std::thread::JoinHandle<()>,
    s: Sender<Vec<IntervalJob>>,
}

impl IntervalJobs {
    pub fn new(
        intervals: &[GenomeInterval],
        min_coords_per_thread: Option<i64>,
        threads: i64,
        output: OutputDataDest,
    ) -> Self {
        let mut map: VecDeque<(GenomeInterval, Vec<IntervalJob>)> = VecDeque::new();
        let mut queue: VecDeque<IntervalJob> = VecDeque::new();
        let mut lock = FILE_MERGE_SINGLETON.lock().unwrap();

        for interval in intervals {
            let chunks = match min_coords_per_thread {
                // we want to split refs into smaller sub chunks
                Some(min_coords_per_thread) => if interval.len() < min_coords_per_thread {
                    interval.chunks(min_coords_per_thread)
                } else {
                    interval.n_chunks(threads)
                }
                .map(|c| Arc::new(IntervalJobInner::new(&c)))
                .collect::<Vec<IntervalJob>>(),

                // or we want one thread to process one entire reference
                None => {
                    vec![Arc::new(IntervalJobInner::new(interval))]
                }
            };

            chunks.iter().for_each(|c| {
                queue.push_back(c.clone());
                lock.push(c.out.clone());
            });

            map.push_back((interval.clone(), chunks.clone()));
        }

        let (s, r): (Sender<Vec<IntervalJob>>, Receiver<Vec<IntervalJob>>) = unbounded();

        let handle = std::thread::spawn(move || {
            let mut main_writer = OutputWriter::new(&output, BUFWRITER_CAP, true, false).unwrap();
            while let Ok(temps) = r.recv() {
                for tmp in temps {
                    match tmp.out {
                        OutputDataDest::Stdout => {
                            panic!("cannot merge from stdout! Critical error")
                        }
                        OutputDataDest::File(ref f) => {
                            match File::open(f) {
                                Err(e) => {
                                    match e.kind() {
                                        std::io::ErrorKind::NotFound => (),
                                        _ => panic!("Failed to open output file for merging: {}", e),
                                    };
                                }

                                Ok(f) => {
                                    let mut reader = BufReader::with_capacity(2 * 1024 * 1024, f);
                                    if let Err(e) = std::io::copy(&mut reader, &mut main_writer.get()) {
                                        match e.kind() {
                                            std::io::ErrorKind::Interrupted => (),
                                            _ => panic!("{e}"),
                                        }
                                    }
                                }
                            }
                            if let Err(e) = std::fs::remove_file(f) {
                                match e.kind() {
                                    std::io::ErrorKind::NotFound => (),
                                    _ => panic!("{}", e),
                                }
                            }
                        }
                    }
                }
            }

            main_writer.flush().expect("Failed to flush merge writer");
        });

        Self { map, handle, queue, s }
    }

    pub fn is_completed(&self) -> bool {
        self.map.is_empty()
    }

    /// Check if we have any intervals with all chunks completed.
    pub fn merge_completed(&mut self) -> Result<(), Error> {
        let mut done = 0;

        if let Some((interval, pending)) = self.map.front() {
            for tmp in pending {
                if *tmp.done.lock().unwrap() {
                    done += 1;
                }
            }

            assert!(done <= pending.len());

            // all chunks have been marked "done" by their assigned workers
            if done == pending.len() {
                info!("Finished ref {}", interval.name);
                let (_, to_merge) = self.map.pop_front().unwrap();
                self.s.send(to_merge)?;
            }
        }

        Ok(())
    }

    /// Should be called when we've done everything (e.g. self.is_completed()). Does one final merge
    /// and signals the writer thread to finish.
    pub fn conclude(mut self) -> Result<(), Error> {
        self.merge_completed()?;
        drop(self.s);
        self.handle.join().expect("Failed to join writer thread");
        Ok(())
    }
}

/// Tell the program on unexpected exit (e.g. SIGTERM, ctrl-c) to delete all temp files it created and hasn't yet merged to final output
pub fn setup_exit_handler() {
    ctrlc::set_handler(|| {
        warn!("Received termination signal. Cleaning up intermediate files...");
        if let Ok(outputs) = FILE_MERGE_SINGLETON.lock() {
            for t in outputs.iter() {
                match t {
                    OutputDataDest::Stdout => (),
                    OutputDataDest::File(ref f) => {
                        if let Err(e) = std::fs::remove_file(f) {
                            match e.kind() {
                                std::io::ErrorKind::NotFound => (),
                                _ => eprintln!("{e}"),
                            }
                        }
                    }
                }
            }
        }

        std::process::exit(130);
    })
    .expect("Failed to set exit handler")
}