baml 0.221.0

#![allow(unsafe_code)]
use std::{
    collections::HashMap,
    ffi::{c_char, c_void, CString},
};

use prost::Message;

use crate::{
    args::FunctionArgs,
    async_stream::AsyncStreamingCall,
    codec::{traits::DecodeHandle, BamlDecode},
    error::BamlError,
    ffi::{self, callbacks},
    proto::baml_cffi_v1::{
        invocation_response::Response as InvResponse,
        invocation_response_success::Result as InvSuccessResult, CffiValueHolder,
        InvocationResponse,
    },
    raw_objects::{Audio, Collector, HTTPRequest, Image, Pdf, TypeBuilder, Video},
    stream::StreamingCall,
};

/// Handle to the BAML runtime
pub struct BamlRuntime {
    ptr: *const c_void,
}

// Safety: The runtime is thread-safe internally (protected by Rust's runtime)
#[allow(unsafe_code)]
unsafe impl Send for BamlRuntime {}
#[allow(unsafe_code)]
unsafe impl Sync for BamlRuntime {}

pub type StaticRuntimeType = once_cell::sync::Lazy<BamlRuntime>;

impl BamlRuntime {
    /// Create a new runtime from embedded BAML source files
    ///
    /// # Arguments
    /// * `baml_src_dir` - Base directory path for BAML sources
    /// * `files` - Map of relative file paths to file contents
    /// * `env` - Environment variables
    pub fn new(
        baml_src_dir: &str,
        files: &HashMap<String, String>,
        env: &HashMap<String, String>,
    ) -> Result<Self, BamlError> {
        // Initialize callbacks first - now returns Result
        callbacks::initialize_callbacks()
            .map_err(|e| BamlError::internal(format!("Failed to load BAML library: {e}")))?;

        // Encode files and env as JSON (matching CFFI format)
        let files_json = json_encode_map(files)?;
        let env_json = json_encode_map(env)?;

        let dir_cstr = CString::new(baml_src_dir)
            .map_err(|_| BamlError::internal("invalid baml_src_dir path (contains null byte)"))?;
        let files_cstr = CString::new(files_json)
            .map_err(|_| BamlError::internal("invalid files json (contains null byte)"))?;
        let env_cstr = CString::new(env_json)
            .map_err(|_| BamlError::internal("invalid env json (contains null byte)"))?;

        #[allow(unsafe_code)]
        let ptr = unsafe {
            ffi::create_baml_runtime(dir_cstr.as_ptr(), files_cstr.as_ptr(), env_cstr.as_ptr())
                .map_err(|e| BamlError::internal(format!("Failed to load BAML library: {e}")))?
        };

        if ptr.is_null() {
            return Err(BamlError::internal("failed to create runtime"));
        }

        Ok(BamlRuntime { ptr })
    }

    /// Call a function synchronously (blocks until complete)
    pub fn call_function<T: BamlDecode>(
        &self,
        name: &str,
        args: &FunctionArgs,
    ) -> Result<T, BamlError> {
        let encoded = args.encode()?;
        let name_cstr =
            CString::new(name).map_err(|_| BamlError::internal("invalid function name"))?;

        let (id, receiver) = callbacks::create_callback();

        #[allow(unsafe_code)]
        let buf = unsafe {
            ffi::call_function_from_c(
                self.ptr,
                name_cstr.as_ptr(),
                encoded.as_ptr().cast::<c_char>(),
                encoded.len(),
                id,
            )
            .map_err(|e| {
                callbacks::remove_callback(id);
                BamlError::internal(format!("Failed to load BAML library: {e}"))
            })?
        };

        // Check for immediate error (decode Buffer response)
        ffi::decode_async_response(buf).map_err(|e| {
            callbacks::remove_callback(id);
            BamlError::internal(e)
        })?;

        // Set up cancellation callback if token provided
        // Guard is dropped when function returns, stopping the watcher
        let _cancel_guard = args.cancellation_token.as_ref().map(|token| {
            token.on_cancel(move || {
                #[allow(unsafe_code)]
                unsafe {
                    let _ = ffi::cancel_function_call(id);
                }
            })
        });

        // Wait for result
        let result = receiver.recv();
        match result {
            Ok(callbacks::CallbackResult::Final(data)) => {
                let holder = CffiValueHolder::decode(&data[..])
                    .map_err(|e| BamlError::internal(format!("decode error: {e}")))?;
                T::baml_decode(&holder)
            }
            Ok(callbacks::CallbackResult::Partial(_)) => Err(BamlError::internal(
                "unexpected partial result in sync call",
            )),
            Ok(callbacks::CallbackResult::Error(e)) => Err(e),
            Err(_) => Err(BamlError::internal("callback channel closed")),
        }
    }

