logprobe 0.4.0

use serde_json::Value;
use std::io::{self, Read};
use thiserror::Error;

use crate::types::{InputFormat, LogprobSequence, TokenLogprob, TopKEntry};

/// Errors returned while parsing logprob input.
///
/// This is a concrete error type so library consumers can match on the failure
/// mode programmatically rather than depending on `anyhow`.
#[derive(Debug, Error)]
pub enum ParseError {
    /// The input was empty or contained only whitespace.
    #[error("empty input")]
    EmptyInput,
    /// The input format could not be determined from its structure.
    #[error("could not detect input format: {0}")]
    UnknownFormat(String),
    /// The input was not valid JSON (or JSONL).
    #[error("invalid JSON: {0}")]
    InvalidJson(String),
    /// A required field was absent (the message names the JSON path or field).
    #[error("missing field: {0}")]
    MissingField(String),
    /// The input parsed as JSON but was structurally inconsistent
    /// (for example, mismatched array lengths).
    #[error("malformed input: {0}")]
    Malformed(String),
    /// Reading from the input reader failed.
    #[error("I/O error: {0}")]
    Io(#[from] io::Error),
}

/// Parse logprobs from a reader, auto-detecting format unless overridden.
///
/// # Errors
///
/// Returns [`ParseError`] if the reader cannot be read, the input is empty,
/// the format cannot be detected, or the JSON is malformed.
pub fn parse_input(
    reader: impl Read,
    format_override: Option<InputFormat>,
    strict: bool,
) -> Result<LogprobSequence, ParseError> {
    let mut buf = String::new();
    let mut reader = io::BufReader::new(reader);
    reader.read_to_string(&mut buf)?;
    parse_string(&buf, format_override, strict)
}

/// Parse logprobs from a string, auto-detecting format unless overridden.
///
/// # Errors
///
/// Returns [`ParseError`] if the input is empty, the format cannot be detected
/// (or is ambiguous under `strict`), or the JSON is malformed.
pub fn parse_string(
    input: &str,
    format_override: Option<InputFormat>,
    strict: bool,
) -> Result<LogprobSequence, ParseError> {
    let trimmed = input.trim();
    if trimmed.is_empty() {
        return Err(ParseError::EmptyInput);
    }

    let format = match format_override {
        Some(f) => f,
        None => detect_format(trimmed, strict)?,
    };

    match format {
        InputFormat::OpenAI => parse_openai(trimmed),
        InputFormat::VllmFlat => parse_vllm(trimmed),
        InputFormat::JsonlStream => parse_jsonl(trimmed),
        InputFormat::Gemini => parse_gemini(trimmed),
        InputFormat::Ollama => parse_ollama(trimmed),
    }
}

/// Detect the input format from content structure.
fn detect_format(input: &str, strict: bool) -> Result<InputFormat, ParseError> {
    // Try parsing as JSON first
    if let Ok(val) = serde_json::from_str::<Value>(input) {
        return detect_format_json(&val, strict);
    }

    // Check if it looks like JSONL (multiple lines, each valid JSON)
    let lines: Vec<&str> = input.lines().filter(|l| !l.trim().is_empty()).collect();
    if !lines.is_empty()
        && lines
            .iter()
            .all(|l| serde_json::from_str::<Value>(l).is_ok())
    {
        return Ok(InputFormat::JsonlStream);
    }

    Err(ParseError::UnknownFormat(
        "not valid JSON or JSONL".to_string(),
    ))
}

fn detect_format_json(val: &Value, strict: bool) -> Result<InputFormat, ParseError> {
    // OpenAI: choices[0].logprobs.content is an array of objects with "token" + "logprob"
    if let Some(content) = val
        .pointer("/choices/0/logprobs/content")
        .and_then(|v| v.as_array())
        && content
            .first()
            .map(|c| c.get("token").is_some() && c.get("logprob").is_some())
            .unwrap_or(false)
    {
        return Ok(InputFormat::OpenAI);
    }

    // vLLM/Together: choices[0].logprobs has "tokens" array and "token_logprobs" array
    if let Some(logprobs) = val.pointer("/choices/0/logprobs")
        && logprobs.get("tokens").and_then(|v| v.as_array()).is_some()
        && logprobs
            .get("token_logprobs")
            .and_then(|v| v.as_array())
            .is_some()
    {
        return Ok(InputFormat::VllmFlat);
    }

    // Gemini: candidates[0].logprobsResult with topCandidates and/or chosenCandidates
    if val
        .pointer("/candidates/0/logprobsResult")
        .and_then(|v| v.as_object())
        .is_some()
    {
        return Ok(InputFormat::Gemini);
    }

    // Ollama: top-level "logprobs" array with token+logprob objects (not nested in choices)
    if let Some(logprobs) = val.get("logprobs").and_then(|v| v.as_array())
        && logprobs
            .first()
            .map(|v| v.get("token").is_some() && v.get("logprob").is_some())
            .unwrap_or(false)
        && val.get("choices").is_none()
    {
        return Ok(InputFormat::Ollama);
    }

    if strict {
        return Err(ParseError::UnknownFormat(
            "--strict-format: could not unambiguously detect format from JSON structure"
                .to_string(),
        ));
    }

    // Fallback: if it's an array of objects with token+logprob, treat as JSONL-style
    if let Some(arr) = val.as_array()
        && arr
            .first()
            .map(|v| v.get("token").is_some() && v.get("logprob").is_some())
            .unwrap_or(false)
    {
        return Ok(InputFormat::JsonlStream);
    }

    Err(ParseError::UnknownFormat(
        "JSON structure does not match any supported format".to_string(),
    ))
}

/// Parse OpenAI Chat Completions format.
fn parse_openai(input: &str) -> Result<LogprobSequence, ParseError> {
    let val: Value =
        serde_json::from_str(input).map_err(|e| ParseError::InvalidJson(e.to_string()))?;
    let model = val.get("model").and_then(|m| m.as_str()).map(String::from);

    let content = val
        .pointer("/choices/0/logprobs/content")
        .and_then(|v| v.as_array())
        .ok_or_else(|| ParseError::MissingField("choices[0].logprobs.content".to_string()))?;

    let mut tokens = Vec::with_capacity(content.len());
    let mut total_logprob = 0.0;

    for item in content {
        let token = item
            .get("token")
            .and_then(|v| v.as_str())
            .ok_or_else(|| ParseError::MissingField("token".to_string()))?
            .to_string();
        let logprob = item
            .get("logprob")
            .and_then(|v| v.as_f64())
            .ok_or_else(|| ParseError::MissingField("logprob".to_string()))?;

        let bytes = item.get("bytes").and_then(|v| v.as_array()).map(|arr| {
            arr.iter()
                .filter_map(|b| b.as_u64().map(|n| n as u8))
                .collect()
        });

        let top_logprobs = item
            .get("top_logprobs")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|entry| {
                        let t = entry.get("token")?.as_str()?.to_string();
                        let lp = entry.get("logprob")?.as_f64()?;
                        Some(TopKEntry {
                            token: t,
                            logprob: lp,
                        })
                    })
                    .collect()
            });

        total_logprob += logprob;
        tokens.push(TokenLogprob {
            token,
            logprob,
            bytes,
            top_logprobs,
        });
    }

    Ok(LogprobSequence {
        tokens,
        model,
        format_detected: InputFormat::OpenAI.to_string(),
        total_logprob,
    })
}

