affine_importer 0.1.2

use std::sync::Arc;

use crate::{
  ImportBatch, ImportBatchLimits, ImportCursor, ImportError, ImportOptions, ImportProgress, ImportRegistry,
  ImportResult, ImportSource,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImportSessionState {
  Ready,
  Drained,
  Cancelled,
}

#[derive(Debug, Clone)]
pub struct ImportSessionOptions {
  pub format: String,
  pub source: ImportSource,
  pub import_options: ImportOptions,
  pub batch_limits: ImportBatchLimits,
}

pub struct ImportSession {
  cursor: Option<Box<dyn ImportCursor>>,
  progress: ImportProgress,
  state: ImportSessionState,
}

impl ImportSession {
  pub fn create(options: ImportSessionOptions, registry: Arc<ImportRegistry>) -> ImportResult<Self> {
    let cursor = registry.create_cursor(
      &options.format,
      options.source,
      options.import_options,
      options.batch_limits,
    )?;
    let progress = ImportProgress {
      completed: 0,
      total: cursor.total(),
    };
    Ok(Self {
      cursor: Some(cursor),
      progress,
      state: ImportSessionState::Ready,
    })
  }

  pub fn state(&self) -> ImportSessionState {
    self.state
  }

  pub fn progress(&self) -> &ImportProgress {
    &self.progress
  }

  pub fn cancel(&mut self) {
    self.cursor = None;
    self.state = ImportSessionState::Cancelled;
  }

  pub fn next_batch(&mut self) -> ImportResult<Option<ImportBatch>> {
    if self.state == ImportSessionState::Cancelled {
      return Err(ImportError::Cancelled);
    }
    let Some(cursor) = self.cursor.as_mut() else {
      self.state = ImportSessionState::Drained;
      return Ok(None);
    };
    let batch = cursor.next_batch()?;
    if let Some(batch) = batch.as_ref() {
      self.progress = batch.progress.clone();
    } else {
      self.state = ImportSessionState::Drained;
      self.cursor = None;
    }
    Ok(batch)
  }
}

#[cfg(test)]
mod tests {
  use std::{
    io::Write,
    path::PathBuf,
    sync::atomic::{AtomicU64, Ordering},
  };

  use zip::{ZipWriter, write::SimpleFileOptions};

  use super::*;
  use crate::ImportError;

  static NEXT_TEST_ARCHIVE_ID: AtomicU64 = AtomicU64::new(1);

  fn archive_path(entries: &[(&str, &[u8])]) -> PathBuf {
    let path = std::env::temp_dir().join(format!(
      "affine-importer-{}-{}.zip",
      std::process::id(),
      NEXT_TEST_ARCHIVE_ID.fetch_add(1, Ordering::Relaxed)
    ));
    let file = std::fs::File::create(&path).unwrap();
    let mut writer = ZipWriter::new(file);
    for (path, bytes) in entries {
      writer.start_file(path, SimpleFileOptions::default()).unwrap();
      writer.write_all(bytes).unwrap();
    }
    writer.finish().unwrap();
    path
  }

  #[test]
  fn session_reports_progress_and_drains_in_batches() {
    let path = archive_path(&[("a.md", b"a"), ("b.md", b"b")]);
    let mut session = ImportSession::create(
      ImportSessionOptions {
        format: "markdownZip".to_string(),
        source: ImportSource::FilePath(path.clone()),
        import_options: ImportOptions::default(),
        batch_limits: ImportBatchLimits {
          max_docs: 1,
          max_blobs: 10,
          max_blob_bytes: u64::MAX,
        },
      },
      Arc::new(ImportRegistry::with_builtin()),
    )
    .unwrap();

    assert_eq!(session.state(), ImportSessionState::Ready);
    assert_eq!(session.progress().completed, 0);
    let first = session.next_batch().unwrap().unwrap();
    assert_eq!(first.docs.len(), 1);
    assert_eq!(first.progress.completed, 1);
    assert!(!first.done);
    let second = session.next_batch().unwrap().unwrap();
    assert_eq!(second.docs.len(), 1);
    assert_eq!(second.progress.completed, 2);
    assert!(second.done);
    assert!(session.next_batch().unwrap().is_none());
    assert_eq!(session.state(), ImportSessionState::Drained);
    let _ = std::fs::remove_file(path);
  }

  #[test]
  fn cancelled_session_rejects_next_batch() {
    let path = archive_path(&[("a.md", b"a")]);
    let mut session = ImportSession::create(
      ImportSessionOptions {
        format: "markdownZip".to_string(),
        source: ImportSource::FilePath(path.clone()),
        import_options: ImportOptions::default(),
        batch_limits: ImportBatchLimits::default(),
      },
      Arc::new(ImportRegistry::with_builtin()),
    )
    .unwrap();
    session.cancel();
    assert!(matches!(session.next_batch(), Err(ImportError::Cancelled)));
    let _ = std::fs::remove_file(path);
  }
}