//! Core library for Krira Augment
//! 
//! Implements high-performance parallel file processing.

use std::fs::File;
use std::io::{BufWriter, Write};


use memmap2::MmapOptions;
use pyo3::prelude::*;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};

mod cleaning;
mod chunker;
mod errors;

use cleaning::RustCleaner;
use chunker::RustChunker;
use errors::KriraError;

// =============================================================================
// structs
// =============================================================================

#[derive(Serialize, Deserialize)]
struct PipelineConfig {
    max_chars: usize,
}

#[derive(Serialize)]
struct ChunkObj {
    text: String,
    length: usize,
}

// =============================================================================
// Core Function
// =============================================================================

/// Process a file in parallel using Rust.
/// 
/// 1. Maps file to memory (Zero-copy read)
/// 2. Iterates lines in parallel (Rayon)
/// 3. Cleans and Chunks lines
/// 4. Writes results to output file
#[pyfunction]
fn process_file_rust(py: Python, input_path: String, output_path: String, config_json: String) -> PyResult<()> {
    
    // 1. Parse Config
    let config: PipelineConfig = serde_json::from_str(&config_json)
        .map_err(|e| KriraError::ConfigError(e.to_string()))?;
        
    let chunker = RustChunker::new(config.max_chars);

    // 2. Prepare Output Writer (Thread-safe via Channel)
    // We use a bounded channel to prevent unlimited memory growth if reading is faster than writing
    let (tx, rx) = std::sync::mpsc::sync_channel::<Vec<ChunkObj>>(5000); // 5000 batches buffer

    // 3. Release GIL for heavy lifting
    py.allow_threads(move || -> Result<(), KriraError> {
        // Spawn Writer Thread
        let writer_handle = std::thread::spawn(move || -> Result<(), KriraError> {
            let output_file = File::create(&output_path).map_err(KriraError::IOError)?;
            let mut writer = BufWriter::new(output_file);
            
            while let Ok(batch) = rx.recv() {
                for res in batch {
                    if let Ok(json_line) = serde_json::to_string(&res) {
                        writeln!(writer, "{}", json_line).map_err(KriraError::IOError)?;
                    }
                }
            }
            writer.flush().map_err(KriraError::IOError)?;
            Ok(())
        });

        // Open Input File
        let file = File::open(&input_path).map_err(KriraError::IOError)?;
        
        // Mmap
        let mmap = unsafe { MmapOptions::new().map(&file).map_err(KriraError::IOError)? };
        
        // Convert to str (assumes UTF-8)
        let content = std::str::from_utf8(&mmap[..])
            .map_err(|e| KriraError::ConfigError(format!("File is not valid UTF-8: {}", e)))?;

        // Parallel Processing with Channel
        content.par_lines()
            .for_each_with(tx, |sender, line| {
                if line.trim().is_empty() { return; }

                // Clean
                let cleaned = RustCleaner::clean(line);
                if cleaned.is_empty() { return; }

                // Chunk
                let chunks = chunker.chunk(&cleaned);

                if chunks.is_empty() { return; }

                // Serialize Results
                let results: Vec<ChunkObj> = chunks.into_iter().map(|c| ChunkObj {
                    length: c.len(),
                    text: c,
                }).collect();

                // Send to Writer Thread
                // This might block if channel is full, providing backpressure
                let _ = sender.send(results);
            });

        // Drop the last sender so the receiver knows we are done
        // (Implicitly done as `content.par_lines` finishes and all `tx` clones are dropped)
        
        // Wait for writer to finish
        match writer_handle.join() {
            Ok(result) => result,
            Err(e) => Err(KriraError::ConfigError(format!("Writer thread panicked: {:?}", e))),
        }
    }).map_err(PyErr::from)
}

// =============================================================================
// Python Module
// =============================================================================

#[pymodule]
fn _rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(process_file_rust, m)?)?;
    Ok(())
}
