use std::io::{self, Read};
use std::mem::ManuallyDrop;
use std::sync::{Arc, Mutex};

use crate::mimetype::MimeTypeGuesser;
use crate::xopen::open_raw;

pub const MAX_K: usize = 31;
const PAGE_SIZE: usize = 65536;
// overlap (MAX_K - 1) + page data (PAGE_SIZE) + 1 byte for the end-of-page terminating 0
const BUF_SIZE: usize = MAX_K + PAGE_SIZE;

// ─── OverlapState ─────────────────────────────────────────────────────────────

pub(crate) struct OverlapState {
    data: [u8; MAX_K],
    len: usize,
    k: usize,
}

impl OverlapState {
    pub(crate) fn new(k: usize) -> Self {
        assert!(k > 0 && k <= MAX_K);
        Self {
            data: [0u8; MAX_K],
            len: 0,
            k,
        }
    }
}

// ─── NucParser trait ──────────────────────────────────────────────────────────

// Transforms a raw page into a compacted nucleotide stream in-place.
//
// Buffer layout on each call:
//   buf[0..overlap_len()]           — overlap bytes copied by write_overlap()
//   buf[overlap_len()..overlap_len()+n] — raw bytes just read from the source
//
// Returns the number of output bytes in buf[0..returned].
pub(crate) trait NucParser {
    // required: format-specific
    fn new(k: usize) -> Self
    where
        Self: Sized;
    fn overlap_state(&self) -> &OverlapState;
    fn overlap_state_mut(&mut self) -> &mut OverlapState;
    fn is_in_seq(&self) -> bool;
    fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize;

    // provided: format-independent overlap management
    fn overlap_len(&self) -> usize {
        self.overlap_state().len
    }

    fn write_overlap(&self, buf: &mut [u8]) {
        let ol = &self.overlap_state();
        buf[..ol.len].copy_from_slice(&ol.data[..ol.len]);
    }

    // Called at end of parse_inplace: saves overlap state and returns adjusted j.
    // seq_start is the j-position where the last sequence started in this call's output.
    fn save_overlap(&mut self, buf: &mut [u8], j: usize, seq_start: usize) -> usize {
        if !self.is_in_seq() {
            self.overlap_state_mut().len = 0;
            return j;
        }
        let seq_len = j - seq_start;
        let k = self.overlap_state().k;
        if seq_len >= k {
            // Sequence long enough: save last k-1 nucleotides, terminate with 0.
            let ol = k - 1;
            self.overlap_state_mut().data[..ol].copy_from_slice(&buf[j - ol..j]);
            self.overlap_state_mut().len = ol;
            // SAFETY: j <= total - 1 < BUF_SIZE = buf.len()
            // (total = overlap_len + n <= (MAX_K-1) + PAGE_SIZE = BUF_SIZE - 1)
            unsafe {
                *buf.get_unchecked_mut(j) = 0;
            }
            j + 1
        } else if seq_len > 0 {
            // Short sequence (< k): save whole fragment, strip from output.
            self.overlap_state_mut().data[..seq_len].copy_from_slice(&buf[seq_start..j]);
            self.overlap_state_mut().len = seq_len;
            seq_start
        } else {
            self.overlap_state_mut().len = 0;
            j
        }
    }
}

// ─── FASTA parser ─────────────────────────────────────────────────────────────

#[derive(Clone, Copy)]
enum FastaState {
    OutSeq,
    InTitle,
    InSeq,
    InAmbiguous,
}

pub(crate) struct FastaParser {
    state: FastaState,
    overlap: OverlapState,
}

impl NucParser for FastaParser {
    fn new(k: usize) -> Self {
        Self {
            state: FastaState::OutSeq,
            overlap: OverlapState::new(k),
        }
    }

    #[inline]
    fn overlap_state(&self) -> &OverlapState {
        &self.overlap
    }

    #[inline]
    fn overlap_state_mut(&mut self) -> &mut OverlapState {
        &mut self.overlap
    }

    #[inline]
    fn is_in_seq(&self) -> bool {
        matches!(self.state, FastaState::InSeq)
    }

    fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize {
        let total = self.overlap.len + n;
        let mut i = 0; // read index
        let mut j = 0; // write index (invariant: j <= i always)
        // j-position where the current sequence started in this call's output;
        // meaningful only when state is InSeq.
        let mut seq_start: usize = 0;

        while i < total {
            // SAFETY: i < total <= BUF_SIZE = buf.len()
            let byte = unsafe { *buf.get_unchecked(i) };

            match self.state {
                FastaState::OutSeq => {
                    if byte == b'>' {
                        self.state = FastaState::InTitle;
                    }
                    i += 1;
                }
                FastaState::InTitle => {
                    if byte == b'\n' || byte == b'\r' {
                        self.state = FastaState::InSeq;
                        seq_start = j;
                    }
                    i += 1;
                }
                FastaState::InSeq => {
                    if byte == b'\n' || byte == b'\r' {
                        i += 1;
                        continue;
                    }
                    let nuc = byte & 0xDF; // to uppercase
                    if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
                        // SAFETY: j <= i < total <= BUF_SIZE = buf.len()
                        unsafe {
                            *buf.get_unchecked_mut(j) = nuc;
                        }
                        j += 1;
                        i += 1;
                    } else if byte == b'>' {
                        if j > seq_start {
                            unsafe {
                                *buf.get_unchecked_mut(j) = 0;
                            }
                            j += 1;
                        }
                        self.state = FastaState::InTitle;
                        i += 1;
                    } else {
                        // first ambiguous base: end current sequence if non-empty
                        if j > seq_start {
                            unsafe {
                                *buf.get_unchecked_mut(j) = 0;
                            }
                            j += 1;
                        }
                        self.state = FastaState::InAmbiguous;
                        i += 1;
                    }
                }
                FastaState::InAmbiguous => {
                    if byte == b'\n' || byte == b'\r' {
                        i += 1;
                        continue;
                    }
                    if byte == b'>' {
                        self.state = FastaState::InTitle;
                        i += 1;
                        continue;
                    }
                    let nuc = byte & 0xDF;
                    if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
                        seq_start = j;
                        // SAFETY: j <= i < total <= BUF_SIZE = buf.len()
                        unsafe {
                            *buf.get_unchecked_mut(j) = nuc;
                        }
                        j += 1;
                        self.state = FastaState::InSeq;
                    }
                    i += 1;
                }
            }
        }

        self.save_overlap(buf, j, seq_start)
    }
}

// ─── FASTQ parser ─────────────────────────────────────────────────────────────

#[derive(Clone, Copy)]
enum FastqState {
    OutSeq,
    InTitle,
    InSeq,
    InAmbiguous,
    InQualTitle,
    InQual,
}

pub(crate) struct FastqParser {
    state: FastqState,
    overlap: OverlapState,
}

impl NucParser for FastqParser {
    fn new(k: usize) -> Self {
        Self {
            state: FastqState::OutSeq,
            overlap: OverlapState::new(k),
        }
    }

    #[inline]
    fn overlap_state(&self) -> &OverlapState {
        &self.overlap
    }

    #[inline]
    fn overlap_state_mut(&mut self) -> &mut OverlapState {
        &mut self.overlap
    }

    #[inline]
    fn is_in_seq(&self) -> bool {
        matches!(self.state, FastqState::InSeq)
    }

    fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize {
        let total = self.overlap.len + n;
        let mut i = 0;
        let mut j = 0;
        let mut seq_start: usize = 0;

        while i < total {
            // SAFETY: i < total <= BUF_SIZE = buf.len()
            let byte = unsafe { *buf.get_unchecked(i) };

            match self.state {
                FastqState::OutSeq => {
                    if byte == b'@' {
                        self.state = FastqState::InTitle;
                    }
                    i += 1;
                }
                FastqState::InTitle => {
                    if byte == b'\n' || byte == b'\r' {
                        self.state = FastqState::InSeq;
                        seq_start = j;
                    }
                    i += 1;
                }
                FastqState::InSeq => {
                    if byte == b'\n' || byte == b'\r' {
                        if j > seq_start {
                            unsafe {
                                *buf.get_unchecked_mut(j) = 0;
                            }
                            j += 1;
                        }
                        self.state = FastqState::InQualTitle;
                        i += 1;
                        continue;
                    }
                    let nuc = byte & 0xDF;
                    if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
                        // SAFETY: j <= i < total <= BUF_SIZE = buf.len()
                        unsafe {
                            *buf.get_unchecked_mut(j) = nuc;
                        }
                        j += 1;
                    } else {
                        if j > seq_start {
                            unsafe {
                                *buf.get_unchecked_mut(j) = 0;
                            }
                            j += 1;
                        }
                        self.state = FastqState::InAmbiguous;
                    }
                    i += 1;
                }
                FastqState::InAmbiguous => {
                    if byte == b'\n' || byte == b'\r' {
                        self.state = FastqState::InQualTitle;
                        i += 1;
                        continue;
                    }
                    let nuc = byte & 0xDF;
                    if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
                        seq_start = j;
                        // SAFETY: j <= i < total <= BUF_SIZE = buf.len()
                        unsafe {
                            *buf.get_unchecked_mut(j) = nuc;
                        }
                        j += 1;
                        self.state = FastqState::InSeq;
                    }
                    i += 1;
                }
                FastqState::InQualTitle => {
                    if byte == b'\n' || byte == b'\r' {
                        self.state = FastqState::InQual;
                    }
                    i += 1;
                }
                FastqState::InQual => {
                    if byte == b'\n' || byte == b'\r' {
                        self.state = FastqState::OutSeq;
                    }
                    i += 1;
                }
            }
        }

        self.save_overlap(buf, j, seq_start)
    }
}