/// Parse vLLM / Together flat format.
fn parse_vllm(input: &str) -> Result<LogprobSequence, ParseError> {
    let val: Value =
        serde_json::from_str(input).map_err(|e| ParseError::InvalidJson(e.to_string()))?;
    let model = val.get("model").and_then(|m| m.as_str()).map(String::from);

    let logprobs_obj = val
        .pointer("/choices/0/logprobs")
        .ok_or_else(|| ParseError::MissingField("choices[0].logprobs".to_string()))?;

    let token_strs = logprobs_obj
        .get("tokens")
        .and_then(|v| v.as_array())
        .ok_or_else(|| ParseError::MissingField("tokens".to_string()))?;

    let token_lps = logprobs_obj
        .get("token_logprobs")
        .and_then(|v| v.as_array())
        .ok_or_else(|| ParseError::MissingField("token_logprobs".to_string()))?;

    let top_lps = logprobs_obj.get("top_logprobs").and_then(|v| v.as_array());

    if token_strs.len() != token_lps.len() {
        return Err(ParseError::Malformed(
            "tokens and token_logprobs arrays have different lengths".to_string(),
        ));
    }

    let mut tokens = Vec::with_capacity(token_strs.len());
    let mut total_logprob = 0.0;

    for (i, (ts, tl)) in token_strs.iter().zip(token_lps.iter()).enumerate() {
        let token = ts
            .as_str()
            .ok_or_else(|| ParseError::Malformed(format!("token at index {i} is not a string")))?
            .to_string();
        let logprob = tl.as_f64().ok_or_else(|| {
            ParseError::Malformed(format!("logprob at index {i} is not a number"))
        })?;

        let top_logprobs = top_lps.and_then(|arr| {
            arr.get(i).and_then(|v| v.as_object()).map(|obj| {
                let mut entries: Vec<TopKEntry> = obj
                    .iter()
                    .map(|(k, v)| TopKEntry {
                        token: k.clone(),
                        logprob: v.as_f64().unwrap_or(f64::NEG_INFINITY),
                    })
                    .collect();
                // JSON objects have arbitrary key order — sort descending
                entries.sort_by(|a, b| {
                    b.logprob
                        .partial_cmp(&a.logprob)
                        .unwrap_or(std::cmp::Ordering::Equal)
                });
                entries
            })
        });

        total_logprob += logprob;
        tokens.push(TokenLogprob {
            token,
            logprob,
            bytes: None,
            top_logprobs,
        });
    }

    Ok(LogprobSequence {
        tokens,
        model,
        format_detected: InputFormat::VllmFlat.to_string(),
        total_logprob,
    })
}

