use hiramu::bedrock::model_info::{ModelInfo, ModelName};
use hiramu::bedrock::models::mistral::mistral_client::{MistralClient, MistralOptions};
use hiramu::bedrock::models::mistral::mistral_request_message::MistralRequestBuilder;

async fn generating_text_with_mistral() {
    let mistral_options = MistralOptions::new()
        .profile_name("bedrock")
        .region("us-west-2");

    let client = MistralClient::new(mistral_options).await;

    let request =
        MistralRequestBuilder::new("<s>[INST] What is the capital of France?[/INST]".to_string())
            .max_tokens(200)
            .temperature(0.8)
            .build();

    let model_id = ModelInfo::from_model_name(ModelName::MistralMixtral8X7BInstruct0x);
    let response = client.generate(model_id, &request).await.unwrap();

    println!("Response: {:?}", response.outputs[0].text);
}

use futures::stream::StreamExt;
use hiramu::bedrock::models::mistral::mistral_client::{MistralClient, MistralOptions};
use hiramu::bedrock::models::mistral::mistral_request_message::MistralRequestBuilder;
use hiramu::bedrock::model_info::{ModelInfo, ModelName};

pub async fn generating_text_with_mistral() {
    let mistral_options = MistralOptions::new()
        .profile_name("bedrock")
        .region("us-west-2");

    let client = MistralClient::new(mistral_options).await;

    let request = MistralRequestBuilder::new("<s>[INST] What is the capital of France?[/INST]".to_string())
        .max_tokens(200)
        .temperature(0.8)
        .build();

    let model_id = ModelInfo::from_model_name(ModelName::MistralMixtral8X7BInstruct0x);
    let mut stream = client.generate_with_stream(model_id, &request).await.unwrap();

    while let Some(result) = stream.next().await {
        match result {
            Ok(response) => {
                println!("Response: {:?}", response.outputs[0].text);
            }
            Err(err) => {
                eprintln!("Error: {:?}", err);
            }
        }
    }
}

use std::io::Write;

use futures::TryStreamExt;

use hiramu::ollama::ollama_client::OllamaClient;
use hiramu::ollama::model::{GenerateRequestBuilder};

async fn generating_text_with_ollama() {
    let client = OllamaClient::new("http://localhost:11434".to_string());
    
    let request = GenerateRequestBuilder::new("mistral".to_string())
        .prompt("Once upon a time".to_string())
        .build();

    let response_stream = client.generate(request).await.unwrap();

    response_stream
        .try_for_each(|chunk| async move {
            print!("{}", chunk.response);
            std::io::stdout().flush()?;
            Ok(())
        })
        .await
        .unwrap();
}

use futures::TryStreamExt;
use std::io::{self, Write};

use hiramu::ollama::{ChatRequestBuilder, Message, OllamaClient, OllamaError, OptionsBuilder};

async fn demo_chat_with_ollama_with_stream() -> Result<(), OllamaError> {
    let client = OllamaClient::new("http://localhost:11434".to_string());

    let messages = vec![Message::new(
        "user".to_string(),
        "What is the capital of France?  "
            .to_string(),
    )];

    let options = OptionsBuilder::new()
        .num_predict(100) // Limit the number of predicted tokens
        .temperature(0.4);

    let request = ChatRequestBuilder::new("mistral".to_string())
        .messages(messages.to_owned())
        .options_from_builder(options)
        .build();

    let response_stream = client.chat(request).await?;

    let result = response_stream
        .try_for_each(|chunk| async {
            let message = chunk.message;
            print!("{}", message.content);
            // Flush the output to ensure the prompt is displayed.
            io::stdout().flush().unwrap();
            Ok(())
        })
        .await;

    result
}

use std::io::Write;

use futures::TryStreamExt;

use hiramu::bedrock::model_info::{ModelInfo, ModelName};
use hiramu::bedrock::models::claude::claude_client::{ClaudeClient, ClaudeOptions};
use hiramu::bedrock::models::claude::claude_request_message::{
    ChatOptions, ContentBlockDelta, ConversationRequest, Message, StreamResultData,
};

pub async fn chat_with_claude() {
    let claude_options = ClaudeOptions::new()
        .profile_name("bedrock")
        .region("us-west-2");

    let client = ClaudeClient::new(claude_options).await;

    let mut conversation_request = ConversationRequest::default();
    conversation_request
        .messages
        .push(Message::new_user_message("Hello, Claude!".to_owned()));

    let chat_options = ChatOptions::default()
        .with_temperature(0.7)
        .with_max_tokens(100)
        .with_model_id(ModelInfo::from_model_name(
            ModelName::AnthropicClaudeHaiku1x,
        ));

    let response_stream = client
        .chat_with_stream(&conversation_request, &chat_options)
        .await
        .unwrap();

    response_stream
        .try_for_each(|chunk| async move {
            match chunk {
                StreamResultData::ContentBlockStart(..) => {
                    println!("\n------------------------------");
                }
                StreamResultData::ContentBlockStop(..) => {
                    println!("\n------------------------------");
                }
                StreamResultData::ContentBlockDelta(ContentBlockDelta { delta, .. }) => {
                    print!("{}", delta.text);
                    std::io::stdout().flush().unwrap();
                }
                _ => {}
            }
            Ok(())
        })
        .await
        .unwrap();
}