    /// Call a function with streaming results
    ///
    /// If `args` contains an `on_tick` callback (set via `FunctionArgs::with_on_tick`),
    /// it will be invoked for each SSE streaming chunk received from the LLM.
    /// A collector is automatically created and injected when on_tick is present.
    pub fn call_function_stream<TPartial, TFinal>(
        &self,
        name: &str,
        args: &FunctionArgs,
    ) -> Result<StreamingCall<TPartial, TFinal>, BamlError>
    where
        TPartial: BamlDecode + Send + 'static,
        TFinal: Clone + BamlDecode + Send + 'static,
    {
        let on_tick_data = args.on_tick.as_ref().map(|cb| {
            let collector = self.new_collector("on-tick-collector");
            let data = callbacks::OnTickData {
                callback: cb.clone(),
                collector: collector.clone(),
            };
            (data, collector)
        });

        let extra_collector = on_tick_data.as_ref().map(|(_, c)| c);
        let encoded = args.encode_with_extra_collector(extra_collector)?;
        let name_cstr =
            CString::new(name).map_err(|_| BamlError::internal("invalid function name"))?;

        let (id, receiver) =
            callbacks::create_callback_with_on_tick(on_tick_data.map(|(d, _)| d));

        #[allow(unsafe_code)]
        let buf = unsafe {
            ffi::call_function_stream_from_c(
                self.ptr,
                name_cstr.as_ptr(),
                encoded.as_ptr().cast::<c_char>(),
                encoded.len(),
                id,
            )
            .map_err(|e| {
                callbacks::remove_callback(id);
                BamlError::internal(format!("Failed to load BAML library: {e}"))
            })?
        };

        ffi::decode_async_response(buf).map_err(|e| {
            callbacks::remove_callback(id);
            BamlError::internal(e)
        })?;

        let cancel_guard = args.cancellation_token.as_ref().map(|token| {
            token.on_cancel(move || {
                #[allow(unsafe_code)]
                unsafe {
                    let _ = ffi::cancel_function_call(id);
                }
            })
        });

        Ok(StreamingCall::new(id, receiver, cancel_guard))
    }

    /// Call a function asynchronously (non-blocking)
    pub async fn call_function_async<T: BamlDecode>(
        &self,
        name: &str,
        args: &FunctionArgs,
    ) -> Result<T, BamlError> {
        let encoded = args.encode()?;
        let name_cstr =
            CString::new(name).map_err(|_| BamlError::internal("invalid function name"))?;

        let (id, receiver) = callbacks::create_async_callback();

        #[allow(unsafe_code)]
        let buf = unsafe {
            ffi::call_function_from_c(
                self.ptr,
                name_cstr.as_ptr(),
                encoded.as_ptr().cast::<c_char>(),
                encoded.len(),
                id,
            )
            .map_err(|e| {
                callbacks::remove_callback(id);
                BamlError::internal(format!("Failed to load BAML library: {e}"))
            })?
        };

        // Check for immediate error (decode Buffer response)
        ffi::decode_async_response(buf).map_err(|e| {
            callbacks::remove_callback(id);
            BamlError::internal(e)
        })?;

        // Set up cancellation callback if token provided
        // Guard is dropped when function returns, stopping the watcher
        let _cancel_guard = args.cancellation_token.as_ref().map(|token| {
            token.on_cancel(move || {
                #[allow(unsafe_code)]
                unsafe {
                    let _ = ffi::cancel_function_call(id);
                }
            })
        });

        // Await result (non-blocking)
        match receiver.recv().await {
            Ok(callbacks::CallbackResult::Final(data)) => {
                let holder = CffiValueHolder::decode(&data[..])
                    .map_err(|e| BamlError::internal(format!("decode error: {e}")))?;
                T::baml_decode(&holder)
            }
            Ok(callbacks::CallbackResult::Partial(_)) => Err(BamlError::internal(
                "unexpected partial result in async call",
            )),
            Ok(callbacks::CallbackResult::Error(e)) => Err(e),
            Err(_) => Err(BamlError::internal("callback channel closed")),
        }
    }