/// Parse JSONL stream format (one {token, logprob} per line).
fn parse_jsonl(input: &str) -> Result<LogprobSequence, ParseError> {
    let mut tokens = Vec::new();
    let mut total_logprob = 0.0;

    // Handle both actual JSONL and a JSON array
    let items: Vec<Value> = if input.trim_start().starts_with('[') {
        serde_json::from_str(input).map_err(|e| ParseError::InvalidJson(e.to_string()))?
    } else {
        input
            .lines()
            .filter(|l| !l.trim().is_empty())
            .map(serde_json::from_str)
            .collect::<Result<Vec<_>, _>>()
            .map_err(|e| ParseError::InvalidJson(e.to_string()))?
    };

    for item in &items {
        let token = item
            .get("token")
            .and_then(|v| v.as_str())
            .ok_or_else(|| ParseError::MissingField("token (JSONL entry)".to_string()))?
            .to_string();
        let logprob = item
            .get("logprob")
            .and_then(|v| v.as_f64())
            .ok_or_else(|| ParseError::MissingField("logprob (JSONL entry)".to_string()))?;

        let bytes = item.get("bytes").and_then(|v| v.as_array()).map(|arr| {
            arr.iter()
                .filter_map(|b| b.as_u64().map(|n| n as u8))
                .collect()
        });

        total_logprob += logprob;
        tokens.push(TokenLogprob {
            token,
            logprob,
            bytes,
            top_logprobs: None,
        });
    }

    Ok(LogprobSequence {
        tokens,
        model: None,
        format_detected: InputFormat::JsonlStream.to_string(),
        total_logprob,
    })
}

