use crate::stdlib::macros::{harn_builtin, register_builtin_defs, VmBuiltinDef};
use crate::value::{VmError, VmValue};
use crate::vm::Vm;

const SELF_CERTAINTY_PROMPT_PREFIX: &str = "Repeat exactly the text between <text> tags. Do not add, omit, or alter any characters.\n<text>\n";
const SELF_CERTAINTY_PROMPT_SUFFIX: &str = "\n</text>";

pub(crate) fn register_rerank_builtins(vm: &mut Vm) {
    register_builtin_defs(vm, RERANK_BUILTINS);
}

const RERANK_BUILTINS: &[&VmBuiltinDef] = &[&SELF_CERTAINTY_BUILTIN_DEF];

/// Return length-normalized confidence from token log probabilities.
#[harn_builtin(
    sig = "__llm_self_certainty(text_or_result: string|dict, options?: dict|nil) -> float",
    kind = "async",
    category = "llm.rerank"
)]
async fn self_certainty_builtin(
    _ctx: crate::vm::AsyncBuiltinCtx,
    args: Vec<VmValue>,
) -> Result<VmValue, VmError> {
    let text = match args.first() {
        Some(VmValue::String(text)) => text.to_string(),
        Some(VmValue::Dict(dict)) => {
            if let Some(logprobs) = dict.get("logprobs") {
                let values = vm_logprobs(logprobs)?;
                return Ok(VmValue::Float(score_logprobs(&values)?));
            }
            return Err(VmError::Runtime(
                "__llm_self_certainty: first argument must be text or an llm result with logprobs"
                    .to_string(),
            ));
        }
        Some(other) => {
            return Err(VmError::Runtime(format!(
                "__llm_self_certainty: text must be a string, got {}",
                other.type_name()
            )))
        }
        None => String::new(),
    };

    let options = args
        .get(1)
        .and_then(VmValue::as_dict)
        .cloned()
        .unwrap_or_default();
    if let Some(value) = options.get("logprobs") {
        if !matches!(value, VmValue::Bool(_) | VmValue::Nil) {
            let values = vm_logprobs(value)?;
            return Ok(VmValue::Float(score_logprobs(&values)?));
        }
    }
    if text.trim().is_empty() {
        return Err(VmError::Runtime(
            "__llm_self_certainty: text must be non-empty when logprobs are not supplied"
                .to_string(),
        ));
    }

    let mut call_options = options;
    call_options.insert("logprobs".to_string(), VmValue::Bool(true));
    call_options
        .entry("stream".to_string())
        .or_insert(VmValue::Bool(false));
    call_options
        .entry("temperature".to_string())
        .or_insert(VmValue::Float(0.0));
    call_options
        .entry("max_tokens".to_string())
        .or_insert(VmValue::Int(estimate_echo_tokens(&text)));

    let prompt = format!("{SELF_CERTAINTY_PROMPT_PREFIX}{text}{SELF_CERTAINTY_PROMPT_SUFFIX}");
    let opts = super::helpers::extract_llm_options(&[
        VmValue::String(std::sync::Arc::from(prompt)),
        VmValue::Nil,
        VmValue::Dict(std::sync::Arc::new(call_options)),
    ])?;
    let result = super::api::vm_call_llm_full(&opts).await?;
    if result.logprobs.is_empty() {
        return Err(VmError::Runtime(format!(
            "__llm_self_certainty: provider '{}' model '{}' did not return token logprobs",
            result.provider, result.model
        )));
    }
    Ok(VmValue::Float(score_logprobs(&result.logprobs)?))
}

fn estimate_echo_tokens(text: &str) -> i64 {
    ((text.chars().count() as i64 + 2) / 3 + 8).clamp(1, 4096)
}

fn vm_logprobs(value: &VmValue) -> Result<Vec<serde_json::Value>, VmError> {
    match value {
        VmValue::List(items) => Ok(items.iter().map(super::helpers::vm_value_to_json).collect()),
        VmValue::Dict(dict) => {
            if let Some(logprobs) = dict.get("logprobs") {
                return vm_logprobs(logprobs);
            }
            Ok(vec![serde_json::Value::Object(
                dict.iter()
                    .map(|(key, value)| (key.clone(), super::helpers::vm_value_to_json(value)))
                    .collect(),
            )])
        }
        other => Err(VmError::Runtime(format!(
            "__llm_self_certainty: logprobs must be a list or dict, got {}",
            other.type_name()
        ))),
    }
}

fn score_logprobs(entries: &[serde_json::Value]) -> Result<f64, VmError> {
    let mut sum = 0.0;
    let mut count = 0usize;
    for entry in entries {
        if let Some(logprob) = logprob_value(entry) {
            if logprob.is_finite() {
                sum += logprob;
                count += 1;
            }
        }
    }
    if count == 0 {
        return Err(VmError::Runtime(
            "__llm_self_certainty: no finite token logprobs were available".to_string(),
        ));
    }
    let mean = sum / count as f64;
    Ok(mean.exp().clamp(0.0, 1.0))
}

fn logprob_value(entry: &serde_json::Value) -> Option<f64> {
    match entry {
        serde_json::Value::Number(number) => number.as_f64(),
        serde_json::Value::Object(object) => object
            .get("logprob")
            .or_else(|| object.get("token_logprob"))
            .and_then(|value| value.as_f64()),
        _ => None,
    }
}

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

    #[test]
    fn score_logprobs_uses_length_normalized_probability() {
        let ln_half = -std::f64::consts::LN_2;
        let entries = vec![
            serde_json::json!({"token": "a", "logprob": ln_half}),
            serde_json::json!({"token": "b", "logprob": ln_half}),
        ];
        let score = score_logprobs(&entries).expect("score");
        assert!((score - 0.5).abs() < 0.000001);
    }
}