    /// Call a function with async streaming results
    ///
    /// If `args` contains an `on_tick` callback (set via `FunctionArgs::with_on_tick`),
    /// it will be invoked for each SSE streaming chunk received from the LLM.
    /// A collector is automatically created and injected when on_tick is present.
    pub fn call_function_stream_async<TPartial, TFinal>(
        &self,
        name: &str,
        args: &FunctionArgs,
    ) -> Result<AsyncStreamingCall<TPartial, TFinal>, BamlError>
    where
        TPartial: BamlDecode + Send + 'static,
        TFinal: Clone + BamlDecode + Send + 'static,
    {
        let on_tick_data = args.on_tick.as_ref().map(|cb| {
            let collector = self.new_collector("on-tick-collector");
            let data = callbacks::OnTickData {
                callback: cb.clone(),
                collector: collector.clone(),
            };
            (data, collector)
        });

        let extra_collector = on_tick_data.as_ref().map(|(_, c)| c);
        let encoded = args.encode_with_extra_collector(extra_collector)?;
        let name_cstr =
            CString::new(name).map_err(|_| BamlError::internal("invalid function name"))?;

        let (id, receiver) =
            callbacks::create_async_callback_with_on_tick(on_tick_data.map(|(d, _)| d));

        #[allow(unsafe_code)]
        let buf = unsafe {
            ffi::call_function_stream_from_c(
                self.ptr,
                name_cstr.as_ptr(),
                encoded.as_ptr().cast::<c_char>(),
                encoded.len(),
                id,
            )
            .map_err(|e| {
                callbacks::remove_callback(id);
                BamlError::internal(format!("Failed to load BAML library: {e}"))
            })?
        };

        // Check for immediate error (decode Buffer response)
        ffi::decode_async_response(buf).map_err(|e| {
            callbacks::remove_callback(id);
            BamlError::internal(e)
        })?;

        // Set up cancellation callback if token provided
        let cancel_guard = args.cancellation_token.as_ref().map(|token| {
            token.on_cancel(move || {
                #[allow(unsafe_code)]
                unsafe {
                    let _ = ffi::cancel_function_call(id);
                }
            })
        });

        Ok(AsyncStreamingCall::new(id, receiver, cancel_guard))
    }

    /// Parse raw LLM output into typed result
    ///
    /// Given the name of a BAML function and the raw text response from an LLM,
    /// this method parses the response according to the function's output type.
    ///
    /// # Arguments
    /// * `function_name` - Name of the BAML function that defines the output
    ///   type
    /// * `llm_response` - Raw text response from the LLM
    ///
    /// # Example
    /// ```ignore
    /// let raw_response = "Hello, World!";
    /// let result: String = runtime.parse("SayHello", raw_response)?;
    /// ```
    pub fn parse<T: BamlDecode>(
        &self,
        function_name: &str,
        llm_response: &str,
        stream: bool,
    ) -> Result<T, BamlError> {
        // Build args using FunctionArgs with parse-specific fields
        let args = FunctionArgs::new().arg("text", llm_response);
        let args = if stream {
            args.arg("stream", true)
        } else {
            args
        };
        let encoded = args.encode()?;
        let name_cstr = CString::new(function_name)
            .map_err(|_| BamlError::internal("invalid function name"))?;

        let (id, receiver) = callbacks::create_callback();

        #[allow(unsafe_code)]
        let buf = unsafe {
            ffi::call_function_parse_from_c(
                self.ptr,
                name_cstr.as_ptr(),
                encoded.as_ptr().cast::<c_char>(),
                encoded.len(),
                id,
            )
            .map_err(|e| {
                callbacks::remove_callback(id);
                BamlError::internal(format!("Failed to load BAML library: {e}"))
            })?
        };

        // Check for immediate error (decode Buffer response)
        ffi::decode_async_response(buf).map_err(|e| {
            callbacks::remove_callback(id);
            BamlError::internal(format!("function parse error: {e}"))
        })?;

        // Wait for result
        match receiver.recv() {
            Ok(callbacks::CallbackResult::Final(data)) => {
                if stream {
                    Err(BamlError::internal("unexpected final result in parse call"))
                } else {
                    let holder = CffiValueHolder::decode(&data[..])
                        .map_err(|e| BamlError::internal(format!("decode error: {e}")))?;
                    T::baml_decode(&holder)
                }
            }
            Ok(callbacks::CallbackResult::Partial(data)) => {
                if stream {
                    let holder = CffiValueHolder::decode(&data[..])
                        .map_err(|e| BamlError::internal(format!("decode error: {e}")))?;
                    T::baml_decode(&holder)
                } else {
                    Err(BamlError::internal(
                        "unexpected partial result in parse call",
                    ))
                }
            }
            Ok(callbacks::CallbackResult::Error(e)) => Err(e),
            Err(_) => Err(BamlError::internal("callback channel closed")),
        }
    }