/// Parse Google Gemini format.
///
/// Gemini uses `candidates[0].logprobsResult` with `chosenCandidates` (chosen tokens)
/// and `topCandidates` (top-k alternatives per position). Field names differ from OpenAI:
/// `logProbability` instead of `logprob`, `tokenId` instead of `bytes`.
fn parse_gemini(input: &str) -> Result<LogprobSequence, ParseError> {
    let val: Value =
        serde_json::from_str(input).map_err(|e| ParseError::InvalidJson(e.to_string()))?;
    let model = val
        .get("modelVersion")
        .or_else(|| val.get("model"))
        .and_then(|m| m.as_str())
        .map(String::from);

    let logprobs_result = val
        .pointer("/candidates/0/logprobsResult")
        .ok_or_else(|| ParseError::MissingField("candidates[0].logprobsResult".to_string()))?;

    let chosen = logprobs_result
        .get("chosenCandidates")
        .and_then(|v| v.as_array())
        .ok_or_else(|| ParseError::MissingField("chosenCandidates".to_string()))?;

    let top_candidates = logprobs_result
        .get("topCandidates")
        .and_then(|v| v.as_array());

    let mut tokens = Vec::with_capacity(chosen.len());
    let mut total_logprob = 0.0;

    for (i, entry) in chosen.iter().enumerate() {
        let token = entry
            .get("token")
            .and_then(|v| v.as_str())
            .ok_or_else(|| ParseError::MissingField(format!("token at position {i}")))?
            .to_string();
        let logprob = entry
            .get("logProbability")
            .and_then(|v| v.as_f64())
            .ok_or_else(|| ParseError::MissingField(format!("logProbability at position {i}")))?;

        // Gemini does not return per-token byte arrays. Deriving them from the
        // token string (token.as_bytes()) is exactly the fallback that BPB
        // refuses elsewhere, so leave bytes unset rather than fabricate them.
        // This keeps the strict-BPB guarantee honest for Gemini input.
        let bytes = None;

        // Extract top-k from topCandidates[i].candidates
        let top_logprobs = top_candidates.and_then(|tc| {
            tc.get(i)
                .and_then(|pos| pos.get("candidates"))
                .and_then(|v| v.as_array())
                .map(|arr| {
                    let mut entries: Vec<TopKEntry> = arr
                        .iter()
                        .filter_map(|c| {
                            let t = c.get("token")?.as_str()?.to_string();
                            let lp = c.get("logProbability")?.as_f64()?;
                            Some(TopKEntry {
                                token: t,
                                logprob: lp,
                            })
                        })
                        .collect();
                    entries.sort_by(|a, b| {
                        b.logprob
                            .partial_cmp(&a.logprob)
                            .unwrap_or(std::cmp::Ordering::Equal)
                    });
                    entries
                })
        });

        total_logprob += logprob;
        tokens.push(TokenLogprob {
            token,
            logprob,
            bytes,
            top_logprobs,
        });
    }

    Ok(LogprobSequence {
        tokens,
        model,
        format_detected: InputFormat::Gemini.to_string(),
        total_logprob,
    })
}

