use std::io::Write as _;

use kproc_llm::prelude::*;

trait Example
{
  type Model: kproc_llm::LargeLanguageModel + Send + Sync;
  async fn create_model() -> Self::Model;
  async fn run(prompt: String) -> kproc_llm::Result<()>
  {
    use kproc_llm::LargeLanguageModel;
    let model = Self::create_model().await;
    let prompt = kproc_llm::GenerationPrompt::prompt(prompt);
    let mut stream = model.generate_stream(prompt)?.await?;

    use futures::stream::StreamExt;
    print!("Answer: ```");
    while let Some(next_token) = stream.next().await
    {
      print!("{}", next_token?);
      let _ = std::io::stdout().flush();
    }
    println!("```");
    Ok(())
  }
}

#[cfg(feature = "ollama")]
struct OllamaSimple {}

#[cfg(feature = "ollama")]
impl Example for OllamaSimple
{
  type Model = kproc_llm::ollama::Ollama;
  async fn create_model() -> Self::Model
  {
    ollama::Ollama::from_model("llama3:instruct")
  }
}

#[cfg(feature = "llama.cpp")]
struct LlamaCppSimple {}

#[cfg(feature = "llama.cpp")]
impl Example for LlamaCppSimple
{
  type Model = kproc_llm::llama_cpp::LlamaCpp;
  async fn create_model() -> Self::Model
  {
    llama_cpp::LlamaCpp::from_model(
      "hf://ggml-org/Meta-Llama-3.1-8B-Instruct-Q4_0-GGUF/meta-llama-3.1-8b-instruct-q4_0.gguf",
    )
    .unwrap()
  }
}

#[cfg(feature = "simple-api")]
struct SimpleApi {}

#[cfg(feature = "simple-api")]
impl Example for SimpleApi
{
  type Model = kproc_llm::simple_api::SimpleApi<yaaral::tokio::Runtime>;
  async fn create_model() -> Self::Model
  {
    simple_api::SimpleApi::new(
      yaaral::tokio::Runtime::current(),
      "http://localhost",
      8080,
      None,
    )
    .unwrap()
  }
}

#[cfg(feature = "candle")]
struct CandleLlama;

#[cfg(feature = "candle")]
impl Example for CandleLlama
{
  type Model = kproc_llm::candle::Candle;
  async fn create_model() -> Self::Model
  {
    kproc_llm::candle::factory::llama2_7b()
      .build()
      .await
      .unwrap()
  }
}

#[cfg(feature = "candle")]
struct CandleSmol;

#[cfg(feature = "candle")]
impl Example for CandleSmol
{
  type Model = kproc_llm::candle::Candle;
  async fn create_model() -> Self::Model
  {
    kproc_llm::candle::factory::smol_lm_2_1_7b()
      .build()
      .await
      .unwrap()
  }
}
enum Backend
{
  #[cfg(feature = "ollama")]
  Ollama,
  #[cfg(feature = "llama.cpp")]
  LlamaCpp,
  #[cfg(feature = "simple-api")]
  SimpleApi,
  #[cfg(feature = "candle")]
  CandleLlama,
  #[cfg(feature = "candle")]
  CandleSmol,
  None,
}

#[tokio::main]
async fn main()
{
  colog::init();
  let mut backend = Backend::None;
  let mut prompt = String::new();
  for argument in std::env::args()
  {
    if argument == "--llama-cpp"
    {
      #[cfg(feature = "llama.cpp")]
      {
        backend = Backend::LlamaCpp;
      }
    }
    else if argument == "--ollama"
    {
      #[cfg(feature = "ollama")]
      {
        backend = Backend::Ollama;
      }
    }
    else if argument == "--simple-api"
    {
      #[cfg(feature = "simple-api")]
      {
        backend = Backend::SimpleApi;
      }
    }
    else if argument == "--candle-llama"
    {
      #[cfg(feature = "candle")]
      {
        backend = Backend::CandleLlama;
      }
    }
    else if argument == "--candle-smol"
    {
      #[cfg(feature = "candle")]
      {
        backend = Backend::CandleSmol;
      }
    }
    else
    {
      prompt = argument;
    }
  }
  match backend
  {
    #[cfg(feature = "ollama")]
    Backend::Ollama => OllamaSimple::run(prompt).await.unwrap(),
    #[cfg(feature = "llama.cpp")]
    Backend::LlamaCpp => LlamaCppSimple::run(prompt).await.unwrap(),
    #[cfg(feature = "simple-api")]
    Backend::SimpleApi => SimpleApi::run(prompt).await.unwrap(),
    #[cfg(feature = "candle")]
    Backend::CandleLlama => CandleLlama::run(prompt).await.unwrap(),
    #[cfg(feature = "candle")]
    Backend::CandleSmol => CandleSmol::run(prompt).await.unwrap(),
    Backend::None => log::error!(
      "Select a backend with --llama-cpp --ollama --simple-api --candle-llama --candle-smol"
    ),
  }
}
