twitcher 0.6.6

use std::{
    fs::File,
    path::{Path, PathBuf},
    time::Duration,
};

use anyhow::{Context, bail};
use bstr::ByteSlice;
use itertools::Itertools;
use rust_htslib::{
    bam,
    bcf::{
        self, HeaderRecord, Read as _,
        header::{HeaderView, TagLength, TagType},
    },
    htslib,
};
use tracing::{debug, instrument, warn};
use url::Url;

use crate::{
    RunnableCommand,
    common::{
        aligner::RUNNING,
        csv::TwitcherCSVWriter,
        list_of_regions::{Regions, RegionsDefinition, Targets},
        reference::ReferenceReader,
    },
    vcf::{
        cli::{Command, OutputOptions},
        pipeline::{
            VCFPipeline,
            reader::VCFReader,
            writer::{OutputWriter, region_writer::RegionWriter},
        },
        strings::{VCF_FORMAT_GT_LINE, VCF_FORMAT_PS_LINE},
    },
};

pub mod cli;
pub mod pipeline;
pub mod strings;

impl RunnableCommand for Command {
    #[instrument(name = "vcf", skip_all)]
    async fn run(self) -> anyhow::Result<()> {
        tokio::spawn(async {
            let mut interval = tokio::time::interval(Duration::from_secs(1));
            loop {
                interval.tick().await;
                let r = RUNNING.load(std::sync::atomic::Ordering::Relaxed);
                debug!("Running tasks: {r}");
            }
        });
        self.validate()?;
        let regions = self.regions.read_regions().await?;
        let input = tokio::task::block_in_place(move || open_input_file(&self.input, regions))?;
        ensure_written_format_fields(input.header());
        let mut header = bcf::Header::from_template(input.header());
        augment_header(&mut header, self.no_version);
        let reference = tokio::task::block_in_place(|| ReferenceReader::try_from(&self.reference))?;
        let targets = self.targets.read_regions().await?.map(Targets::new);
        let region_output = RegionWriter::from_parameter(self.ts_regions.as_deref()).await?;
        let output = tokio::task::block_in_place(move || open_output_file(&self.output, &header))?;
        let csv_output = if let Some(f) = self.csv_output {
            Some(TwitcherCSVWriter::new(Box::new(File::create(f)?)))
        } else {
            None
        };
        let bam = open_phasing_bam_file(self.phasing_bam.as_deref(), &self.reference.file)?;
        VCFPipeline {
            input,
            output,
            reference,
            targets,
            region_output,
            csv_output,
            settings: self.pipeline_settings,
            bam,
            phasing: self.phasing,
        }
        .run()
        .await?;

        Ok(())
    }
}

impl Command {
    fn validate(&self) -> anyhow::Result<()> {
        if self.output.output_file.as_ref().is_none_or(|o| o == "-")
            && self.ts_regions.as_ref().is_some_and(|o| o == "-")
        {
            bail!("Cannot output both vcf and regions output on stdout.");
        }
        self.pipeline_settings.cluster.cluster_strategy.validate()?;
        Ok(())
    }
}

fn augment_header(header: &mut bcf::Header, no_version: bool) {
    #[allow(clippy::wildcard_imports)]
    use strings::*;
    // Eventually, what we add in this method will likely depend on the settings etc ...
    header.push_record(
        format!(
            "##INFO=<ID={VCF_TS_CIGARETS_KEY},Number=1,Type=String,Description=\"{VCF_TS_CIGARETS_DESC}\">"
        )
        .as_bytes(),
    );
    header.push_record(
        format!(
            "##INFO=<ID={VCF_TS_COST_KEY},Number=1,Type=Integer,Description=\"{VCF_TS_COST_DESC}\">"
        )
        .as_bytes(),
    );
    header.push_record(
        format!(
            "##INFO=<ID={VCF_TS_INNER_LEN_KEY},Number=.,Type=Integer,Description=\"{VCF_TS_INNER_LEN_DESC}\">"
        )
        .as_bytes(),
    );
    header.push_record(
        format!(
            "##INFO=<ID={VCF_TS_JUMP_KEY},Number=.,Type=Integer,Description=\"{VCF_TS_JUMP_DESC}\">"
        )
        .as_bytes(),
    );
    header.push_record(
        format!(
            "##INFO=<ID={VCF_TS_ID_KEY},Number=1,Type=String,Description=\"{VCF_TS_ID_DESC}\">"
        )
        .as_bytes(),
    );
    header.push_record(
        format!(
            "##INFO=<ID={VCF_CLUSTER_GRP_KEY},Number=1,Type=String,Description=\"{VCF_CLUSTER_GRP_DESC}\">"
        )
        .as_bytes(),
    );
    header.push_record(
        format!(
            "##INFO=<ID={VCF_TWITCHER_PHASE_KEY},Number=0,Type=Flag,Description=\"{VCF_TWITCHER_PHASE_DESC}\">"
        )
        .as_bytes(),
    );
    header.push_record(VCF_FORMAT_GT_LINE);
    header.push_record(VCF_FORMAT_PS_LINE);
    if !no_version {
        header.push_record(
            format!("##{VCF_HEADER_TWITCHER_VERSION_KEY}={TWITCHER_VERSION}",).as_bytes(),
        );
        header.push_record(
            format!("##{VCF_HEADER_TWITCHER_CMD}={}", std::env::args().join(" ")).as_bytes(),
        );
    }
}

