RustAI_Full/src/pdf_import.rs

use crate::training::websearch::derive_samples_from_pdf_text;
use anyhow::{anyhow, Context, Result};
use lopdf::Document;
use ollama_rs::Ollama;
use pdf_extract::extract_text;
use serde_json;
use std::fs::{self, File};
use std::io::{self, BufWriter, Write};
use std::path::{Path, PathBuf};

/// Scan the given directory and return all PDF file paths.
fn find_pdf_files(input_dir: &Path) -> Result<Vec<PathBuf>> {
    if !input_dir.is_dir() {
        return Err(anyhow!("{} is not a directory", input_dir.display()));
    }

    let mut pdf_files = Vec::new();
    for entry in fs::read_dir(input_dir)
        .with_context(|| format!("Failed to read directory {}", input_dir.display()))?
    {
        let entry = entry?;
        let path = entry.path();
        if path.is_file() {
            if let Some(ext) = path.extension() {
                if ext.to_string_lossy().eq_ignore_ascii_case("pdf") {
                    pdf_files.push(path);
                }
            }
        }
    }

    Ok(pdf_files)
}

const MAX_CHUNK_CHARS: usize = 10000;

/// Split large PDF text into manageable chunks while preserving whole paragraphs.
pub fn split_text_into_chunks(text: &str, max_chars: usize) -> Vec<String> {
    let mut chunks = Vec::new();
    let mut current = String::new();

    for paragraph in text.split('\n').map(str::trim).filter(|p| !p.is_empty()) {
        if current.len() + paragraph.len() + 1 > max_chars {
            if !current.is_empty() {
                chunks.push(current);
                current = String::new();
            }
        }

        if paragraph.len() > max_chars {
            for piece in paragraph.chars().collect::<Vec<_>>().chunks(max_chars) {
                let piece_str: String = piece.iter().collect();
                if !current.is_empty() {
                    chunks.push(current);
                    current = String::new();
                }
                chunks.push(piece_str);
            }
        } else {
            if !current.is_empty() {
                current.push(' ');
            }
            current.push_str(paragraph);
        }
    }

    if !current.is_empty() {
        chunks.push(current);
    }

    chunks
}

pub async fn import_pdf_directory(
    input_dir: &Path,
    output_file: &Path,
    ollama: &Ollama,
    model_name: &str,
    prompt_each_chunk: bool,
) -> Result<bool> {
    let pdf_files = find_pdf_files(input_dir)?;
    if pdf_files.is_empty() {
        return Err(anyhow!("No PDF files found in {}", input_dir.display()));
    }

    if let Some(parent) = output_file.parent() {
        fs::create_dir_all(parent)?;
    }

    let mut writer = BufWriter::new(File::create(output_file)?);
    println!(
        "📦 Found {} PDF file(s) in {}",
        pdf_files.len(),
        input_dir.display()
    );

    let mut total_samples = 0;
    for pdf_path in &pdf_files {
        println!("📄 Converting {}", pdf_path.display());
        let text = extract_text(pdf_path)
            .with_context(|| format!("Failed to extract text from {}", pdf_path.display()))?;

        let pages_count = count_pdf_pages(pdf_path).unwrap_or(0);
        let chunks = split_text_into_chunks(&text, MAX_CHUNK_CHARS);
        if pages_count > 0 {
            println!(
                "   ➜ PDF produced {} chunk(s) from {} pages",
                chunks.len(),
                pages_count
            );
        } else {
            println!(
                "   ➜ PDF produced {} chunk(s); could not detect page count",
                chunks.len()
            );
        }

        let (count, completed) = process_pdf_chunks(
            pdf_path,
            &chunks,
            &mut writer,
            ollama,
            model_name,
            pages_count,
            prompt_each_chunk,
        )
        .await?;
        total_samples += count;
        if !completed {
            println!(
                "⚠️ Stopping directory conversion after partial output for {}",
                pdf_path.display()
            );
            return Ok(false);
        }
    }

    writer.flush()?;
    println!(
        "✅ Wrote {} JSONL sample(s) to {}",
        total_samples,
        output_file.display()
    );
    Ok(true)
}

