twitcher 0.7.0

use std::collections::HashMap;
use std::hash::Hash;
use std::io::Write;
use std::process::Stdio;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU8};
use std::sync::{LazyLock, Mutex, MutexGuard};
use std::time::{Duration, Instant};

use anyhow::Context;
use bstr::ByteSlice;
use compact_genome::implementation::alphabets::dna_alphabet_or_n::DnaAlphabetOrN;
use lib_tsalign::a_star_aligner::alignment_geometry::AlignmentRange;
use lib_tsalign::config::TemplateSwitchConfig;
use lib_tsalign::costs::U64Cost;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt};
use tokio::sync::watch::error::RecvError;
use tokio::sync::{Semaphore, watch};
use tracing::{error, instrument, trace};

use crate::common::aligner::cli::MLSSelector;
use crate::common::aligner::db::{Database, StaticAlignmentKey};
use crate::common::aligner::exec_stats::{
    AlignmentSource, AlignmentStatsContext, AlignmentStatsWriter, ExecStats, MemoryUsage,
};
use crate::common::aligner::result::{SoftFailureReason, TwitcherAlignmentWithStatistics};
use crate::common::aligner::{
    fpa::FourPointAligner,
    result::{AlignmentFailure, TwitcherAlignment},
};
use crate::common::coords::GenomeRegion;
use crate::common::{ImmutableSequence, SequencePair};
use crate::counter;
use crate::worker::{WorkerQuery, WorkerQueryMetadata, WorkerResult};

pub mod cli;
mod db;
pub mod exec_stats;
pub mod fpa;
pub mod result;

/// Use this to monitor how many tasks are running
pub static RUNNING: LazyLock<AtomicU8> = LazyLock::new(|| AtomicU8::new(0));

pub struct InMemoryCache {
    in_progress: HashMap<AlignmentKey, watch::Sender<Option<Arc<TwitcherAlignment>>>>,
    finished: HashMap<AlignmentKey, Arc<TwitcherAlignment>>,
}

pub struct AlignmentOrchestrator {
    aligners: Arc<AlignerSelector>,
    pub costs: Arc<TemplateSwitchConfig<DnaAlphabetOrN, U64Cost>>,
    parallelism: Arc<Semaphore>,
    per_alignment_settings: PerAlignmentSettings,
    in_memory_cache: Option<Arc<Mutex<InMemoryCache>>>,
    database: Option<Arc<Mutex<Database>>>,
    failed_alignment_writer: Option<Arc<Mutex<Box<dyn Write + Send>>>>,
    stats_writer: Option<Arc<AlignmentStatsWriter>>,
}

#[derive(PartialEq, Eq, Clone)]
pub struct AlignmentKey {
    // implicit part of key:
    // - reference
    // - aligner plus settings
    // - cost function
    reference_region: GenomeRegion,
    alignment_ranges: AlignmentRange,
    query_sequence: ImmutableSequence,
}

impl Hash for AlignmentKey {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.reference_region.hash(state);
        self.alignment_ranges.reference_offset().hash(state);
        self.alignment_ranges.reference_limit().hash(state);
        self.alignment_ranges.query_offset().hash(state);
        self.alignment_ranges.query_limit().hash(state);
        self.query_sequence.hash(state);
    }
}

pub struct PerAlignmentSettings {
    memory_allowance: usize,
    timeout: Option<Duration>,
}

/// Everything needed to report one finished alignment to the optional output files.
///
/// This is per-request rather than per-computation: a result served from the cache or the database
/// is reported again, with its `source` set accordingly.
struct ResultReporting {
    failed_writer: Option<Arc<Mutex<Box<dyn Write + Send>>>>,
    stats_writer: Option<Arc<AlignmentStatsWriter>>,
    source: AlignmentSource,
    reference_length: usize,
    query_length: usize,
}

pub struct InProgress {
    receiver: watch::Receiver<Option<Arc<TwitcherAlignment>>>,
    cluster_region: GenomeRegion,
    reporting: ResultReporting,
    /// Guards against reporting twice if a consumer calls [`InProgress::recv`] more than once.
    reported: bool,
}