/// Declare the FORMAT fields this program writes itself on the header of the already-open
/// input file, unless the input declares them already.
///
/// Records are written with the (augmented) output header, but local read-based phasing
/// (`--phasing-bam`) sets GT and PS while the records still carry the *input* header, and
/// htslib rejects a tag that the record's own header does not know. Declaring the fields here
/// rather than only in [`augment_header`] covers both: the output header is duplicated from
/// this one, so it inherits whatever is added.
fn ensure_written_format_fields(header: &HeaderView) {
    ensure_format_field(
        header,
        b"GT",
        (TagType::String, TagLength::Fixed(1)),
        VCF_FORMAT_GT_LINE,
    );
    ensure_format_field(
        header,
        b"PS",
        (TagType::Integer, TagLength::Fixed(1)),
        VCF_FORMAT_PS_LINE,
    );
}

/// Append `line` to an open header unless `tag` is already declared there with type `expected`.
///
/// A tag declared with a different type cannot be fixed up — htslib keeps the first declaration
/// of an ID — so that case is only reported; writing the field is skipped later on.
fn ensure_format_field(
    header: &HeaderView,
    tag: &[u8],
    expected: (TagType, TagLength),
    line: &[u8],
) {
    match header.format_type(tag) {
        Ok(found) if found == expected => return,
        Ok(found) => {
            warn!(
                "Input header declares FORMAT/{} as {found:?}, but this program writes {expected:?}. \
                 The field will not be written.",
                tag.as_bstr()
            );
            return;
        }
        Err(_) => {}
    }

    let Ok(line) = std::ffi::CString::new(line) else {
        warn!(
            "Cannot declare FORMAT/{}: header line contains a nul byte.",
            tag.as_bstr()
        );
        return;
    };
    // SAFETY: the pointer is used only for the duration of this call, while `header` is alive.
    // Appending only extends the header dictionaries, so ids already referenced by records read
    // from this header keep their meaning.
    let failed = unsafe {
        let inner = header.as_ptr();
        htslib::bcf_hdr_append(inner, line.as_ptr()) != 0 || htslib::bcf_hdr_sync(inner) != 0
    };
    if failed {
        warn!(
            "Failed to add a FORMAT/{} declaration to the input header.",
            tag.as_bstr()
        );
    } else {
        debug!(
            "Added a missing FORMAT/{} declaration to the header.",
            tag.as_bstr()
        );
    }
}

fn open_input_file(input: &str, regions: Option<Regions>) -> anyhow::Result<VCFReader> {
    let reader = if let Some(regions) = regions {
        let inner = match input {
            "-" => bail!("Cannot create indexed reader from stdin"),
            path_or_url => {
                if let Ok(url) = Url::parse(path_or_url) {
                    bcf::IndexedReader::from_url(&url)
                } else {
                    bcf::IndexedReader::from_path(path_or_url)
                }
            }
        }?;

        let regions =
            regions.into_linear(|chr| inner.header().name2rid(chr.as_ref()).unwrap() as usize);
        VCFReader::with_regions(inner, regions, input.to_string())
    } else {
        let inner = match input {
            "-" => bcf::Reader::from_stdin(),
            path_or_url => {
                if let Ok(url) = Url::parse(path_or_url) {
                    bcf::Reader::from_url(&url)
                } else {
                    bcf::Reader::from_path(path_or_url)
                }
            }
        }?;
        VCFReader::entire_file(inner, input.to_string())
    };

    let sample_count = reader.header().sample_count();
    if sample_count > 1 {
        let first_sample = reader.header().samples()[0].as_bstr();
        warn!(
            "The input file has {sample_count} samples, but twitcher only supports one sample per file. All but the first sample ({first_sample}) will be ignored."
        );
    }
    if sample_count == 0 {
        bail!(
            "There needs to be a sample present in the input VCF file. Sites-only VCFs are currently not supported."
        );
    }

    validate_version(get_vcf_version(reader.header())?)?;

    Ok(reader)
}

