harmony-protocol 0.1.0

use wasm_bindgen::prelude::*;
use wasm_bindgen_futures::future_to_promise;
use serde_wasm_bindgen::{from_value, to_value};
use js_sys::{Array, Object, Promise};

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,
};

// Convert Rust error to JS Error
fn rust_error_to_js(err: anyhow::Error) -> JsValue {
    JsValue::from_str(&err.to_string())
}

// WASM wrapper for Role
#[wasm_bindgen]
pub struct WasmRole {
    inner: Role,
}

#[wasm_bindgen]
impl WasmRole {
    #[wasm_bindgen(constructor)]
    pub fn new(role: &str) -> Result<WasmRole, JsValue> {
        let inner = Role::try_from(role)
            .map_err(|e| JsValue::from_str(e))?;
        Ok(WasmRole { inner })
    }

    #[wasm_bindgen(js_name = "user")]
    pub fn user() -> WasmRole {
        WasmRole { inner: Role::User }
    }

    #[wasm_bindgen(js_name = "assistant")]
    pub fn assistant() -> WasmRole {
        WasmRole { inner: Role::Assistant }
    }

    #[wasm_bindgen(js_name = "system")]
    pub fn system() -> WasmRole {
        WasmRole { inner: Role::System }
    }

    #[wasm_bindgen(js_name = "developer")]
    pub fn developer() -> WasmRole {
        WasmRole { inner: Role::Developer }
    }

    #[wasm_bindgen(js_name = "tool")]
    pub fn tool() -> WasmRole {
        WasmRole { inner: Role::Tool }
    }

    #[wasm_bindgen(js_name = "toString")]
    pub fn to_string(&self) -> String {
        self.inner.as_str().to_string()
    }
}

// WASM wrapper for ReasoningEffort
#[wasm_bindgen]
pub struct WasmReasoningEffort {
    inner: ReasoningEffort,
}

#[wasm_bindgen]
impl WasmReasoningEffort {
    #[wasm_bindgen(js_name = "low")]
    pub fn low() -> WasmReasoningEffort {
        WasmReasoningEffort { inner: ReasoningEffort::Low }
    }

    #[wasm_bindgen(js_name = "medium")]
    pub fn medium() -> WasmReasoningEffort {
        WasmReasoningEffort { inner: ReasoningEffort::Medium }
    }

    #[wasm_bindgen(js_name = "high")]
    pub fn high() -> WasmReasoningEffort {
        WasmReasoningEffort { inner: ReasoningEffort::High }
    }

    #[wasm_bindgen(js_name = "toString")]
    pub fn to_string(&self) -> String {
        match self.inner {
            ReasoningEffort::Low => "low".to_string(),
            ReasoningEffort::Medium => "medium".to_string(),
            ReasoningEffort::High => "high".to_string(),
        }
    }
}

// WASM wrapper for ToolDescription
#[wasm_bindgen]
pub struct WasmToolDescription {
    inner: ToolDescription,
}