// ─── GenBank parser ───────────────────────────────────────────────────────────

const ORIGIN_TAIL: &[u8] = b"RIGIN";

#[derive(Clone, Copy)]
enum GenbankState {
    OutSeq,
    MatchOrigin,
    SkipOriginLine,
    InSeq,
    InSlash,
    InAmbiguous,
}

pub(crate) struct GenbankParser {
    state: GenbankState,
    overlap: OverlapState,
    keyword_pos: usize,
    at_line_start: bool,
}

impl NucParser for GenbankParser {
    fn new(k: usize) -> Self {
        Self {
            state: GenbankState::OutSeq,
            overlap: OverlapState::new(k),
            keyword_pos: 0,
            at_line_start: true,
        }
    }

    #[inline]
    fn overlap_state(&self) -> &OverlapState {
        &self.overlap
    }

    #[inline]
    fn overlap_state_mut(&mut self) -> &mut OverlapState {
        &mut self.overlap
    }

    #[inline]
    fn is_in_seq(&self) -> bool {
        matches!(self.state, GenbankState::InSeq)
    }

    fn parse_inplace(&mut self, buf: &mut [u8], n: usize) -> usize {
        let total = self.overlap.len + n;
        let mut i = 0;
        let mut j = 0;
        let mut seq_start: usize = 0;

        while i < total {
            // SAFETY: i < total <= BUF_SIZE = buf.len()
            let byte = unsafe { *buf.get_unchecked(i) };

            match self.state {
                GenbankState::OutSeq => {
                    if byte == b'\n' || byte == b'\r' {
                        self.at_line_start = true;
                    } else if self.at_line_start && byte == b'O' {
                        self.state = GenbankState::MatchOrigin;
                        self.keyword_pos = 1;
                        self.at_line_start = false;
                    } else {
                        self.at_line_start = false;
                    }
                    i += 1;
                }
                GenbankState::MatchOrigin => {
                    if byte == b'\n' || byte == b'\r' {
                        self.state = GenbankState::OutSeq;
                        self.at_line_start = true;
                    } else if byte == ORIGIN_TAIL[self.keyword_pos - 1] {
                        self.keyword_pos += 1;
                        if self.keyword_pos == 6 {
                            self.state = GenbankState::SkipOriginLine;
                        }
                    } else {
                        self.state = GenbankState::OutSeq;
                        self.at_line_start = false;
                    }
                    i += 1;
                }
                GenbankState::SkipOriginLine => {
                    if byte == b'\n' || byte == b'\r' {
                        self.state = GenbankState::InSeq;
                        seq_start = j;
                    }
                    i += 1;
                }
                GenbankState::InSeq => {
                    if byte == b'\n' || byte == b'\r' {
                        self.at_line_start = true;
                        i += 1;
                        continue;
                    }
                    if self.at_line_start && byte == b'/' {
                        self.state = GenbankState::InSlash;
                        self.at_line_start = false;
                        i += 1;
                        continue;
                    }
                    self.at_line_start = false;
                    let nuc = byte & 0xDF;
                    if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
                        // SAFETY: j <= i < total <= BUF_SIZE = buf.len()
                        unsafe {
                            *buf.get_unchecked_mut(j) = nuc;
                        }
                        j += 1;
                    } else if byte.is_ascii_digit() || byte == b' ' {
                        // position numbers and spacing between groups: skip
                    } else {
                        // ambiguous base: end current sequence if non-empty
                        if j > seq_start {
                            unsafe {
                                *buf.get_unchecked_mut(j) = 0;
                            }
                            j += 1;
                        }
                        self.state = GenbankState::InAmbiguous;
                    }
                    i += 1;
                }
                GenbankState::InSlash => {
                    if byte == b'/' {
                        // confirmed "//": end of sequence record
                        if j > seq_start {
                            unsafe {
                                *buf.get_unchecked_mut(j) = 0;
                            }
                            j += 1;
                        }
                        self.state = GenbankState::OutSeq;
                        self.at_line_start = false;
                    } else if byte == b'\n' || byte == b'\r' {
                        // single '/' line: back to sequence
                        self.state = GenbankState::InSeq;
                        self.at_line_start = true;
                    } else {
                        // false positive: single '/' mid-line, resume sequence
                        self.state = GenbankState::InSeq;
                        self.at_line_start = false;
                    }
                    i += 1;
                }
                GenbankState::InAmbiguous => {
                    if byte == b'\n' || byte == b'\r' {
                        self.at_line_start = true;
                        i += 1;
                        continue;
                    }
                    if self.at_line_start && byte == b'/' {
                        self.state = GenbankState::InSlash;
                        self.at_line_start = false;
                        i += 1;
                        continue;
                    }
                    self.at_line_start = false;
                    let nuc = byte & 0xDF;
                    if nuc == b'A' || nuc == b'C' || nuc == b'G' || nuc == b'T' {
                        seq_start = j;
                        // SAFETY: j <= i < total <= BUF_SIZE = buf.len()
                        unsafe {
                            *buf.get_unchecked_mut(j) = nuc;
                        }
                        j += 1;
                        self.state = GenbankState::InSeq;
                    }
                    // digits, spaces, other ambiguous codes: skip
                    i += 1;
                }
            }
        }