fn validate_version((major, minor): (u8, u8)) -> anyhow::Result<()> {
    if major != 4 || (!(2..5).contains(&minor)) {
        bail!("VCF version unsupported. Twitcher only handles v4.2 <= version <= v4.5");
    }
    Ok(())
}

fn get_vcf_version(header: &HeaderView) -> anyhow::Result<(u8, u8)> {
    let records = header.header_records();
    let first = records
        .first()
        .with_context(|| "Header does not have at least one record.")?;
    if let HeaderRecord::Generic { key, value } = first
        && key == "fileformat"
        && value.starts_with("VCFv")
    {
        let dot = value
            .find('.')
            .with_context(|| "No separator between minor and major version")?;
        let major: u8 = value[4..dot].parse()?;
        let minor: u8 = value[dot + 1..].parse()?;
        return Ok((major, minor));
    }
    bail!("The first header record of the VCF input file must be the VCF version.");
}

fn open_phasing_bam_file<P: AsRef<Path>>(
    arg: Option<&str>,
    reference_path: P,
) -> anyhow::Result<Option<bam::IndexedReader>> {
    let Some(path) = arg else {
        return Ok(None);
    };

    let mut reader = if let Ok(url) = Url::parse(path) {
        bam::IndexedReader::from_url(&url)?
    } else {
        bam::IndexedReader::from_path(path)?
    };

    reader.set_reference(reference_path)?;

    Ok(Some(reader))
}

fn open_output_file(output: &OutputOptions, header: &bcf::Header) -> anyhow::Result<OutputWriter> {
    let inner = match output.output_file.as_deref() {
        None | Some("-") => bcf::Writer::from_stdout(header, true, bcf::Format::Vcf),
        Some(path_or_url) => {
            let url = Url::parse(path_or_url);
            let path = PathBuf::from(match &url {
                Ok(url) => url.path(),
                _ => path_or_url,
            });
            let compressed = matches!(
                path.extension().and_then(|e| e.to_str()),
                Some("gz" | "bgz")
            );
            let format = if path
                .file_name()
                .is_none_or(|n| n.to_str().is_some_and(|s| s.contains(".vcf")))
            {
                bcf::Format::Vcf
            } else {
                bcf::Format::Bcf
            };
            match url {
                Ok(url) => bcf::Writer::from_url(&url, header, !compressed, format),
                Err(_) => bcf::Writer::from_path(path, header, !compressed, format),
            }
        }
    }?;

    let writer = if let Some(plus_minus) = output.output_filter.realigned_and_context {
        OutputWriter::new_buffered(inner, i64::try_from(plus_minus)?)
    } else {
        OutputWriter::new_native(inner, output.output_filter.only_realigned)
    };
    Ok(writer)
}

#[cfg(test)]
mod tests {
    use std::io::Read as _;

    use rust_htslib::{bcf::Header, bgzf};

    use crate::common::reference::CliReferenceArg;

    use super::*;

    #[test]
    fn validate_rejects_both_outputs_on_stdout() {
        let cmd = Command {
            output: None::<&str>.into(), // stdout (default)
            ts_regions: Some("-".into()),
            ..Default::default()
        };
        assert!(cmd.validate().is_err());
    }

