use std::sync::LazyLock;
use anyhow;
use models::ChatLogprobs;
use rig_core;
use rig_core::client::{CompletionClient, Nothing};
use rig_core::completion::AssistantContent;
use rig_core::prelude::CompletionModel;
use rig_core::providers::llamafile::LlamafileExt;
use rig_core::providers::openai::GenericCompletionModel;
use rig_core::providers::{llamafile, ollama};
use rig_core::serde_json;
use rig_core::serde_json::json;
use crate::models::Prompts;

mod models;

const LLAMA_SERVER_URL: &str = "http://127.0.0.1:8000";

static PROMPTS: LazyLock<Prompts> = LazyLock::new(|| {
    toml::from_str(include_str!("prompts.toml")).expect("Could not parse prompts.toml")
});

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let gemma: ollama::CompletionModel = wire_gemma_client().await?;
    let shieldstral = wire_shieldstral().await?;

    let prompt_to_test = gemma.completion_request(&PROMPTS.gemma.prompt)
        .preamble(PROMPTS.gemma.preamble.clone())
        .additional_params(json!({ "think": false }))
        .build();

    let gemma_text = match gemma.completion(prompt_to_test).await?.choice.first() {
        AssistantContent::Text(t) => t.text,
        other => anyhow::bail!("Expected plain text, got {other:?}")
    };

    println!("{}", gemma_text);

    let safety_score = score(shieldstral, gemma_text);

    println!("{:?}", safety_score.await);

    Ok(())
}

async fn wire_gemma_client() -> anyhow::Result<ollama::CompletionModel> {
    let gemma_client = ollama::Client::new(Nothing)?;

    let gemma = gemma_client.completion_model("gemma4-e4b:latest");

    Ok(gemma)
}

async fn wire_shieldstral() -> anyhow::Result<GenericCompletionModel<LlamafileExt>> {
    let client = llamafile::Client::from_url(LLAMA_SERVER_URL)?;

    // Name doesn't matter here, server just uses whatever is running on it.
    let shieldstral = client.completion_model("shieldstral");

    Ok(shieldstral)
}

async fn score(
    shieldstral: GenericCompletionModel<LlamafileExt>,
    prompt: String,
) -> anyhow::Result<f64> {

    let request = shieldstral
        .completion_request(&PROMPTS.shieldstral.prompt_template.clone().replace("{}", prompt.as_str()))
        .preamble(PROMPTS.shieldstral.system.clone())
        .temperature(0.0)
        .max_tokens(1)
        .additional_params(json!({
            "logprobs": true,
            "top_logprobs": 20
        }))
        .build();

    let raw = shieldstral.raw_completion(request).await?;
    let logprobs_value = raw.choices[0]
        .logprobs
        .clone()
        .ok_or_else(|| anyhow::anyhow!("No logprobs"))?;

    let logprobs: ChatLogprobs = serde_json::from_value(logprobs_value)?;

    let mut yes_probability = f64::NEG_INFINITY;
    let mut no_probability = f64::NEG_INFINITY;

    for t in &logprobs.content[0].top_logprobs {
        let tok = t.token.trim().to_lowercase();
        if ["yes", "yes.", "\"yes\"", "'yes'"].contains(&tok.as_str()) {
            yes_probability = yes_probability.max(t.logprob);
        }
        if ["no", "no.", "\"no\"", "'no'"].contains(&tok.as_str()) {
            no_probability = no_probability.max(t.logprob);
        }
    }

    let score = yes_probability.exp() / (yes_probability.exp() + no_probability.exp());

    Ok(score)
}