        self.save_overlap(buf, j, seq_start)
    }
}

// ─── NucPage ──────────────────────────────────────────────────────────────────

/// Owned page of compacted nucleotides: uppercase A/C/G/T bytes separated by `0`
/// at sequence boundaries. Automatically returns its buffer to the pool on drop.
pub struct NucPage {
    data: ManuallyDrop<Vec<u8>>,
    len: usize,
    pool: Arc<Mutex<Vec<Vec<u8>>>>,
}

impl std::ops::Deref for NucPage {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        &self.data[..self.len]
    }
}

impl Drop for NucPage {
    fn drop(&mut self) {
        // SAFETY: data is never accessed after this point
        let buf = unsafe { ManuallyDrop::take(&mut self.data) };
        self.pool.lock().unwrap().push(buf);
    }
}

// ─── NucPageCursor ────────────────────────────────────────────────────────────

/// A forward cursor over the normalised bytes of a [`NucPage`].
///
/// Provides the `next_byte` / `rewind` interface consumed by
/// [`obiskbuilder::SuperKmerStreamIter`].
pub struct NucPageCursor<'a> {
    data: &'a [u8],
    pos: usize,
}

impl NucPageCursor<'_> {
    /// Returns the next byte in the page, or `None` at end.
    #[inline]
    pub fn next_byte(&mut self) -> Option<u8> {
        if self.pos < self.data.len() {
            let b = self.data[self.pos];
            self.pos += 1;
            Some(b)
        } else {
            None
        }
    }

    /// Steps the cursor back by `n` bytes.
    ///
    /// The caller guarantees that the last `n` bytes were all `ACGT`
    /// (no `0x00` separators), so they are still in the page buffer.
    #[inline]
    pub fn rewind(&mut self, n: usize) {
        self.pos -= n;
    }

    /// Total number of bytes in the underlying page.
    #[inline]
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Returns `true` if the page contains no bytes.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }
}

impl NucPage {
    /// Creates a forward cursor positioned at the start of this page.
    pub fn cursor(&self) -> NucPageCursor<'_> {
        NucPageCursor { data: self, pos: 0 }
    }
}

// ─── NucStream ────────────────────────────────────────────────────────────────