    // =========================================================================
    // Build Request Methods
    // =========================================================================

    /// Build an HTTP request for a BAML function without executing it (sync, non-streaming).
    /// The `stream` arg should already be set in the FunctionArgs.
    pub fn build_request(&self, name: &str, args: &FunctionArgs) -> Result<HTTPRequest, BamlError> {
        self.build_request_inner(name, args)
    }

    /// Build an HTTP request for a streaming BAML function without executing it (sync).
    /// The `stream` arg should already be set in the FunctionArgs.
    pub fn build_request_stream(
        &self,
        name: &str,
        args: &FunctionArgs,
    ) -> Result<HTTPRequest, BamlError> {
        self.build_request_inner(name, args)
    }

    /// Build an HTTP request for a BAML function without executing it (async, non-streaming).
    /// The `stream` arg should already be set in the FunctionArgs.
    pub async fn build_request_async(
        &self,
        name: &str,
        args: &FunctionArgs,
    ) -> Result<HTTPRequest, BamlError> {
        self.build_request_inner_async(name, args).await
    }

    /// Build an HTTP request for a streaming BAML function without executing it (async).
    /// The `stream` arg should already be set in the FunctionArgs.
    pub async fn build_request_stream_async(
        &self,
        name: &str,
        args: &FunctionArgs,
    ) -> Result<HTTPRequest, BamlError> {
        self.build_request_inner_async(name, args).await
    }

    /// Internal sync implementation for build_request
    fn build_request_inner(
        &self,
        name: &str,
        args: &FunctionArgs,
    ) -> Result<HTTPRequest, BamlError> {
        let encoded = args.encode()?;
        let name_cstr =
            CString::new(name).map_err(|_| BamlError::internal("invalid function name"))?;

        let (id, receiver) = callbacks::create_callback();

        #[allow(unsafe_code)]
        let buf = unsafe {
            ffi::build_request_from_c(
                self.ptr,
                name_cstr.as_ptr(),
                encoded.as_ptr().cast::<c_char>(),
                encoded.len(),
                id,
            )
            .map_err(|e| {
                callbacks::remove_callback(id);
                BamlError::internal(format!("Failed to load BAML library: {e}"))
            })?
        };

        // Check for immediate error (decode Buffer response)
        ffi::decode_async_response(buf).map_err(|e| {
            callbacks::remove_callback(id);
            BamlError::internal(e)
        })?;

        // Wait for result
        match receiver.recv() {
            Ok(callbacks::CallbackResult::Final(data)) => {
                Self::decode_http_request_from_invocation_response(&data, self.ptr)
            }
            Ok(callbacks::CallbackResult::Partial(_)) => Err(BamlError::internal(
                "unexpected partial result in build_request call",
            )),
            Ok(callbacks::CallbackResult::Error(e)) => Err(e),
            Err(_) => Err(BamlError::internal("callback channel closed")),
        }
    }

    /// Internal async implementation for build_request
    async fn build_request_inner_async(
        &self,
        name: &str,
        args: &FunctionArgs,
    ) -> Result<HTTPRequest, BamlError> {
        let encoded = args.encode()?;
        let name_cstr =
            CString::new(name).map_err(|_| BamlError::internal("invalid function name"))?;

        let (id, receiver) = callbacks::create_async_callback();

        #[allow(unsafe_code)]
        let buf = unsafe {
            ffi::build_request_from_c(
                self.ptr,
                name_cstr.as_ptr(),
                encoded.as_ptr().cast::<c_char>(),
                encoded.len(),
                id,
            )
            .map_err(|e| {
                callbacks::remove_callback(id);
                BamlError::internal(format!("Failed to load BAML library: {e}"))
            })?
        };

        // Check for immediate error (decode Buffer response)
        ffi::decode_async_response(buf).map_err(|e| {
            callbacks::remove_callback(id);
            BamlError::internal(e)
        })?;

        // Await result (non-blocking)
        match receiver.recv().await {
            Ok(callbacks::CallbackResult::Final(data)) => {
                Self::decode_http_request_from_invocation_response(&data, self.ptr)
            }
            Ok(callbacks::CallbackResult::Partial(_)) => Err(BamlError::internal(
                "unexpected partial result in build_request async call",
            )),
            Ok(callbacks::CallbackResult::Error(e)) => Err(e),
            Err(_) => Err(BamlError::internal("callback channel closed")),
        }
    }