impl InProgress {
    const fn new(
        receiver: watch::Receiver<Option<Arc<TwitcherAlignment>>>,
        cluster_region: GenomeRegion,
        reporting: ResultReporting,
    ) -> Self {
        Self {
            receiver,
            cluster_region,
            reporting,
            reported: false,
        }
    }

    pub async fn recv(&mut self) -> Result<Arc<TwitcherAlignment>, RecvError> {
        #[expect(
            clippy::expect_used,
            reason = "compiler cannot statically verify, but correctness clear"
        )]
        let result = (*self.receiver.wait_for(Option::is_some).await?)
            .as_ref()
            .expect("there is a value, because the condition that we wait on is Option::is_some")
            .clone();
        if !self.reported {
            self.reported = true;
            count_result(&result);
            self.report(&result);
        }
        Ok(result)
    }

    fn report(&self, result: &TwitcherAlignment) {
        if result.outcome.is_err()
            && let Some(w) = &self.reporting.failed_writer
        {
            let _ = writeln!(w.lock().unwrap(), "{}", self.cluster_region);
        }
        if let Some(w) = &self.reporting.stats_writer {
            w.write(
                result,
                &AlignmentStatsContext {
                    cluster_region: &self.cluster_region,
                    source: self.reporting.source,
                    reference_length: self.reporting.reference_length,
                    query_length: self.reporting.query_length,
                },
            );
        }
    }
}

impl AlignmentOrchestrator {
    pub fn enable_cache(&mut self) {
        let cache = InMemoryCache {
            in_progress: HashMap::new(),
            finished: HashMap::new(),
        };
        self.in_memory_cache = Some(Arc::new(Mutex::new(cache)));
    }