#[wasm_bindgen]
impl WasmToolDescription {
    #[wasm_bindgen(constructor)]
    pub fn new(name: String, description: String, parameters: Option<JsValue>) -> Result<WasmToolDescription, JsValue> {
        let parameters = if let Some(params) = parameters {
            Some(from_value(params).map_err(|e| JsValue::from_str(&e.to_string()))?)
        } else {
            None
        };

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

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

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

// WASM wrapper for SystemContent
#[wasm_bindgen]
pub struct WasmSystemContent {
    inner: SystemContent,
}

#[wasm_bindgen]
impl WasmSystemContent {
    #[wasm_bindgen(constructor)]
    pub fn new() -> WasmSystemContent {
        WasmSystemContent {
            inner: SystemContent::new(),
        }
    }

    #[wasm_bindgen(js_name = "withModelIdentity")]
    pub fn with_model_identity(mut self, model_identity: String) -> WasmSystemContent {
        self.inner = self.inner.with_model_identity(model_identity);
        self
    }

    #[wasm_bindgen(js_name = "withReasoningEffort")]
    pub fn with_reasoning_effort(mut self, effort: WasmReasoningEffort) -> WasmSystemContent {
        self.inner = self.inner.with_reasoning_effort(effort.inner);
        self
    }

    #[wasm_bindgen(js_name = "withRequiredChannels")]
    pub fn with_required_channels(mut self, channels: JsValue) -> Result<WasmSystemContent, JsValue> {
        let channels: Vec<String> = from_value(channels)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        self.inner = self.inner.with_required_channels(channels);
        Ok(self)
    }

    #[wasm_bindgen(js_name = "withBrowserTool")]
    pub fn with_browser_tool(mut self) -> WasmSystemContent {
        self.inner = self.inner.with_browser_tool();
        self
    }

    #[wasm_bindgen(js_name = "withPythonTool")]
    pub fn with_python_tool(mut self) -> WasmSystemContent {
        self.inner = self.inner.with_python_tool();
        self
    }

    #[wasm_bindgen(js_name = "withConversationStartDate")]
    pub fn with_conversation_start_date(mut self, date: String) -> WasmSystemContent {
        self.inner = self.inner.with_conversation_start_date(date);
        self
    }

    #[wasm_bindgen(js_name = "withKnowledgeCutoff")]
    pub fn with_knowledge_cutoff(mut self, cutoff: String) -> WasmSystemContent {
        self.inner = self.inner.with_knowledge_cutoff(cutoff);
        self
    }
}

// WASM wrapper for DeveloperContent
#[wasm_bindgen]
pub struct WasmDeveloperContent {
    inner: DeveloperContent,
}

#[wasm_bindgen]
impl WasmDeveloperContent {
    #[wasm_bindgen(constructor)]
    pub fn new() -> WasmDeveloperContent {
        WasmDeveloperContent {
            inner: DeveloperContent::new(),
        }
    }

    #[wasm_bindgen(js_name = "withInstructions")]
    pub fn with_instructions(mut self, instructions: String) -> WasmDeveloperContent {
        self.inner = self.inner.with_instructions(instructions);
        self
    }

    #[wasm_bindgen(js_name = "withFunctionTools")]
    pub fn with_function_tools(mut self, tools: JsValue) -> Result<WasmDeveloperContent, JsValue> {
        let tool_descriptions: Vec<WasmToolDescription> = from_value(tools)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        let tools: Vec<ToolDescription> = tool_descriptions.into_iter().map(|t| t.inner).collect();
        self.inner = self.inner.with_function_tools(tools);
        Ok(self)
    }
}

// WASM wrapper for Message
#[wasm_bindgen]
pub struct WasmMessage {
    inner: Message,
}

#[wasm_bindgen]
impl WasmMessage {
    #[wasm_bindgen(js_name = "fromRoleAndContent")]
    pub fn from_role_and_content(role: WasmRole, content: JsValue) -> Result<WasmMessage, JsValue> {
        let content = if content.is_instance_of::<WasmSystemContent>() {
            let sys_content: WasmSystemContent = from_value(content)
                .map_err(|e| JsValue::from_str(&e.to_string()))?;
            Content::SystemContent(sys_content.inner)
        } else if content.is_instance_of::<WasmDeveloperContent>() {
            let dev_content: WasmDeveloperContent = from_value(content)
                .map_err(|e| JsValue::from_str(&e.to_string()))?;
            Content::DeveloperContent(dev_content.inner)
        } else {
            let text: String = from_value(content)
                .map_err(|e| JsValue::from_str(&e.to_string()))?;
            Content::Text(TextContent { text })
        };

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

    #[wasm_bindgen(js_name = "withChannel")]
    pub fn with_channel(mut self, channel: String) -> WasmMessage {
        self.inner = self.inner.with_channel(channel);
        self
    }

    #[wasm_bindgen(js_name = "withRecipient")]
    pub fn with_recipient(mut self, recipient: String) -> WasmMessage {
        self.inner = self.inner.with_recipient(recipient);
        self
    }

    #[wasm_bindgen(js_name = "withContentType")]
    pub fn with_content_type(mut self, content_type: String) -> WasmMessage {
        self.inner = self.inner.with_content_type(content_type);
        self
    }

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

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

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

// WASM wrapper for Conversation
#[wasm_bindgen]
pub struct WasmConversation {
    inner: Conversation,
}

#[wasm_bindgen]
impl WasmConversation {
    #[wasm_bindgen(js_name = "fromMessages")]
    pub fn from_messages(messages: JsValue) -> Result<WasmConversation, JsValue> {
        let wasm_messages: Vec<WasmMessage> = from_value(messages)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        let messages: Vec<Message> = wasm_messages.into_iter().map(|m| m.inner).collect();
        Ok(WasmConversation {
            inner: Conversation::from_messages(messages),
        })
    }

    #[wasm_bindgen(getter)]
    pub fn length(&self) -> usize {
        self.inner.messages.len()
    }
}

// WASM wrapper for HarmonyEncoding
#[wasm_bindgen]
pub struct WasmHarmonyEncoding {
    inner: HarmonyEncoding,
}

#[wasm_bindgen]
impl WasmHarmonyEncoding {
    #[wasm_bindgen(getter)]
    pub fn name(&self) -> String {
        self.inner.name().to_string()
    }

    #[wasm_bindgen(js_name = "renderConversationForCompletion")]
    pub fn render_conversation_for_completion(
        &self,
        conversation: WasmConversation,
        next_turn_role: WasmRole
    ) -> Result<Array, JsValue> {
        let tokens = self.inner
            .render_conversation_for_completion(&conversation.inner, next_turn_role.inner, None)
            .map_err(rust_error_to_js)?;

        let js_array = Array::new();
        for token in tokens {
            js_array.push(&JsValue::from(token));
        }
        Ok(js_array)
    }

    #[wasm_bindgen(js_name = "renderConversation")]
    pub fn render_conversation(&self, conversation: WasmConversation) -> Result<Array, JsValue> {
        let tokens = self.inner
            .render_conversation(&conversation.inner, None)
            .map_err(rust_error_to_js)?;

        let js_array = Array::new();
        for token in tokens {
            js_array.push(&JsValue::from(token));
        }
        Ok(js_array)
    }

    #[wasm_bindgen(js_name = "parseMessagesFromCompletionTokens")]
    pub fn parse_messages_from_completion_tokens(
        &self,
        tokens: JsValue,
        role: Option<WasmRole>
    ) -> Result<Array, JsValue> {
        let tokens: Vec<u32> = from_value(tokens)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        let role = role.map(|r| r.inner);

        let messages = self.inner
            .parse_messages_from_completion_tokens(tokens, role)
            .map_err(rust_error_to_js)?;

        let js_array = Array::new();
        for message in messages {
            let wasm_message = WasmMessage { inner: message };
            js_array.push(&to_value(&wasm_message).map_err(|e| JsValue::from_str(&e.to_string()))?);
        }
        Ok(js_array)
    }
}

// WASM wrapper for HarmonyEncodingName
#[wasm_bindgen]
pub struct WasmHarmonyEncodingName {
    inner: HarmonyEncodingName,
}

#[wasm_bindgen]
impl WasmHarmonyEncodingName {
    #[wasm_bindgen(js_name = "harmonyGptOss")]
    pub fn harmony_gpt_oss() -> WasmHarmonyEncodingName {
        WasmHarmonyEncodingName {
            inner: HarmonyEncodingName::HarmonyGptOss,
        }
    }

    #[wasm_bindgen(js_name = "toString")]
    pub fn to_string(&self) -> String {
        self.inner.to_string()
    }
}

// WASM wrapper for StreamableParser
#[wasm_bindgen]
pub struct WasmStreamableParser {
    inner: StreamableParser,
}

#[wasm_bindgen]
impl WasmStreamableParser {
    #[wasm_bindgen(constructor)]
    pub fn new(encoding: WasmHarmonyEncoding, role: Option<WasmRole>) -> Result<WasmStreamableParser, JsValue> {
        let role = role.map(|r| r.inner);
        let inner = StreamableParser::new(encoding.inner, role)
            .map_err(rust_error_to_js)?;
        Ok(WasmStreamableParser { inner })
    }

    pub fn process(&mut self, token: u32) -> Result<(), JsValue> {
        self.inner.process(token).map_err(rust_error_to_js)
    }

    #[wasm_bindgen(js_name = "processEos")]
    pub fn process_eos(&mut self) -> Result<(), JsValue> {
        self.inner.process_eos().map_err(rust_error_to_js)
    }

    #[wasm_bindgen(js_name = "currentContent")]
    pub fn current_content(&self) -> Result<String, JsValue> {
        self.inner.current_content().map_err(rust_error_to_js)
    }

    #[wasm_bindgen(js_name = "currentRole")]
    pub fn current_role(&self) -> Option<WasmRole> {
        self.inner.current_role().map(|r| WasmRole { inner: r })
    }

    #[wasm_bindgen(js_name = "lastContentDelta")]
    pub fn last_content_delta(&self) -> Result<Option<String>, JsValue> {
        self.inner.last_content_delta().map_err(rust_error_to_js)
    }

    #[wasm_bindgen(js_name = "intoMessages")]
    pub fn into_messages(self) -> Result<Array, JsValue> {
        let messages = self.inner.into_messages();
        let js_array = Array::new();
        for message in messages {
            let wasm_message = WasmMessage { inner: message };
            js_array.push(&to_value(&wasm_message).map_err(|e| JsValue::from_str(&e.to_string()))?);
        }
        Ok(js_array)
    }

    #[wasm_bindgen(js_name = "currentChannel")]
    pub fn current_channel(&self) -> Option<String> {
        self.inner.current_channel()
    }

    #[wasm_bindgen(js_name = "currentRecipient")]
    pub fn current_recipient(&self) -> Option<String> {
        self.inner.current_recipient()
    }
}

// Main loading function (async for WASM)
#[wasm_bindgen(js_name = "loadHarmonyEncoding")]
pub fn load_harmony_encoding_wasm(name: WasmHarmonyEncodingName) -> Promise {
    let inner_name = name.inner;
    future_to_promise(async move {
        let inner = crate::load_harmony_encoding(inner_name)
            .await
            .map_err(rust_error_to_js)?;
        let wasm_encoding = WasmHarmonyEncoding { inner };
        to_value(&wasm_encoding).map_err(|e| JsValue::from_str(&e.to_string()))
    })
}

// Set panic hook for better error messages in WASM
#[wasm_bindgen(start)]
pub fn main() {
    #[cfg(feature = "console_error_panic_hook")]
    console_error_panic_hook::set_once();
}