#![doc = include_str!("../README.MD")]
#![warn(missing_docs)]

use futures::Stream;
use std::future::Future;

pub mod prelude;

#[cfg_attr(not(any(feature = "ollama", feature = "llama.cpp")), deny(warnings))]
mod error;
mod message;
mod prompts;

#[cfg(feature = "candle")]
pub mod candle;

#[cfg(feature = "llama.cpp")]
pub mod llama_cpp;

#[cfg(feature = "ollama")]
pub mod ollama;

#[cfg(feature = "simple-api")]
pub mod simple_api;

#[cfg(feature = "template")]
pub mod template;

pub use error::Error;
pub use message::{Message, Role};
pub use prompts::{ChatPrompt, GenerationPrompt};

/// Export Result type.
pub type Result<T, E = Error> = std::result::Result<T, E>;

/// Format
#[derive(Debug, Clone, Copy)]
pub enum Format
{
  /// Output text
  Text,
  /// Output Json
  Json,
}

/// String stream
pub type StringStream = ccutils::futures::BoxedStream<Result<String>>;

/// Convenient function for coercing a box pin stream to a StringStream
#[allow(dead_code)]
pub(crate) fn pin_stream<T: 'static + Send + Stream<Item = Result<String>>>(t: T) -> StringStream
{
  Box::pin(t)
}

fn accumulate<T>(stream_maker: T) -> Result<impl Future<Output = Result<String>> + Send>
where
  T: Future<Output = Result<StringStream>> + Send,
{
  use futures::stream::StreamExt;
  Ok(async {
    let mut result: String = Default::default();
    let mut stream = Box::pin(stream_maker.await?);
    while let Some(next_token) = stream.next().await
    {
      if result.is_empty()
      {
        result = next_token?;
      }
      else
      {
        result.push_str(&next_token?);
      }
    }
    Ok(result)
  })
}

#[allow(dead_code)]
pub(crate) fn generate_with_chat<LLM>(
  llm: &LLM,
  prompt: GenerationPrompt,
) -> Result<impl Future<Output = Result<StringStream>> + Send + use<'_, LLM>>
where
  LLM: LargeLanguageModel,
{
  let chat_prompt = ChatPrompt::new()
    .system_opt(prompt.system)
    .assistant_opt(prompt.assistant)
    .user(prompt.user)
    .format(prompt.format);
  llm.chat_stream(chat_prompt)
}

/// LLM
pub trait LargeLanguageModel
{
  /// Chat with a model, returning a stream
  fn chat_stream(
    &self,
    prompt: ChatPrompt,
  ) -> Result<impl Future<Output = Result<StringStream>> + Send>;
  /// Run chat on a model, return once the complete answer has been computed.
  /// The default implementation call `chat_stream` until completion of the stream.
  fn chat(&self, prompt: ChatPrompt) -> Result<impl Future<Output = Result<String>> + Send>
  {
    let stream = self.chat_stream(prompt)?;
    accumulate(stream)
  }
  /// Run inference on a model, returning a stream
  fn generate_stream(
    &self,
    prompt: GenerationPrompt,
  ) -> Result<impl Future<Output = Result<StringStream>> + Send>;
  /// Run inference on a model, return once the complete answer has been computed.
  /// The default implementation call `infer_stream` until completion of the stream.
  fn generate(
    &self,
    prompt: GenerationPrompt,
  ) -> Result<impl Future<Output = Result<String>> + Send>
  {
    let stream = self.generate_stream(prompt)?;
    accumulate(stream)
  }
}