/// Parse Ollama native format.
///
/// Ollama's `/api/generate` and `/api/chat` return `logprobs` as a top-level array
/// (not nested under `choices`). Token objects use the same field names as OpenAI:
/// `token`, `logprob`, `bytes`, `top_logprobs`.
fn parse_ollama(input: &str) -> Result<LogprobSequence, ParseError> {
    let val: Value =
        serde_json::from_str(input).map_err(|e| ParseError::InvalidJson(e.to_string()))?;
    let model = val.get("model").and_then(|m| m.as_str()).map(String::from);

    let logprobs_arr = val
        .get("logprobs")
        .and_then(|v| v.as_array())
        .ok_or_else(|| ParseError::MissingField("top-level logprobs array".to_string()))?;

    let mut tokens = Vec::with_capacity(logprobs_arr.len());
    let mut total_logprob = 0.0;

    for (i, item) in logprobs_arr.iter().enumerate() {
        let token = item
            .get("token")
            .and_then(|v| v.as_str())
            .ok_or_else(|| ParseError::MissingField(format!("token at position {i}")))?
            .to_string();
        let logprob = item
            .get("logprob")
            .and_then(|v| v.as_f64())
            .ok_or_else(|| ParseError::MissingField(format!("logprob at position {i}")))?;

        let bytes = item.get("bytes").and_then(|v| v.as_array()).map(|arr| {
            arr.iter()
                .filter_map(|b| b.as_u64().map(|n| n as u8))
                .collect()
        });

        let top_logprobs = item
            .get("top_logprobs")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|entry| {
                        let t = entry.get("token")?.as_str()?.to_string();
                        let lp = entry.get("logprob")?.as_f64()?;
                        Some(TopKEntry {
                            token: t,
                            logprob: lp,
                        })
                    })
                    .collect()
            });

        total_logprob += logprob;
        tokens.push(TokenLogprob {
            token,
            logprob,
            bytes,
            top_logprobs,
        });
    }

    Ok(LogprobSequence {
        tokens,
        model,
        format_detected: InputFormat::Ollama.to_string(),
        total_logprob,
    })
}

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

    #[test]
    fn test_detect_openai() {
        let input = r#"{"choices":[{"logprobs":{"content":[{"token":"Hi","logprob":-0.5}]}}]}"#;
        let seq = parse_string(input, None, false).unwrap();
        assert_eq!(seq.format_detected, "openai");
        assert_eq!(seq.tokens.len(), 1);
    }

    #[test]
    fn test_detect_vllm() {
        let input = r#"{"choices":[{"logprobs":{"tokens":["Hi"],"token_logprobs":[-0.5]}}]}"#;
        let seq = parse_string(input, None, false).unwrap();
        assert_eq!(seq.format_detected, "vllm");
    }

    #[test]
    fn test_detect_jsonl() {
        let input = "{\"token\":\"Hi\",\"logprob\":-0.5}\n{\"token\":\" there\",\"logprob\":-1.0}";
        let seq = parse_string(input, None, false).unwrap();
        assert_eq!(seq.format_detected, "jsonl");
        assert_eq!(seq.tokens.len(), 2);
    }

    #[test]
    fn test_detect_gemini() {
        let input = r#"{
            "candidates": [{
                "content": {"parts": [{"text": "Paris"}], "role": "model"},
                "logprobsResult": {
                    "topCandidates": [
                        {"candidates": [
                            {"token": "Paris", "tokenId": 1, "logProbability": -0.05},
                            {"token": "The", "tokenId": 2, "logProbability": -3.12}
                        ]}
                    ],
                    "chosenCandidates": [
                        {"token": "Paris", "tokenId": 1, "logProbability": -0.05}
                    ]
                }
            }]
        }"#;
        let seq = parse_string(input, None, false).unwrap();
        assert_eq!(seq.format_detected, "gemini");
        assert_eq!(seq.tokens.len(), 1);
        assert_eq!(seq.tokens[0].token, "Paris");
        assert!((seq.tokens[0].logprob - (-0.05)).abs() < 1e-6);
        let top = seq.tokens[0].top_logprobs.as_ref().unwrap();
        assert_eq!(top.len(), 2);
        assert_eq!(top[0].token, "Paris"); // sorted descending
    }

    #[test]
    fn test_detect_ollama() {
        let input = r#"{
            "model": "gemma3",
            "response": "Hello",
            "done": true,
            "logprobs": [
                {
                    "token": "Hello",
                    "logprob": -0.523,
                    "bytes": [72, 101, 108, 108, 111],
                    "top_logprobs": [
                        {"token": "Hello", "logprob": -0.523, "bytes": [72, 101, 108, 108, 111]},
                        {"token": "Hi", "logprob": -1.8, "bytes": [72, 105]}
                    ]
                }
            ]
        }"#;
        let seq = parse_string(input, None, false).unwrap();
        assert_eq!(seq.format_detected, "ollama");
        assert_eq!(seq.tokens.len(), 1);
        assert_eq!(seq.model.unwrap(), "gemma3");
        assert_eq!(
            seq.tokens[0].bytes.as_ref().unwrap(),
            &vec![72, 101, 108, 108, 111]
        );
    }

    #[test]
    fn test_strict_rejects_ambiguous() {
        let input = r#"[{"token":"Hi","logprob":-0.5}]"#;
        // This is valid as jsonl-style array — strict should still parse it
        // since it falls through to the array detection
        let result = parse_string(input, None, true);
        // strict mode rejects if JSON structure doesn't match openai/vllm
        assert!(result.is_err());
    }
}