[content] /home/franz.srambical/crowd-pilot-serializer/crates/python/src/lib.rs (rust)
//! Python bindings for the crowd-pilot serializer.

use pyo3::prelude::*;
use pyo3::types::{PyDict, PyModule};
use std::sync::Mutex;

use crowd_pilot_serializer_core::{
    ConversationMessage as CoreMessage, ConversationStateManager as CoreManager,
    ConversationStateManagerConfig,
};

/// A message in the conversation.
#[pyclass]
#[derive(Clone)]
pub struct ConversationMessage {
    #[pyo3(get)]
    pub from_: String,
    #[pyo3(get)]
    pub value: String,
}

#[pymethods]
impl ConversationMessage {
    fn __repr__(&self) -> String {
        format!(
            "ConversationMessage(from_='{}', value='{}...')",
            self.from_,
            &self.value[..self.value.len().min(50)]
        )
    }

    /// Convert to a dictionary.
    fn to_dict(&self, py: Python<'_>) -> PyResult<Py<PyDict>> {
        let dict = PyDict::new(py);
        dict.set_item("from", &self.from_)?;
        dict.set_item("value", &self.value)?;
        Ok(dict.into())
    }
}

impl From<CoreMessage> for ConversationMessage {
    fn from(msg: CoreMessage) -> Self {
        Self {
            from_: msg.from,
            value: msg.value,
        }
    }
}

/// Token counter function type.
type TokenCounterFn = Box<dyn Fn(&str) -> usize + Send>;

/// Manages conversation state for serializing IDE events.
///
/// The token counter is provided at construction time, allowing the caller to use
/// either a character-based approximation (for runtime) or an accurate tokenizer.
#[pyclass]
pub struct ConversationStateManager {
    inner: Mutex<CoreManager<TokenCounterFn>>,
}

#[pymethods]
impl ConversationStateManager {
    /// Create a new ConversationStateManager.
    ///
    /// Args:
    ///     count_tokens: Function that returns the token count for a given string.
    ///                   For preprocessing: `lambda text: len(tokenizer.encode(text))`
    ///     viewport_radius: Lines above/below cursor to show. Default: 10
    ///     coalesce_radius: Radius for grouping nearby edits. Default: 5
    ///     max_tokens_per_message: Maximum tokens per message. Default: 2048
    ///     max_tokens_per_terminal_output: Maximum tokens per terminal output. Default: 256
    #[new]
    #[pyo3(signature = (count_tokens, viewport_radius=None, coalesce_radius=None, max_tokens_per_message=None, max_tokens_per_terminal_output=None))]
    fn new(
        count_tokens: PyObject,
        viewport_radius: Option<usize>,
        coalesce_radius: Option<usize>,
        max_tokens_per_message: Option<usize>,
        max_tokens_per_terminal_output: Option<usize>,
    ) -> PyResult<Self> {
        let token_counter: TokenCounterFn = Box::new(move |text: &str| {
            Python::with_gil(|py| {
                count_tokens
                    .call1(py, (text,))
                    .and_then(|result| result.extract::<usize>(py))
                    .unwrap_or_else(|_| text.len() / 4)
            })
        });

        let config = ConversationStateManagerConfig {
            viewport_radius: viewport_radius.unwrap_or(10),
            coalesce_radius: coalesce_radius.unwrap_or(5),
            max_tokens_per_message: max_tokens_per_message.unwrap_or(2048),
            max_tokens_per_terminal_output: max_tokens_per_terminal_output.unwrap_or(256),
        };

        Ok(Self {
            inner: Mutex::new(CoreManager::new(token_counter, config)),
        })
    }

    /// Reset all state.
    fn reset(&self) -> PyResult<()> {
        let mut inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        inner.reset();
        Ok(())
    }

    /// Get a copy of all messages.
    fn get_messages(&self) -> PyResult<Vec<ConversationMessage>> {
        let inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        Ok(inner.get_messages().into_iter().map(Into::into).collect())
    }

    /// Get the current content of a file.
    fn get_file_content(&self, file_path: &str) -> PyResult<String> {
        let inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        Ok(inner.get_file_content(file_path))
    }

