harmony-protocol 0.1.0

use pyo3::prelude::*;
use pyo3::types::PyDict;
use std::collections::HashMap;

use crate::{
    chat::{
        Author, ChannelConfig, Content, Conversation, DeveloperContent, Message,
        ReasoningEffort, Role, SystemContent, TextContent, ToolDescription, ToolNamespaceConfig
    },
    encoding::{HarmonyEncoding, RenderConversationConfig, StreamableParser},
    load_harmony_encoding, HarmonyEncodingName,
    tiktoken::Rank,
};

// Python wrapper for Role
#[pyclass]
#[derive(Clone)]
pub struct PyRole {
    inner: Role,
}

#[pymethods]
impl PyRole {
    #[new]
    fn new(role: &str) -> PyResult<Self> {
        let inner = Role::try_from(role)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))?;
        Ok(PyRole { inner })
    }

    #[staticmethod]
    fn user() -> PyRole {
        PyRole { inner: Role::User }
    }

    #[staticmethod]
    fn assistant() -> PyRole {
        PyRole { inner: Role::Assistant }
    }

    #[staticmethod]
    fn system() -> PyRole {
        PyRole { inner: Role::System }
    }

    #[staticmethod]
    fn developer() -> PyRole {
        PyRole { inner: Role::Developer }
    }

    #[staticmethod]
    fn tool() -> PyRole {
        PyRole { inner: Role::Tool }
    }

    fn __str__(&self) -> String {
        self.inner.as_str().to_string()
    }

    fn __repr__(&self) -> String {
        format!("Role('{}')", self.inner.as_str())
    }
}

// Python wrapper for ReasoningEffort
#[pyclass]
#[derive(Clone)]
pub struct PyReasoningEffort {
    inner: ReasoningEffort,
}

#[pymethods]
impl PyReasoningEffort {
    #[staticmethod]
    fn low() -> PyReasoningEffort {
        PyReasoningEffort { inner: ReasoningEffort::Low }
    }

    #[staticmethod]
    fn medium() -> PyReasoningEffort {
        PyReasoningEffort { inner: ReasoningEffort::Medium }
    }

    #[staticmethod]
    fn high() -> PyReasoningEffort {
        PyReasoningEffort { inner: ReasoningEffort::High }
    }

    fn __str__(&self) -> String {
        match self.inner {
            ReasoningEffort::Low => "low".to_string(),
            ReasoningEffort::Medium => "medium".to_string(),
            ReasoningEffort::High => "high".to_string(),
        }
    }
}

// Python wrapper for ToolDescription
#[pyclass]
#[derive(Clone)]
pub struct PyToolDescription {
    inner: ToolDescription,
}

#[pymethods]
impl PyToolDescription {
    #[new]
    fn new(name: String, description: String, parameters: Option<PyObject>) -> PyResult<Self> {
        let parameters = if let Some(params) = parameters {
            Python::with_gil(|py| {
                let json_str = params.call_method0(py, "to_json")
                    .or_else(|_| {
                        // Fallback to serde_json serialization
                        let dict = params.extract::<&PyDict>(py)?;
                        let json_value = pythonize::depythonize(dict)?;
                        Ok::<PyObject, PyErr>(serde_json::to_string(&json_value)?.into_py(py))
                    })?;
                let json_str: String = json_str.extract(py)?;
                serde_json::from_str(&json_str).map_err(PyErr::from)
            })?
        } else {
            None
        };

        Ok(PyToolDescription {
            inner: ToolDescription::new(name, description, parameters),
        })
    }

    #[getter]
    fn name(&self) -> String {
        self.inner.name.clone()
    }

    #[getter]
    fn description(&self) -> String {
        self.inner.description.clone()
    }
}

// Python wrapper for SystemContent
#[pyclass]
#[derive(Clone)]
pub struct PySystemContent {
    inner: SystemContent,
}

#[pymethods]
impl PySystemContent {
    #[new]
    fn new() -> Self {
        PySystemContent {
            inner: SystemContent::new(),
        }
    }