    fn lock_cache(&self) -> Option<MutexGuard<'_, InMemoryCache>> {
        self.in_memory_cache.as_ref().map(|c| c.lock().unwrap())
    }

    fn lock_database(&self) -> Option<MutexGuard<'_, Database>> {
        self.database.as_ref().map(|c| c.lock().unwrap())
    }

    #[expect(clippy::significant_drop_tightening, reason = "false positive")]
    #[instrument(name = "get_alignment", skip_all, fields(pos = %cluster_region))]
    pub fn get_or_compute_alignment(
        &self,
        reference_sequence_name: &str,
        reference_region: &GenomeRegion,
        cluster_region: GenomeRegion,
        query: AlignmentQuery,
    ) -> anyhow::Result<InProgress> {
        counter!("alignments").inc(1);
        trace!("Starting alignment for cluster");
        let key = AlignmentKey {
            reference_region: reference_region.clone(),
            alignment_ranges: query.ranges.clone(),
            query_sequence: query.sequences.query.clone(),
        };
        let reporting = |source| ResultReporting {
            failed_writer: self.failed_alignment_writer.clone(),
            stats_writer: self.stats_writer.clone(),
            source,
            reference_length: query.sequences.reference.len(),
            query_length: query.sequences.query.len(),
        };

        let cache_lock = self.lock_cache();

        if let Some(ref cache) = cache_lock {
            if let Some(sender) = cache.in_progress.get(&key) {
                // Alignment is already in progress; subscribe and report on recv.
                counter!("alignments.from_cache").inc(1);
                return Ok(InProgress::new(
                    sender.subscribe(),
                    cluster_region,
                    reporting(AlignmentSource::MemCache),
                ));
            }

            if let Some(result) = cache.finished.get(&key) {
                // Finished result in cache; report on recv.
                counter!("alignments.from_cache").inc(1);
                let (_, rx) = watch::channel(Some(result.clone()));
                return Ok(InProgress::new(
                    rx,
                    cluster_region,
                    reporting(AlignmentSource::MemCache),
                ));
            }
        }

        // Next, check if the DB has a result
        if let Some(mut db) = self.lock_database() {
            if db.needs_init() {
                db.init_with_config(&StaticAlignmentKey {
                    reference_name: reference_sequence_name,
                    aligner_config: self.aligners.describe()?,
                })?;
            }
            if let Ok(Some(result)) = tokio::task::block_in_place(|| {
                db.lookup(
                    &key,
                    self.per_alignment_settings.memory_allowance,
                    self.per_alignment_settings.timeout,
                )
            }) {
                counter!("alignments.from_db").inc(1);
                let result = Arc::new(result);
                // write to in-memory cache
                if let Some(mut cache) = cache_lock {
                    cache.finished.insert(key, result.clone());
                }
                let (_, rx) = watch::channel(Some(result));
                return Ok(InProgress::new(
                    rx,
                    cluster_region,
                    reporting(AlignmentSource::Db),
                ));
            }
        }

        let metadata = WorkerQueryMetadata {
            cluster_region: cluster_region.clone(),
        };
        let (tx, rx) = watch::channel(None);

        if let Some(mut cache) = cache_lock {
            cache.in_progress.insert(key.clone(), tx.clone());
        }

        let reporting = reporting(AlignmentSource::Computed);
        self.start_realignment_with_callback(query, metadata, key, tx);
        Ok(InProgress::new(rx, cluster_region, reporting))
    }

    fn start_realignment_with_callback(
        &self,
        query: AlignmentQuery,
        metadata: WorkerQueryMetadata,
        key: AlignmentKey,
        sender: watch::Sender<Option<Arc<TwitcherAlignment>>>,
    ) {
        let aligner = self.aligners.clone();
        let log_level = *crate::THIS_LOG_LEVEL.get_or_init(Default::default);
        let memory = self.per_alignment_settings.memory_allowance;
        let timeout = self.per_alignment_settings.timeout;

        let queue = self.parallelism.clone().acquire_owned();
        let state_mutex = self.in_memory_cache.clone();
        let db_mutex = self.database.clone();
        tokio::spawn(async move {
            let Ok(_permit) = queue.await else {
                error!(
                    "Cannot aquire concurrency permit to align {}",
                    metadata.cluster_region
                );
                return;
            };
            RUNNING.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            let worker_query = WorkerQuery {
                aligner: (&*aligner).into(),
                log_level,
                memory,
                query,
                metadata,
            };

            counter!("alignments.computations").inc(1);
            let start = Instant::now();
            let WorkerResult {
                result,
                memory: memory_usage,
            } = Self::run_alignment(&worker_query, timeout).await;
            let res = Arc::new(TwitcherAlignment::new(
                result,
                ExecStats {
                    wall: start.elapsed(),
                    memory: memory_usage,
                    memory_limit: memory,
                },
            ));

            // Move key from in_progress to finished while holding the lock so
            // there is no window where a duplicate lookup misses both maps.
            // Lock is released before the blocking DB write.
            if let Some(cache) = state_mutex.as_ref().map(|c| c.lock().unwrap()).as_mut() {
                cache.in_progress.remove(&key);
                cache.finished.insert(key.clone(), res.clone());
            }
            let _ = sender
                .send(Some(res.clone()))
                .inspect_err(|e| error!("Error sending result: {e:?}"));
            if let Some(mut db) = db_mutex.as_ref().map(|db| db.lock().unwrap())
                && let Err(e) = tokio::task::block_in_place(|| {
                    db.store(key.clone(), res.clone(), memory, timeout)
                })
            {
                error!("Can't write alignment to database: {e}");
            }
            RUNNING.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
        });
    }

    /// Run one alignment in a worker process, and observe how long it took and how much memory it
    /// used. A failure to even run the aligner is reported as an [`AlignmentFailure`], so that a
    /// single broken alignment never takes down the whole run.
    #[instrument(skip_all)]
    async fn run_alignment(wq: &WorkerQuery<'_>, timeout: Option<Duration>) -> WorkerResult {
        Self::spawn_aligner(wq, timeout)
            .await
            .unwrap_or_else(|failure| WorkerResult {
                result: Err(failure),
                memory: MemoryUsage::default(),
            })
    }

    async fn spawn_aligner(
        wq: &WorkerQuery<'_>,
        timeout: Option<Duration>,
    ) -> Result<WorkerResult, AlignmentFailure> {
        let this_exe = (*crate::THIS_EXE)
            .as_ref()
            .map_err(AlignmentFailure::error)?;
        let mut cmd = tokio::process::Command::new(this_exe);
        cmd.arg("worker");
        cmd.stdin(Stdio::piped());
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());

        let mut child = cmd.spawn()?;

        let mut stdin = child.stdin.take().context("stdin not piped")?;
        let msg = rmp_serde::to_vec(&wq).map_err(|e| {
            AlignmentFailure::error(&format!("Can't encode and write worker query: {e}"))
        })?;
        stdin.write_all(&msg).await?;
        drop(stdin);

        let stderr = child.stderr.take().context("stderr not captured")?;
        let mut stderr_writer_handle = None;
        let oom = Arc::new(AtomicBool::new(false));
        let oom2 = oom.clone();
        if let Some(mut log_w) = crate::STDERR_LOG_WRITER.get().cloned() {
            stderr_writer_handle = Some(tokio::spawn(async move {
                static WARNED_ALREADY_ABOUT_FW_DIR: AtomicBool = AtomicBool::new(false);
                // Forward any stderr lines to the stderr writer that lies on top of the progress bar
                let mut err_br = tokio::io::BufReader::new(stderr).lines();
                loop {
                    let l = match err_br.next_line().await {
                        Ok(Some(l)) => l,
                        Ok(None) => break,
                        Err(e) => {
                            error!("{e}");
                            continue;
                        }
                    };
                    if l.contains("memory allocation of") {
                        oom2.store(true, std::sync::atomic::Ordering::Relaxed);
                        continue;
                    }
                    if l.contains("Forward direction not yet supported in PreprocessedTemplateSwitchMinLengthStrategy") {
                        if WARNED_ALREADY_ABOUT_FW_DIR.load(std::sync::atomic::Ordering::Relaxed) {
                           continue;
                        }
                        WARNED_ALREADY_ABOUT_FW_DIR.store(true, std::sync::atomic::Ordering::Relaxed);
                    }

                    let _ = tokio::task::block_in_place(|| writeln!(log_w, "{l}"));
                }
            }));
        }

        // Only observable while the process is alive, so it has to be read before the kill below.
        let mut memory = MemoryUsage::default();
        let exit = if let Some(timeout) = timeout {
            match tokio::time::timeout(timeout, child.wait()).await {
                Ok(exit) => Ok(exit?),
                Err(_elapsed) => {
                    if let Some(pid) = child.id() {
                        memory = MemoryUsage::of_process(pid);
                    }
                    child.kill().await?;
                    Err(timeout)
                }
            }
        } else {
            Ok(child.wait().await?)
        };

        if let Some(h) = stderr_writer_handle {
            let _ = h.await;
        }

        let result = match exit {
            Ok(exit_status) => {
                if exit_status.success() {
                    let mut result_bytes = Vec::new();
                    let mut stdout = child.stdout.take().context("stdout not captured")?;
                    stdout.read_to_end(&mut result_bytes).await?;
                    // Deserialize the result and return it
                    let worker_result = rmp_serde::from_slice::<WorkerResult>(&result_bytes)
                        .map_err(|e| {
                            AlignmentFailure::error(&format!("Can't read result: {e:?}"))
                        })?;
                    return Ok(worker_result);
                } else if let Some(exit_code) = exit_status.code() {
                    Err(AlignmentFailure::error(&format!(
                        "Aligner exited with a non-zero exit code: {exit_code}",
                    )))
                } else if oom.load(std::sync::atomic::Ordering::Relaxed) {
                    Err(AlignmentFailure::oom())
                } else {
                    Err(AlignmentFailure::error(
                        "Aligner exited abnormally (no exit code). Perhaps it ran out of memory?",
                    ))
                }
            }
            Err(timeout) => Err(AlignmentFailure::timeout(timeout)),
        };

        Ok(WorkerResult { result, memory })
    }

    pub fn clear_cache(&self) {
        if let Some(cache) = &self.in_memory_cache {
            let mut cache = cache.lock().unwrap();
            cache.finished.clear();
        }
    }
}