use std::io::Write;

use futures::TryStreamExt;

use hiramu::bedrock::models::claude::claude_client::{ClaudeClient, ClaudeOptions};
use hiramu::bedrock::models::claude::claude_request_message::{ChatOptions, ContentBlockDelta, ConversationRequest, Message, StreamResultData};
use hiramu::fetch_and_base64_encode_image;

async fn image_with_claude() {
    let claude_options = ClaudeOptions::new()
        .profile_name("bedrock")
        .region("us-west-2");

    let client = ClaudeClient::new(claude_options).await;

    let image_url = "./data/mario.png";
    let input_text = "What's in this image?".to_string();
    let image = fetch_and_base64_encode_image(image_url).await.unwrap().to_string();
    let mime_type = "image/png".to_string();

    let message = Message::new_user_message_with_image(&input_text, &image, &mime_type);

    let mut conversation_request = ConversationRequest::default();
    conversation_request.messages.push(message);

    let chat_options = ChatOptions::default()
        .with_temperature(0.7)
        .with_max_tokens(100);

    let response_stream = client
        .chat_with_stream(&conversation_request, &chat_options)
        .await
        .unwrap();

        response_stream
        .try_for_each(|chunk| async move {
            match chunk {
                StreamResultData::ContentBlockStart(..) => {
                    println!("\n------------------------------");
                }
                StreamResultData::ContentBlockStop(..) => {
                    println!("\n------------------------------");
                }

                StreamResultData::ContentBlockDelta(ContentBlockDelta { delta, .. }) => {
                    print!("{}", delta.text);
                    std::io::stdout().flush().unwrap();
                }
                _ => {}
            }
            Ok(())
        })
        .await
        .unwrap();
}

use hiramu::bedrock::bedrock_client::{BedrockClient, BedrockClientOptions};
use hiramu::bedrock::model_info::{ModelInfo, ModelName};

#[tokio::main]
async fn main() {
    let model_id = ModelInfo::from_model_name(ModelName::AnthropicClaudeHaiku1x);
    let profile_name = "bedrock";
    let region = "us-west-2";

    let prompt = "Hi. In a short paragraph, explain what you can do.";

    let payload = serde_json::json!({
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 1000,
        "messages": [{
            "role": "user",
            "content": [{
                "type": "text",
                "text": prompt
            }]
        }]
    });

    let options = BedrockClientOptions::new()
        .profile_name(profile_name)
        .region(region);

    let client = BedrockClient::new(options).await;

    let result = client
        .generate_raw(model_id.to_string(), payload)
        .await
        .unwrap();

    println!("{:?}", result);
}

use futures::TryStreamExt;
use hiramu::bedrock::bedrock_client::{BedrockClient, BedrockClientOptions};
use hiramu::bedrock::model_info::{ModelInfo, ModelName};

#[tokio::main]
async fn main() {
    let model_id = ModelInfo::from_model_name(ModelName::AnthropicClaudeHaiku1x);
    let profile_name = "bedrock";
    let region = "us-west-2";

    let prompt = "Hi. In a short paragraph, explain what you can do.";

    let payload = serde_json::json!({
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 1000,
        "messages": [{
            "role": "user",
            "content": [{
                "type": "text",
                "text": prompt
            }]
        }]
    });

    let options = BedrockClientOptions::new()
        .profile_name(profile_name)
        .region(region);

    let client = BedrockClient::new(options).await;

    let stream = client
        .generate_raw_stream(model_id.to_string(), payload)
        .await
        .unwrap();

    stream
        .try_for_each(|chunk| async move {
            println!("{:?}", chunk);
            Ok(())
        })
        .await
        .unwrap();
}

use hiramu::ollama::{EmbeddingsRequestBuilder, OllamaClient};

pub async fn demo_ollama_embedding() -> Result<(), Box<dyn std::error::Error>> {
    let client = OllamaClient::new("http://localhost:11434".to_string());

    let prompt = "The quick brown fox jumps over the lazy dog.";

    let request = EmbeddingsRequestBuilder::new("nomic-embed-text".to_string(), prompt.to_string())
        .keep_alive("10m".to_string())
        .build();

    match client.embeddings(request).await {
        Ok(response) => {
            // Print embeddings dimensions
            println!("Embeddings dimensions: {:?}", response.embedding.len());
            println!("Embeddings: {:?}", response);
        }
        Err(error) => {
            eprintln!("Error: {:?}", error);
        }
    }

    Ok(())
}