    /// Decode an InvocationResponse containing a BamlObjectHandle into an HTTPRequest
    fn decode_http_request_from_invocation_response(
        data: &[u8],
        runtime_ptr: *const c_void,
    ) -> Result<HTTPRequest, BamlError> {
        let response = InvocationResponse::decode(data)
            .map_err(|e| BamlError::internal(format!("decode InvocationResponse error: {e}")))?;

        match response.response {
            Some(InvResponse::Success(success)) => match success.result {
                Some(InvSuccessResult::Object(handle)) => {
                    HTTPRequest::decode_handle(handle, runtime_ptr)
                }
                other => Err(BamlError::internal(format!(
                    "expected object handle in InvocationResponse, got: {other:?}"
                ))),
            },
            Some(InvResponse::Error(msg)) => Err(BamlError::internal(msg)),
            None => Err(BamlError::internal(
                "empty response in InvocationResponse for build_request",
            )),
        }
    }

    // =========================================================================
    // Media Factory Methods
    // =========================================================================

    /// Create an Image from a URL
    pub fn new_image_from_url(&self, url: &str, mime_type: Option<&str>) -> Image {
        Image::from_url(self.ptr, url, mime_type)
    }

    /// Create an Image from base64-encoded data
    pub fn new_image_from_base64(&self, base64: &str, mime_type: Option<&str>) -> Image {
        Image::from_base64(self.ptr, base64, mime_type)
    }

    /// Create Audio from a URL
    pub fn new_audio_from_url(&self, url: &str, mime_type: Option<&str>) -> Audio {
        Audio::from_url(self.ptr, url, mime_type)
    }

    /// Create Audio from base64-encoded data
    pub fn new_audio_from_base64(&self, base64: &str, mime_type: Option<&str>) -> Audio {
        Audio::from_base64(self.ptr, base64, mime_type)
    }

    /// Create a PDF from a URL
    pub fn new_pdf_from_url(&self, url: &str, mime_type: Option<&str>) -> Pdf {
        Pdf::from_url(self.ptr, url, mime_type)
    }

    /// Create a PDF from base64-encoded data
    pub fn new_pdf_from_base64(&self, base64: &str, mime_type: Option<&str>) -> Pdf {
        Pdf::from_base64(self.ptr, base64, mime_type)
    }

    /// Create a Video from a URL
    pub fn new_video_from_url(&self, url: &str, mime_type: Option<&str>) -> Video {
        Video::from_url(self.ptr, url, mime_type)
    }

    /// Create a Video from base64-encoded data
    pub fn new_video_from_base64(&self, base64: &str, mime_type: Option<&str>) -> Video {
        Video::from_base64(self.ptr, base64, mime_type)
    }

    // =========================================================================
    // Collector Factory Methods
    // =========================================================================

    /// Create a new collector for telemetry
    pub fn new_collector(&self, name: &str) -> Collector {
        Collector::new(self.ptr, name)
    }

    // =========================================================================
    // TypeBuilder Factory Methods
    // =========================================================================

    /// Create a new `TypeBuilder` for dynamic type construction
    pub fn new_type_builder(&self) -> TypeBuilder {
        TypeBuilder::new(self.ptr)
    }
}

impl Drop for BamlRuntime {
    fn drop(&mut self) {
        #[allow(unsafe_code)]
        // Ignore errors during drop - the library should already be loaded at this point
        // and we can't do much about errors during cleanup anyway
        let _ = unsafe { ffi::destroy_baml_runtime(self.ptr) };
    }
}

/// Simple JSON encoding for maps
///
/// This is a minimal implementation to avoid adding `serde_json` as a
/// dependency. For simplicity, we assume keys and values don't contain
/// problematic characters that would require complex escaping beyond basic
/// escapes.
fn json_encode_map(map: &HashMap<String, String>) -> Result<String, BamlError> {
    serde_json::to_string(map)
        .map_err(|e| BamlError::internal(format!("failed to encode map: {e}")))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_json_encode_empty_map() {
        let map: HashMap<String, String> = HashMap::new();
        let result = json_encode_map(&map).unwrap();
        assert_eq!(result, "{}");
    }

    #[test]
    fn test_json_encode_simple_map() {
        let mut map = HashMap::new();
        map.insert("key".to_string(), "value".to_string());
        let result = json_encode_map(&map).unwrap();
        assert_eq!(result, "{\"key\":\"value\"}");
    }
}