pub(crate) struct NucStream<R: Read, P: NucParser> {
    reader: R,
    parser: P,
    pool: Arc<Mutex<Vec<Vec<u8>>>>,
    eof: bool,
}

impl<R: Read, P: NucParser> NucStream<R, P> {
    pub(crate) fn new(reader: R, k: usize) -> Self {
        Self {
            reader,
            parser: P::new(k),
            pool: Arc::new(Mutex::new(Vec::new())),
            eof: false,
        }
    }

    pub(crate) fn read_page(&mut self) -> Option<NucPage> {
        loop {
            if self.eof {
                return None;
            }
            // take a buffer from the pool, or allocate fresh if all are in-flight
            let mut buf = self
                .pool
                .lock()
                .unwrap()
                .pop()
                .unwrap_or_else(|| vec![0u8; BUF_SIZE]);

            let ol = self.parser.overlap_len();
            self.parser.write_overlap(&mut buf[..ol]);
            let n = self.reader.read(&mut buf[ol..ol + PAGE_SIZE]).unwrap_or(0);
            if n == 0 {
                self.eof = true;
                if ol == 0 {
                    self.pool.lock().unwrap().push(buf);
                    return None;
                }
            }
            let out_len = self.parser.parse_inplace(&mut buf, n);
            if out_len > 0 {
                return Some(NucPage {
                    data: ManuallyDrop::new(buf),
                    len: out_len,
                    pool: Arc::clone(&self.pool),
                });
            }
            // empty page (all headers/ambiguous): return buf to pool and loop
            self.pool.lock().unwrap().push(buf);
        }
    }
}

impl<R: Read, P: NucParser> Iterator for NucStream<R, P> {
    type Item = NucPage;
    fn next(&mut self) -> Option<NucPage> {
        self.read_page()
    }
}

// ─── FastaNucStream ───────────────────────────────────────────────────────────

pub(crate) type FastaNucStream<R> = NucStream<R, FastaParser>;
pub(crate) type FastqNucStream<R> = NucStream<R, FastqParser>;
pub(crate) type GenbankNucStream<R> = NucStream<R, GenbankParser>;

// ─── AnyNucStream ─────────────────────────────────────────────────────────────

pub(crate) enum AnyNucStream<R: Read> {
    Fasta(FastaNucStream<R>),
    Fastq(FastqNucStream<R>),
    Genbank(GenbankNucStream<R>),
}

impl<R: Read> Iterator for AnyNucStream<R> {
    type Item = NucPage;
    fn next(&mut self) -> Option<NucPage> {
        match self {
            AnyNucStream::Fasta(s) => s.next(),
            AnyNucStream::Fastq(s) => s.next(),
            AnyNucStream::Genbank(s) => s.next(),
        }
    }
}

fn dispatch<R: Read>(
    mut guesser: MimeTypeGuesser<R>,
    k: usize,
) -> Option<AnyNucStream<MimeTypeGuesser<R>>> {
    match guesser.mime_type() {
        Some("text/fasta") => Some(AnyNucStream::Fasta(NucStream::new(guesser, k))),
        Some("text/fastq") => Some(AnyNucStream::Fastq(NucStream::new(guesser, k))),
        Some("text/gbff") => Some(AnyNucStream::Genbank(NucStream::new(guesser, k))),
        _ => None,
    }
}

/// Wraps an already-open reader in a nucleotide stream, detecting its format.
/// Returns `None` if the format is not recognised.
pub(crate) fn nuc_stream<R: Read>(reader: R, k: usize) -> Option<AnyNucStream<MimeTypeGuesser<R>>> {
    dispatch(MimeTypeGuesser::new(reader), k)
}

/// Opens a nucleotide stream from any source (file path, URL, or `-` for stdin),
/// with transparent decompression and automatic format detection.
///
/// # Errors
/// Returns an `io::Error` if the source cannot be opened, decompression fails,
/// or the format is not recognised.
pub fn open_nuc_stream(
    source: &str,
    k: usize,
) -> io::Result<Box<dyn Iterator<Item = NucPage> + Send>> {
    let reader = open_raw(source)?;
    nuc_stream(reader, k)
        .map(|s| Box::new(s) as Box<dyn Iterator<Item = NucPage> + Send>)
        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "unknown sequence format"))
}

#[cfg(test)]
#[path = "tests/nucstream.rs"]
mod tests;