    #[test]
    fn validate_allows_vcf_on_file_regions_on_stdout() {
        let cmd = Command {
            output: Some("/tmp/out.vcf").into(),
            ts_regions: Some("-".into()),
            ..Default::default()
        };
        assert!(cmd.validate().is_ok());
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_pipeline() {
        let dir = tempfile::tempdir().unwrap();
        let outpath = format!("{}/out.vcf", dir.path().to_str().unwrap());
        let cmd = Command {
            input: "./test_files/test.vcf".to_string(),
            reference: CliReferenceArg::from("./test_files/test.fa"),
            output: Some(outpath.clone()).into(),
            ..Default::default()
        };
        cmd.run().await.unwrap();
        let output = std::fs::read_to_string(outpath).unwrap();
        let out_records = output.lines().filter(|l| !l.starts_with('#'));
        assert_eq!(out_records.count(), 2);
    }

    /// An input header that does not declare PS must not stop the clusterizer from writing PS
    /// onto records, which still carry that very header (see `ensure_written_format_fields`).
    #[test]
    fn ps_is_writable_after_ensuring_format_fields() {
        let mut reader =
            open_input_file("./test_files/test_no_ps_format_header.vcf", None).unwrap();
        assert!(
            reader.header().format_type(b"PS").is_err(),
            "fixture must not declare PS"
        );

        let mut record = reader.empty_record();
        reader.read(&mut record).unwrap().unwrap();
        assert!(
            record.push_format_integer(b"PS", &[201]).is_err(),
            "without the declaration, htslib rejects PS"
        );

        ensure_written_format_fields(record.header());
        record.push_format_integer(b"PS", &[201]).unwrap();
        assert_eq!(record.format(b"PS").integer().unwrap()[0], [201]);
    }

    /// A PS declaration of an incompatible type cannot be overridden, and must be left alone
    /// instead of producing a header with two conflicting PS lines.
    #[test]
    fn ensure_format_field_keeps_incompatible_declaration() {
        let reader = open_input_file("./test_files/test.vcf", None).unwrap();
        ensure_format_field(
            reader.header(),
            b"PS",
            (TagType::String, TagLength::Fixed(1)),
            b"##FORMAT=<ID=PS,Number=1,Type=String,Description=\"Phase set\">",
        );
        assert_eq!(
            reader.header().format_type(b"PS").unwrap(),
            (TagType::Integer, TagLength::Fixed(1))
        );
    }

    #[test]
    fn test_auto_detect_output_format_vcf() {
        let reader = open_input_file("./test_files/test.vcf", None).unwrap();
        let header = Header::from_template(reader.header());
        let dir = tempfile::tempdir().unwrap();
        let outpath_vcf = format!("{}/out.vcf", dir.path().to_str().unwrap());
        let writer = open_output_file(&Some(&outpath_vcf).into(), &header).unwrap();
        std::mem::drop(writer);

        let written_file = std::fs::read_to_string(outpath_vcf).unwrap();
        assert!(written_file.starts_with("##fileformat=VCFv4"));
    }

    #[test]
    fn test_auto_detect_output_format_vcf_gz() {
        let reader = open_input_file("./test_files/test.vcf", None).unwrap();
        let header = Header::from_template(reader.header());
        let dir = tempfile::tempdir().unwrap();
        let outpath_vcf_gz = format!("{}/out.vcf.gz", dir.path().to_str().unwrap());
        let writer = open_output_file(&Some(&outpath_vcf_gz).into(), &header).unwrap();
        std::mem::drop(writer);

        let mut written_file = Vec::new();
        bgzf::Reader::from_path(outpath_vcf_gz)
            .unwrap()
            .read_to_end(&mut written_file)
            .unwrap();
        assert!(written_file.starts_with(b"##fileformat=VCFv4"));
    }

    #[test]
    fn test_auto_detect_output_format_bcf() {
        let reader = open_input_file("./test_files/test.vcf", None).unwrap();
        let header = Header::from_template(reader.header());
        let dir = tempfile::tempdir().unwrap();
        let outpath_bcf = format!("{}/out.bcf", dir.path().to_str().unwrap());
        let writer = open_output_file(&Some(&outpath_bcf).into(), &header).unwrap();
        std::mem::drop(writer);

        let mut written_file = Vec::new();
        bgzf::Reader::from_path(outpath_bcf)
            .unwrap()
            .read_to_end(&mut written_file)
            .unwrap();
        assert!(written_file.starts_with(b"BCF"));
    }

    #[test]
    fn test_auto_detect_output_format_bcf_gz() {
        // .bcf.gz should be the same as bcf
        let reader = open_input_file("./test_files/test.vcf", None).unwrap();
        let header = Header::from_template(reader.header());
        let dir = tempfile::tempdir().unwrap();
        let outpath_bcf_gz = format!("{}/out.bcf.gz", dir.path().to_str().unwrap());
        let writer = open_output_file(&Some(&outpath_bcf_gz).into(), &header).unwrap();
        std::mem::drop(writer);

        let mut written_file = Vec::new();
        bgzf::Reader::from_path(outpath_bcf_gz)
            .unwrap()
            .read_to_end(&mut written_file)
            .unwrap();
        assert!(written_file.starts_with(b"BCF"));
    }
}