fn count_result(res: &TwitcherAlignment) {
    let key = match &res.outcome {
        Ok(TwitcherAlignmentWithStatistics { alignment, .. }) if alignment.has_ts() => {
            "alignments.results.successful.with_ts"
        }
        Ok(TwitcherAlignmentWithStatistics { .. }) => "alignments.results.successful.without_ts",
        Err(AlignmentFailure::SoftFailure {
            reason: SoftFailureReason::OutOfMemory,
        }) => "alignments.results.failed.oom",
        Err(AlignmentFailure::SoftFailure {
            reason: SoftFailureReason::Timeout(_),
        }) => "alignments.results.failed.timeout",
        Err(AlignmentFailure::SoftFailure {
            reason: SoftFailureReason::Other(_),
        }) => "alignments.results.failed.other",
        Err(AlignmentFailure::Error { .. }) => "alignments.results.error",
    };
    counter!(key).inc(1);
}

impl TryFrom<&cli::CliAlignmentArgs> for AlignmentOrchestrator {
    type Error = anyhow::Error;

    fn try_from(value: &cli::CliAlignmentArgs) -> Result<Self, Self::Error> {
        let (sem, mem_per_thread) = value.init_semaphore()?;
        let (alns, costs) = value.init_aligner()?;
        let database: Option<Database> = (&value.database).try_into()?;
        let failed_alignment_writer = value
            .failed_alignments_output
            .as_deref()
            .map(|path| -> anyhow::Result<_> { Ok(Arc::new(Mutex::new(create_file(path)?))) })
            .transpose()?;
        let stats_writer = value
            .alignment_stats_output
            .as_deref()
            .map(|path| -> anyhow::Result<_> {
                Ok(Arc::new(AlignmentStatsWriter::new(
                    create_file(path)?,
                    alns.kind(),
                )))
            })
            .transpose()?;
        Ok(Self {
            aligners: alns.into(),
            costs: Arc::new(costs),
            parallelism: sem.into(),
            per_alignment_settings: PerAlignmentSettings {
                memory_allowance: mem_per_thread,
                timeout: value.aligner_timeout,
            },
            in_memory_cache: None,
            database: database.map(|db| Arc::new(Mutex::new(db))),
            failed_alignment_writer,
            stats_writer,
        })
    }
}