    fn with_model_identity(mut self, model_identity: String) -> Self {
        self.inner = self.inner.with_model_identity(model_identity);
        self
    }

    fn with_reasoning_effort(mut self, effort: PyReasoningEffort) -> Self {
        self.inner = self.inner.with_reasoning_effort(effort.inner);
        self
    }

    fn with_required_channels(mut self, channels: Vec<String>) -> Self {
        self.inner = self.inner.with_required_channels(channels);
        self
    }

    fn with_browser_tool(mut self) -> Self {
        self.inner = self.inner.with_browser_tool();
        self
    }

    fn with_python_tool(mut self) -> Self {
        self.inner = self.inner.with_python_tool();
        self
    }

    fn with_conversation_start_date(mut self, date: String) -> Self {
        self.inner = self.inner.with_conversation_start_date(date);
        self
    }

    fn with_knowledge_cutoff(mut self, cutoff: String) -> Self {
        self.inner = self.inner.with_knowledge_cutoff(cutoff);
        self
    }
}

// Python wrapper for DeveloperContent
#[pyclass]
#[derive(Clone)]
pub struct PyDeveloperContent {
    inner: DeveloperContent,
}

#[pymethods]
impl PyDeveloperContent {
    #[new]
    fn new() -> Self {
        PyDeveloperContent {
            inner: DeveloperContent::new(),
        }
    }

    fn with_instructions(mut self, instructions: String) -> Self {
        self.inner = self.inner.with_instructions(instructions);
        self
    }

    fn with_function_tools(mut self, tools: Vec<PyToolDescription>) -> Self {
        let tools: Vec<ToolDescription> = tools.into_iter().map(|t| t.inner).collect();
        self.inner = self.inner.with_function_tools(tools);
        self
    }
}

// Python wrapper for Message
#[pyclass]
#[derive(Clone)]
pub struct PyMessage {
    inner: Message,
}

#[pymethods]
impl PyMessage {
    #[staticmethod]
    fn from_role_and_content(role: PyRole, content: PyObject) -> PyResult<Self> {
        Python::with_gil(|py| {
            let content = if content.is_instance_of::<PySystemContent>(py)? {
                let sys_content: PySystemContent = content.extract(py)?;
                Content::SystemContent(sys_content.inner)
            } else if content.is_instance_of::<PyDeveloperContent>(py)? {
                let dev_content: PyDeveloperContent = content.extract(py)?;
                Content::DeveloperContent(dev_content.inner)
            } else {
                let text: String = content.extract(py)?;
                Content::Text(TextContent { text })
            };

            Ok(PyMessage {
                inner: Message::from_role_and_content(role.inner, content),
            })
        })
    }

    fn with_channel(mut self, channel: String) -> Self {
        self.inner = self.inner.with_channel(channel);
        self
    }

    fn with_recipient(mut self, recipient: String) -> Self {
        self.inner = self.inner.with_recipient(recipient);
        self
    }

    fn with_content_type(mut self, content_type: String) -> Self {
        self.inner = self.inner.with_content_type(content_type);
        self
    }

    #[getter]
    fn channel(&self) -> Option<String> {
        self.inner.channel.clone()
    }

    #[getter]
    fn recipient(&self) -> Option<String> {
        self.inner.recipient.clone()
    }

    #[getter]
    fn role(&self) -> String {
        self.inner.author.role.as_str().to_string()
    }
}

// Python wrapper for Conversation
#[pyclass]
#[derive(Clone)]
pub struct PyConversation {
    inner: Conversation,
}

#[pymethods]
impl PyConversation {
    #[staticmethod]
    fn from_messages(messages: Vec<PyMessage>) -> Self {
        let messages: Vec<Message> = messages.into_iter().map(|m| m.inner).collect();
        PyConversation {
            inner: Conversation::from_messages(messages),
        }
    }

    fn __len__(&self) -> usize {
        self.inner.messages.len()
    }
}

// Python wrapper for HarmonyEncoding
#[pyclass]
pub struct PyHarmonyEncoding {
    inner: HarmonyEncoding,
}

