use std::{path::Path, pin::Pin};
use anyhow::Context;
use bstr::{BStr, ByteSlice};
use rust_htslib::bcf::Record;
use tokio::io::{AsyncWrite, AsyncWriteExt, BufWriter};
enum RegionsFormat {
Bed,
Tab,
}
pub struct RegionWriter {
writer: Pin<Box<dyn AsyncWrite + Send>>,
format: RegionsFormat,
}
impl RegionWriter {
pub async fn from_parameter(param: Option<&str>) -> anyhow::Result<Option<Self>> {
if let Some(param) = param {
let rw = if param == "-" {
let writer = Box::pin(tokio::io::stdout());
RegionWriter {
writer,
format: RegionsFormat::Tab,
}
} else {
let file = tokio::fs::File::create(param).await?;
let writer = Box::pin(BufWriter::new(file));
let format = if Path::new(param)
.extension()
.is_some_and(|ext| ext.eq_ignore_ascii_case("bed"))
{
RegionsFormat::Bed
} else {
RegionsFormat::Tab
};
RegionWriter { writer, format }
};
Ok(Some(rw))
} else {
Ok(None)
}
}
pub async fn write<'a>(
&mut self,
records: impl DoubleEndedIterator<Item = &'a Record>,
) -> anyhow::Result<()> {
let (reg, pos, end) = self
.get_coordinates(records)
.with_context(|| "Couldn't extract coordinates")?;
eprintln!("Now writing regions ...");
self.writer
.write_all(format!("{reg}\t{pos}\t{end}\n").as_bytes())
.await?;
Ok(())
}
pub async fn flush(&mut self) -> std::io::Result<()> {
self.writer.flush().await
}
fn get_coordinates<'a>(
&self,
mut records: impl DoubleEndedIterator<Item = &'a Record>,
) -> Option<(&'a BStr, i64, i64)> {
let first = records.next()?;
let header = first.header();
let rstr = header.rid2name(first.rid()?).ok()?.as_bstr();
let pos_0_incl = first.pos();
let end_0_excl = records.next_back().unwrap_or(first).end();
let (pos, end) = match self.format {
RegionsFormat::Bed => (pos_0_incl, end_0_excl),
RegionsFormat::Tab => (pos_0_incl + 1, end_0_excl ),
};
Some((rstr, pos, end))
}
}