pub async fn import_pdf_path(
    input_path: &Path,
    output_file: &Path,
    ollama: &Ollama,
    model_name: &str,
    prompt_each_chunk: bool,
) -> Result<bool> {
    if input_path.is_dir() {
        import_pdf_directory(
            input_path,
            output_file,
            ollama,
            model_name,
            prompt_each_chunk,
        )
        .await
    } else if input_path.is_file() {
        if input_path
            .extension()
            .map(|ext| ext.to_string_lossy().eq_ignore_ascii_case("pdf"))
            .unwrap_or(false)
        {
            if let Some(parent) = output_file.parent() {
                fs::create_dir_all(parent)?;
            }

            let mut writer = BufWriter::new(File::create(output_file)?);
            println!("📄 Converting {}", input_path.display());
            let text = extract_text(input_path)
                .with_context(|| format!("Failed to extract text from {}", input_path.display()))?;

            let pages_count = count_pdf_pages(input_path).unwrap_or(0);
            let chunks = split_text_into_chunks(&text, MAX_CHUNK_CHARS);
            if pages_count > 0 {
                println!(
                    "   ➜ PDF produced {} chunk(s) from {} pages",
                    chunks.len(),
                    pages_count
                );
            } else {
                println!(
                    "   ➜ PDF produced {} chunk(s); could not detect page count",
                    chunks.len()
                );
            }

            let (total_samples, completed) = process_pdf_chunks(
                input_path,
                &chunks,
                &mut writer,
                ollama,
                model_name,
                pages_count,
                prompt_each_chunk,
            )
            .await?;
            writer.flush()?;
            println!(
                "✅ Wrote {} JSONL sample(s) to {}",
                total_samples,
                output_file.display()
            );
            if completed {
                println!("✅ Done! Full PDF distillation complete.");
            } else {
                println!("⚠️ Partial output created; import stopped early.");
            }
            Ok(completed)
        } else {
            Err(anyhow!(
                "The selected file is not a PDF: {}",
                input_path.display()
            ))
        }
    } else {
        Err(anyhow!(
            "The selected path is neither a PDF file nor a directory: {}",
            input_path.display()
        ))
    }
}

async fn process_pdf_chunks(
    pdf_path: &Path,
    chunks: &[String],
    writer: &mut BufWriter<File>,
    ollama: &Ollama,
    model_name: &str,
    pages_count: usize,
    prompt_each_chunk: bool,
) -> Result<(usize, bool)> {
    let mut total_samples = 0;
    let batch_size: usize = 3;

    for batch_start in (0..chunks.len()).step_by(batch_size) {
        let batch_end = std::cmp::min(batch_start + batch_size, chunks.len());
        let batch = &chunks[batch_start..batch_end];

        for (offset, _chunk) in batch.iter().enumerate() {
            let index = batch_start + offset;
            let progress = format!("chunk {}/{}", index + 1, chunks.len());
            let page_label = if pages_count > 0 {
                let current_page = ((index * pages_count) / chunks.len()) + 1;
                format!("page {}/{}", current_page, pages_count)
            } else {
                "page unknown".to_string()
            };
            println!("   ➜ Processing {} ({})", progress, page_label);
        }

        let owned_chunks: Vec<String> = batch.to_vec();
        let samples =
            derive_samples_from_pdf_text(pdf_path, owned_chunks, ollama, model_name).await?;
        for sample in &samples {
            let line = serde_json::to_string(&sample)?;
            writeln!(writer, "{}", line)?;
            total_samples += 1;
        }

        if batch_end < chunks.len() && prompt_each_chunk {
            if !prompt_continue().await? {
                println!("⚠️ Stopped early. Partial output written to disk.");
                return Ok((total_samples, false));
            }
        }
    }
    Ok((total_samples, true))
}

fn count_pdf_pages(pdf_path: &Path) -> Result<usize> {
    let doc = Document::load(pdf_path)?;
    Ok(doc.get_pages().len())
}

async fn prompt_continue() -> Result<bool> {
    tokio::task::spawn_blocking(|| {
        print!("Press Enter to continue to the next chunk, or type 'quit' to stop: ");
        let stdout = io::stdout();
        let mut handle = stdout.lock();
        handle.flush()?;

        let mut input = String::new();
        io::stdin().read_line(&mut input)?;
        Ok(!input.trim().eq_ignore_ascii_case("quit"))
    })
    .await?
}