mdict-rs 0.1.4

use std::path::Path;
use std::sync::Arc;

use crate::core::{DecodedKeyEntry, MdictFile};
use crate::error::{Error, Result};
use crate::types::{ContainerKind, Header, KeyOrdinal, OpenOptions};

/// Opened `.mdx` dictionary file.
#[derive(Debug)]
pub struct MdxFile {
    inner: MdictFile,
}

/// A decoded MDX key/value record.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MdxRecord {
    pub key: String,
    pub text: String,
}

/// A decoded MDX key with its stable ordinal.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MdxOrdinalKey {
    pub ordinal: KeyOrdinal,
    pub key: String,
}

/// A lazy iterator over decoded MDX keys without reading record payloads.
pub struct MdxKeyIter<'a> {
    file: &'a MdxFile,
    block_index: usize,
    entry_index: usize,
    current_block: Option<Arc<Vec<DecodedKeyEntry>>>,
}

/// A lazy iterator over decoded MDX keys with stable ordinals.
pub struct MdxOrdinalKeyIter<'a> {
    file: &'a MdxFile,
    block_index: usize,
    entry_index: usize,
    current_block: Option<Arc<Vec<DecodedKeyEntry>>>,
}

impl MdxFile {
    /// Opens an MDX file with default options.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        Self::open_with_options(path, OpenOptions::default())
    }

    /// Opens an MDX file with explicit options.
    pub fn open_with_options(path: impl AsRef<Path>, options: OpenOptions) -> Result<Self> {
        Ok(Self {
            inner: MdictFile::open(path, ContainerKind::Mdx, &options)?,
        })
    }

    /// Returns the parsed top-level header metadata.
    pub fn header(&self) -> &Header {
        &self.inner.header
    }

    /// Returns the number of dictionary entries declared by the file.
    pub fn len(&self) -> u64 {
        self.inner.len()
    }

    /// Performs a lazy lookup for a single MDX key.
    pub fn lookup(&self, key: &str) -> Result<Option<MdxRecord>> {
        let Some((block_index, entry_index, entries)) = self.inner.find_entry(key)? else {
            return Ok(None);
        };
        let entry = &entries[entry_index];
        let end = self.inner.record_end(block_index, entry_index, &entries)?;
        let bytes = self.inner.read_record_span(entry.record_start, end)?;
        let text = decode_record_text(self.inner.key_encoding, &bytes)?;
        Ok(Some(MdxRecord {
            key: entry.key.clone(),
            text,
        }))
    }

    /// Returns a lazy iterator over decoded keys without reading records.
    pub fn keys(&self) -> MdxKeyIter<'_> {
        MdxKeyIter::new(self)
    }

    /// Returns a lazy iterator over decoded keys with stable ordinals.
    pub fn keys_with_ordinals(&self) -> MdxOrdinalKeyIter<'_> {
        MdxOrdinalKeyIter::new(self)
    }

    /// Returns the original key at a stable ordinal without reading records.
    pub fn key_at(&self, ordinal: KeyOrdinal) -> Result<Option<String>> {
        self.inner.key_at_ordinal(ordinal)
    }

    /// Returns original keys for a batch of ordinals in input order.
    pub fn keys_at(&self, ordinals: &[KeyOrdinal]) -> Result<Vec<Option<String>>> {
        self.inner.keys_at_ordinals(ordinals)
    }

    pub fn entries(&self) -> MdxEntryIter<'_> {
        MdxEntryIter::new(self)
    }
}

impl<'a> MdxKeyIter<'a> {
    fn new(file: &'a MdxFile) -> Self {
        Self {
            file,
            block_index: 0,
            entry_index: 0,
            current_block: None,
        }
    }
}

impl<'a> MdxOrdinalKeyIter<'a> {
    fn new(file: &'a MdxFile) -> Self {
        Self {
            file,
            block_index: 0,
            entry_index: 0,
            current_block: None,
        }
    }
}

impl Iterator for MdxKeyIter<'_> {
    type Item = Result<String>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.block_index >= self.file.inner.key_block_count() {
                return None;
            }

            if self.current_block.is_none() {
                match self.file.inner.decode_key_block(self.block_index) {
                    Ok(entries) => {
                        self.current_block = Some(entries);
                        self.entry_index = 0;
                    }
                    Err(error) => return Some(Err(error)),
                }
            }

            let entries = self.current_block.as_ref().unwrap();
            if self.entry_index >= entries.len() {
                self.block_index += 1;
                self.current_block = None;
                continue;
            }

            let key = entries[self.entry_index].key.clone();
            self.entry_index += 1;
            return Some(Ok(key));
        }
    }
}

impl Iterator for MdxOrdinalKeyIter<'_> {
    type Item = Result<MdxOrdinalKey>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.block_index >= self.file.inner.key_block_count() {
                return None;
            }

            if self.current_block.is_none() {
                match self.file.inner.decode_key_block(self.block_index) {
                    Ok(entries) => {
                        self.current_block = Some(entries);
                        self.entry_index = 0;
                    }
                    Err(error) => return Some(Err(error)),
                }
            }

            let entries = self.current_block.as_ref().unwrap();
            if self.entry_index >= entries.len() {
                self.block_index += 1;
                self.current_block = None;
                continue;
            }

            let block = &self.file.inner.key_index.blocks[self.block_index];
            let result = (|| -> Result<MdxOrdinalKey> {
                let ordinal = block
                    .entry_start_index
                    .checked_add(self.entry_index as u64)
                    .ok_or(Error::InvalidFormat("key ordinal overflow"))?;
                Ok(MdxOrdinalKey {
                    ordinal: KeyOrdinal::from(ordinal),
                    key: entries[self.entry_index].key.clone(),
                })
            })();
            self.entry_index += 1;
            return Some(result);
        }
    }
}

/// Lazy iterator over decoded MDX entries.
pub struct MdxEntryIter<'a> {
    file: &'a MdxFile,
    block_index: usize,
    entry_index: usize,
    current_block: Option<Arc<Vec<DecodedKeyEntry>>>,
}

impl<'a> MdxEntryIter<'a> {
    fn new(file: &'a MdxFile) -> Self {
        Self {
            file,
            block_index: 0,
            entry_index: 0,
            current_block: None,
        }
    }
}

impl Iterator for MdxEntryIter<'_> {
    type Item = Result<MdxRecord>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.block_index >= self.file.inner.key_block_count() {
                return None;
            }

            if self.current_block.is_none() {
                match self.file.inner.decode_key_block(self.block_index) {
                    Ok(entries) => {
                        self.current_block = Some(entries);
                        self.entry_index = 0;
                    }
                    Err(error) => return Some(Err(error)),
                }
            }

            let entries = self.current_block.as_ref().unwrap();
            if self.entry_index >= entries.len() {
                self.block_index += 1;
                self.current_block = None;
                continue;
            }

            let entry = &entries[self.entry_index];
            let result = (|| -> Result<MdxRecord> {
                let end =
                    self.file
                        .inner
                        .record_end(self.block_index, self.entry_index, entries)?;
                let bytes = self.file.inner.read_record_span(entry.record_start, end)?;
                let text = decode_record_text(self.file.inner.key_encoding, &bytes)?;
                Ok(MdxRecord {
                    key: entry.key.clone(),
                    text,
                })
            })();
            self.entry_index += 1;
            return Some(result);
        }
    }
}

fn decode_record_text(encoding: crate::encoding::TextEncoding, bytes: &[u8]) -> Result<String> {
    let text = encoding.decode(bytes, "mdx record")?;
    Ok(text.trim_end_matches('\0').to_owned())
}