    /// Handle a tab (file switch) event.
    ///
    /// Args:
    ///     file_path: The path to the file.
    ///     text_content: The file contents, or None if switching to an already-open file.
    #[pyo3(signature = (file_path, text_content=None))]
    fn handle_tab_event(&self, file_path: &str, text_content: Option<&str>) -> PyResult<()> {
        let mut inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        inner.handle_tab_event(file_path, text_content);
        Ok(())
    }

    /// Handle a content change event.
    ///
    /// Args:
    ///     file_path: The path to the file.
    ///     offset: The character offset where the change starts.
    ///     length: The number of characters being replaced.
    ///     new_text: The new text being inserted.
    fn handle_content_event(
        &self,
        file_path: &str,
        offset: usize,
        length: usize,
        new_text: &str,
    ) -> PyResult<()> {
        let mut inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        inner.handle_content_event(file_path, offset, length, new_text);
        Ok(())
    }

    /// Handle a selection event.
    ///
    /// Args:
    ///     file_path: The path to the file.
    ///     offset: The character offset of the selection start.
    fn handle_selection_event(&self, file_path: &str, offset: usize) -> PyResult<()> {
        let mut inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        inner.handle_selection_event(file_path, offset);
        Ok(())
    }

    /// Handle a terminal command event.
    ///
    /// Args:
    ///     command: The command that was executed.
    fn handle_terminal_command_event(&self, command: &str) -> PyResult<()> {
        let mut inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        inner.handle_terminal_command_event(command);
        Ok(())
    }

    /// Handle a terminal output event.
    ///
    /// Args:
    ///     output: The terminal output.
    fn handle_terminal_output_event(&self, output: &str) -> PyResult<()> {
        let mut inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        inner.handle_terminal_output_event(output);
        Ok(())
    }

    /// Handle a terminal focus event.
    fn handle_terminal_focus_event(&self) -> PyResult<()> {
        let mut inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        inner.handle_terminal_focus_event();
        Ok(())
    }

    /// Handle a git branch checkout event.
    ///
    /// Args:
    ///     branch_info: The git checkout message containing the branch name.
    fn handle_git_branch_checkout_event(&self, branch_info: &str) -> PyResult<()> {
        let mut inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        inner.handle_git_branch_checkout_event(branch_info);
        Ok(())
    }

    /// Finalize and get conversation ready for model.
    fn finalize_for_model(&self) -> PyResult<Vec<ConversationMessage>> {
        let mut inner = self
            .inner
            .lock()
            .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Lock poisoned"))?;
        Ok(inner
            .finalize_for_model()
            .into_iter()
            .map(Into::into)
            .collect())
    }
}

/// Estimate tokens using character approximation.
/// Uses ~4 characters per token as a rough approximation.
#[pyfunction]
fn estimate_tokens(text: &str) -> usize {
    (text.len() + 3) / 4
}

/// Clean text by normalizing line endings.
#[pyfunction]
fn clean_text(text: &str) -> String {
    crowd_pilot_serializer_core::clean_text(text)
}

/// Create a fenced code block.
#[pyfunction]
#[pyo3(signature = (content, language=None))]
fn fenced_block(content: &str, language: Option<&str>) -> String {
    crowd_pilot_serializer_core::fenced_block(language, content)
}

/// Normalize terminal output.
#[pyfunction]
fn normalize_terminal_output(raw: &str) -> String {
    crowd_pilot_serializer_core::normalize_terminal_output(raw)
}

/// Generate line-numbered output.
#[pyfunction]
#[pyo3(signature = (content, start_line=None, end_line=None))]
fn line_numbered_output(content: &str, start_line: Option<usize>, end_line: Option<usize>) -> String {
    crowd_pilot_serializer_core::line_numbered_output(content, start_line, end_line)
}

/// Python module definition.
#[pymodule]
fn crowd_pilot_serializer(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_class::<ConversationMessage>()?;
    m.add_class::<ConversationStateManager>()?;
    m.add_function(wrap_pyfunction!(estimate_tokens, m)?)?;
    m.add_function(wrap_pyfunction!(clean_text, m)?)?;
    m.add_function(wrap_pyfunction!(fenced_block, m)?)?;
    m.add_function(wrap_pyfunction!(normalize_terminal_output, m)?)?;
    m.add_function(wrap_pyfunction!(line_numbered_output, m)?)?;
    Ok(())
}