#[pymethods]
impl PyHarmonyEncoding {
    fn name(&self) -> String {
        self.inner.name().to_string()
    }

    fn render_conversation_for_completion(
        &self,
        conversation: PyConversation,
        next_turn_role: PyRole
    ) -> PyResult<Vec<u32>> {
        let tokens = self.inner
            .render_conversation_for_completion(&conversation.inner, next_turn_role.inner, None)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
        Ok(tokens)
    }

    fn render_conversation(
        &self,
        conversation: PyConversation
    ) -> PyResult<Vec<u32>> {
        let tokens = self.inner
            .render_conversation(&conversation.inner, None)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
        Ok(tokens)
    }

    fn parse_messages_from_completion_tokens(
        &self,
        tokens: Vec<u32>,
        role: Option<PyRole>
    ) -> PyResult<Vec<PyMessage>> {
        let role = role.map(|r| r.inner);
        let messages = self.inner
            .parse_messages_from_completion_tokens(tokens, role)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

        Ok(messages.into_iter().map(|m| PyMessage { inner: m }).collect())
    }
}

// Python wrapper for HarmonyEncodingName
#[pyclass]
#[derive(Clone)]
pub struct PyHarmonyEncodingName {
    inner: HarmonyEncodingName,
}

#[pymethods]
impl PyHarmonyEncodingName {
    #[staticmethod]
    fn harmony_gpt_oss() -> PyHarmonyEncodingName {
        PyHarmonyEncodingName {
            inner: HarmonyEncodingName::HarmonyGptOss,
        }
    }

    fn __str__(&self) -> String {
        self.inner.to_string()
    }
}

// Python wrapper for StreamableParser
#[pyclass]
pub struct PyStreamableParser {
    inner: StreamableParser,
}

#[pymethods]
impl PyStreamableParser {
    #[new]
    fn new(encoding: PyHarmonyEncoding, role: Option<PyRole>) -> PyResult<Self> {
        let role = role.map(|r| r.inner);
        let inner = StreamableParser::new(encoding.inner, role)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
        Ok(PyStreamableParser { inner })
    }

    fn process(&mut self, token: u32) -> PyResult<()> {
        self.inner.process(token)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
        Ok(())
    }

    fn process_eos(&mut self) -> PyResult<()> {
        self.inner.process_eos()
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
        Ok(())
    }

    fn current_content(&self) -> PyResult<String> {
        self.inner.current_content()
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
    }

    fn current_role(&self) -> Option<PyRole> {
        self.inner.current_role().map(|r| PyRole { inner: r })
    }

    fn last_content_delta(&self) -> PyResult<Option<String>> {
        self.inner.last_content_delta()
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
    }

    fn into_messages(self) -> Vec<PyMessage> {
        self.inner.into_messages()
            .into_iter()
            .map(|m| PyMessage { inner: m })
            .collect()
    }

    fn current_channel(&self) -> Option<String> {
        self.inner.current_channel()
    }

    fn current_recipient(&self) -> Option<String> {
        self.inner.current_recipient()
    }
}

// Main loading function
#[pyfunction]
fn load_harmony_encoding(name: PyHarmonyEncodingName) -> PyResult<PyHarmonyEncoding> {
    let inner = crate::load_harmony_encoding(name.inner)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
    Ok(PyHarmonyEncoding { inner })
}

// Module definition
#[pymodule]
fn harmony_protocol(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<PyRole>()?;
    m.add_class::<PyReasoningEffort>()?;
    m.add_class::<PyToolDescription>()?;
    m.add_class::<PySystemContent>()?;
    m.add_class::<PyDeveloperContent>()?;
    m.add_class::<PyMessage>()?;
    m.add_class::<PyConversation>()?;
    m.add_class::<PyHarmonyEncoding>()?;
    m.add_class::<PyHarmonyEncodingName>()?;
    m.add_class::<PyStreamableParser>()?;
    m.add_function(wrap_pyfunction!(load_harmony_encoding, m)?)?;
    Ok(())
}