use needletail::{FastxReader, parse_fastx_file, parse_fastx_stdin};
use sassy::CachedRev;
use std::path::{Path, PathBuf};
use std::sync::Mutex; //Todo: could use parking_lot mutex - faster

/// Each batch of text records will be at most this size if possible.
const DEFAULT_BATCH_BYTES: usize = 1024 * 1024; // 1 MiB

/// Type alias for fasta record IDs.
pub type ID = String;

/// A search pattern, with ID from fasta file.
#[derive(Clone, Debug)]
pub struct PatternRecord {
    pub id: ID,
    pub seq: Vec<u8>,
}

/// A text to be searched, with ID from fasta file.
/// TODO: Reduce the number of allocations here.
#[derive(Debug)]
pub struct TextRecord {
    pub id: ID,
    pub seq: CachedRev<Vec<u8>>,
    pub quality: Vec<u8>,
}

/// A batch of alignment tasks, with total text size around `DEFAULT_BATCH_BYTES`.
/// This avoids lock contention of sending too small items across threads.
///
/// Each task represents searching _all_ patterns against _all_ text records.
pub type TaskBatch<'a> = (&'a Path, &'a [PatternRecord], Vec<TextRecord>);

struct RecordState {
    /// The current fasta reader.
    reader: Box<dyn FastxReader + Send>,
    /// The next file index.
    cur_file_index: usize,
    /// The next batch id.
    next_batch_id: usize,
    /// Current text record, that can be send to multiple threads.
    current_record: Option<TextRecord>,
}

/// Thread-safe iterator giving *batches* of (pattern, text) pairs.
/// Each batch searches at least `batch_byte_limit` bytes of text.
///
/// Created using `TaskIterator::new` from a list of patterns and a path to a Fasta file to be searched.
pub struct InputIterator<'a> {
    patterns: &'a [PatternRecord],
    paths: &'a Vec<PathBuf>,
    state: Mutex<RecordState>,
    batch_byte_limit: usize,
    rev: bool,
}

fn parse_file(path: &PathBuf) -> Box<dyn FastxReader> {
    if path == Path::new("") || path == Path::new("-") {
        return parse_fastx_stdin().unwrap();
    } else {
        return parse_fastx_file(path).unwrap();
    }
}

impl<'a> InputIterator<'a> {
    /// Create a new iterator over `fasta_path`, going through `patterns`.
    /// `max_batch_bytes` controls how many texts are bundled together.
    pub fn new(
        paths: &'a Vec<PathBuf>,
        patterns: &'a [PatternRecord],
        max_batch_bytes: Option<usize>,
        rev: bool,
    ) -> Self {
        let reader = parse_file(&paths[0]);
        // Just empty state when we create the iterator
        let state = RecordState {
            reader,
            cur_file_index: 0,
            next_batch_id: 0,
            current_record: None,
        };
        Self {
            patterns,
            paths,
            state: Mutex::new(state),
            batch_byte_limit: max_batch_bytes.unwrap_or(DEFAULT_BATCH_BYTES),
            rev,
        }
    }

    /// Get the next batch, or returns None when done.
    pub fn next_batch(&self) -> Option<(usize, TaskBatch<'a>)> {
        let mut state = self.state.lock().unwrap();
        let batch_id = state.next_batch_id;
        state.next_batch_id += 1;
        let mut batch: TaskBatch<'a> =
            (&self.paths[state.cur_file_index], self.patterns, Vec::new());
        let mut bytes_in_batch = 0usize;

        // Effectively this gets a record, add all patterns, then tries
        // to push another text record, if possible. This way texts
        // are only 'read' from the Fasta file once.

        loop {
            // Make sure we have a current record, just so we can unwrap
            loop {
                match state.reader.next() {
                    Some(Ok(rec)) => {
                        let id = String::from_utf8(rec.id().to_vec()).unwrap().to_string();
                        let seq = rec.seq().into_owned();
                        let static_text = CachedRev::new(seq, self.rev);
                        state.current_record = Some(TextRecord {
                            id,
                            seq: static_text,
                            quality: rec.qual().unwrap_or(&[]).to_vec(),
                        });
                        break;
                    }
                    Some(Err(e)) => panic!("Error reading FASTA record: {e}"),
                    None => {
                        // Reached end of reader, initialize for next file.
                        let end_of_files = state.cur_file_index + 1 >= self.paths.len();
                        if !end_of_files {
                            state.cur_file_index += 1;
                            if state.cur_file_index < self.paths.len() {
                                state.reader = parse_file(&self.paths[state.cur_file_index]);
                            }
                        }

                        // Return last batch for the current file.
                        if !batch.2.is_empty() {
                            return Some((batch_id, batch));
                        }
                        if end_of_files {
                            return None;
                        }

                        // Start reading next file.
                        continue;
                    }
                }
            }

            let current_record = &mut state.current_record;

            // We get the ref to the current record we have available
            let record_len = current_record.as_ref().unwrap().seq.text.len();

            // If no space left for next record, we return current batch
            if !batch.2.is_empty()
                && bytes_in_batch + record_len * self.patterns.len() > self.batch_byte_limit
            {
                break; // return current batch, keep state for next call
            }

            // Add next pattern
            batch.2.push(current_record.take().unwrap());
            bytes_in_batch += record_len * self.patterns.len();
        }

        Some((batch_id, batch))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::Rng;
    use std::io::Write;
    use tempfile::NamedTempFile;

    fn random_dan_seq(len: usize) -> Vec<u8> {
        let mut rng = rand::rng();
        let mut seq = Vec::new();
        let bases = b"ACGT";
        for _ in 0..len {
            seq.push(bases[rng.random_range(0..bases.len())]);
        }
        seq
    }

    #[test]
    fn test_record_iterator() {
        // Create 100 different random sequences with length of 100-1000
        let mut rng = rand::rng();
        let mut seqs = Vec::new();
        for _ in 0..100 {
            seqs.push(random_dan_seq(rng.random_range(100..1000)));
        }

        // Create a temporary file to write the fasta file to
        let mut file = NamedTempFile::new().unwrap();
        for (i, seq) in seqs.into_iter().enumerate() {
            file.write_all(format!(">seq_{}\n{}\n", i, String::from_utf8(seq).unwrap()).as_bytes())
                .unwrap();
        }
        file.flush().unwrap();

        // Create 10 different random patterns
        let mut patterns = Vec::new();
        for i in 0..10 {
            patterns.push(PatternRecord {
                id: format!("pattern_{}", i),
                seq: random_dan_seq(rng.random_range(250..1000)),
            });
        }

        // Create the iterator
        let paths = vec![file.path().to_path_buf()];
        let iter = InputIterator::new(&paths, &patterns, Some(500), true);

        // Pull 10 batches
        let mut batch_id = 0;
        while let Some(batch) = iter.next_batch() {
            batch_id += 1;
            // Get unique texts, and then their length sum
            let unique_texts = batch
                .1
                .2
                .iter()
                .map(|item| item.seq.text.clone())
                .collect::<std::collections::HashSet<_>>();
            let text_len = unique_texts.iter().map(|text| text.len()).sum::<usize>();
            let n_patterns = batch
                .1
                .1
                .iter()
                .map(|item| item.id.clone())
                .collect::<std::collections::HashSet<_>>()
                .len();
            let n_texts = unique_texts.len();
            println!(
                "Batch {batch_id} (tot_size: {text_len}, n_texts: {n_texts}): {n_patterns} patterns"
            );
        }
        drop(file);
    }
}