fn create_file(path: &str) -> anyhow::Result<Box<dyn Write + Send>> {
    Ok(Box::new(std::io::BufWriter::new(std::fs::File::create(
        path,
    )?)))
}

#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct AlignmentQuery {
    pub sequences: SequencePair,
    pub ranges: AlignmentRange,
}

impl AlignmentQuery {
    #[expect(unused)]
    pub fn visualize(&self) -> anyhow::Result<(String, String)> {
        let rs = format!(
            "{}|{}|{}",
            self.sequences
                .reference
                .get(0..self.ranges.reference_offset())
                .context("ranges oob")?
                .as_bstr(),
            self.sequences
                .reference
                .get(self.ranges.reference_range())
                .context("ranges oob")?
                .as_bstr(),
            self.sequences
                .reference
                .get(self.ranges.reference_limit()..)
                .context("ranges oob")?
                .as_bstr()
        );
        let qs = format!(
            "{}|{}|{}",
            self.sequences
                .query
                .get(0..self.ranges.query_offset())
                .context("ranges oob")?
                .as_bstr(),
            self.sequences
                .query
                .get(self.ranges.query_range())
                .context("ranges oob")?
                .as_bstr(),
            self.sequences
                .query
                .get(self.ranges.query_limit()..)
                .context("ranges oob")?
                .as_bstr()
        );
        Ok((rs, qs))
    }
}

#[derive(Deserialize, Serialize)]
#[allow(clippy::large_enum_variant)]
pub enum AlignerSelector {
    AStar {
        costs: TemplateSwitchConfig<DnaAlphabetOrN, U64Cost>,
        min_length_strategy: MLSSelector,
        allow_mixed_descendants: bool,
        no_ts: bool,
    },
    Fpa(FourPointAligner),
}

pub type AlignerSelectorDescription = Vec<u8>;

impl AlignerSelector {
    pub fn describe(&self) -> anyhow::Result<AlignerSelectorDescription> {
        Ok(rmp_serde::to_vec(self)?)
    }

    /// Which aligner this is, for reporting purposes.
    pub const fn kind(&self) -> &'static str {
        match self {
            Self::AStar { .. } => "astar",
            Self::Fpa(_) => "fpa",
        }
    }
}