//! Text-level roundtrip tests for the real chat-template and output-processor pairing.
//!
//! The invariant under test is that a structured assistant message rendered as history can be
//! parsed from the generated assistant completion and then rendered back to the exact same
//! assistant-completion text.

use std::pin::Pin;
use std::sync::Arc;

use anyhow::{Context as _, Result, bail, ensure};
use futures::{Stream, StreamExt as _, stream};
use serde_json_fmt::JsonFormat as JsonFmt;
use serial_test::file_serial;
use vllm_chat::{
    AssistantContentBlock, AssistantMessage, AssistantMessageExt as _, AssistantToolCall,
    ChatEvent, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice, FinishReason,
    GenerationPromptMode, LoadModelBackendsOptions, NewChatOutputProcessorOptions, ParserSelection,
    RendererSelection, load_model_backends,
};
use vllm_text::{DecodedTextEvent, Finished, Prompt};

/// One model/parser configuration used to run the fixed roundtrip fixtures.
#[derive(Clone)]
struct RoundtripCase {
    /// Hugging Face model id resolved through the production backend loader.
    model_id: &'static str,
    /// Final assistant-history suffix rendered by the chat template but not
    /// generated by the model body consumed by the output processor.
    // TODO: we should adopt `ContinueFinalAssistant` mode to naturally handle this.
    assistant_stop_suffix: &'static str,
    /// Tool parser selection used by the output processor.
    tool_call_parser: ParserSelection,
    /// Reasoning parser selection used by the output processor.
    reasoning_parser: ParserSelection,
    /// How this model's chat template handles thinking mode.
    thinking_behavior: ThinkingBehavior,
    /// JSON formatting expected after this model's template has materialized
    /// tool-call arguments.
    json_fmt: JsonFmt,
    /// Whether the template renders tool-call argument object keys in sorted order.
    sort_json_keys: bool,
}

#[derive(Clone, Copy)]
enum ThinkingBehavior {
    /// The chat template accepts explicit thinking on/off kwargs, and uses
    /// `default` when the request does not specify either kwarg.
    Toggleable { default: bool },
    /// The chat template always behaves as `value` for this fixture.
    Always { value: bool },
}

impl ThinkingBehavior {
    fn default(self) -> bool {
        match self {
            Self::Toggleable { default } => default,
            Self::Always { value } => value,
        }
    }

    fn fixtures(self) -> Vec<Option<bool>> {
        match self {
            Self::Toggleable { .. } => vec![
                Some(true),  // explicitly enable thinking
                Some(false), // explicitly disable thinking
                None,        // use default template behavior
            ],
            Self::Always { value } => vec![
                Some(value), // explicitly request the supported thinking behavior
                None,        // use default template behavior
            ],
        }
    }
}

impl RoundtripCase {
    /// Qwen3 XML tool-call format with `qwen3` reasoning tags.
    fn qwen3() -> Self {
        Self {
            model_id: "Qwen/Qwen3-0.6B",
            assistant_stop_suffix: "<|im_end|>\n",
            tool_call_parser: ParserSelection::Auto,
            reasoning_parser: ParserSelection::Auto,
            thinking_behavior: ThinkingBehavior::Toggleable { default: true },
            json_fmt: spaced_json_fmt(),
            sort_json_keys: false,
        }
    }

    /// Qwen3.5 coder-style JSON tool-call format with `qwen3` reasoning tags.
    fn qwen35() -> Self {
        Self {
            model_id: "Qwen/Qwen3.5-4B",
            assistant_stop_suffix: "<|im_end|>\n",
            tool_call_parser: ParserSelection::Auto,
            reasoning_parser: ParserSelection::Auto,
            thinking_behavior: ThinkingBehavior::Toggleable { default: true },
            json_fmt: compact_json_fmt(),
            sort_json_keys: false,
        }
    }

    /// MiniMax M2.5 XML invoke format with `<think>` reasoning tags.
    fn minimax_m25() -> Self {
        Self {
            model_id: "MiniMaxAI/MiniMax-M2.5",
            assistant_stop_suffix: "[e~[\n",
            tool_call_parser: ParserSelection::Auto,
            reasoning_parser: ParserSelection::Auto,
            thinking_behavior: ThinkingBehavior::Always { value: true },
            json_fmt: compact_json_fmt(),
            sort_json_keys: false,
        }
    }

    /// DeepSeek V4 DSML tool-call format.
    fn deepseek_v4() -> Self {
        Self {
            model_id: "deepseek-ai/DeepSeek-V4-Flash",
            assistant_stop_suffix: "<｜end▁of▁sentence｜>",
            tool_call_parser: ParserSelection::Auto,
            reasoning_parser: ParserSelection::Auto,
            thinking_behavior: ThinkingBehavior::Toggleable { default: false },
            json_fmt: compact_json_fmt(),
            sort_json_keys: false,
        }
    }

    /// GLM-4.7 XML-like argument format with `<think>` reasoning tags.
    fn glm47() -> Self {
        Self {
            model_id: "zai-org/GLM-4.7-Flash",
            assistant_stop_suffix: "",
            tool_call_parser: ParserSelection::Auto,
            reasoning_parser: ParserSelection::Auto,
            thinking_behavior: ThinkingBehavior::Toggleable { default: true },
            json_fmt: compact_json_fmt(),
            sort_json_keys: false,
        }
    }

    /// Gemma4 channel reasoning with custom function-call arguments.
    fn gemma4() -> Self {
        Self {
            model_id: "google/gemma-4-E4B-it",
            assistant_stop_suffix: "<|tool_response>",
            tool_call_parser: ParserSelection::Auto,
            reasoning_parser: ParserSelection::Auto,
            thinking_behavior: ThinkingBehavior::Always { value: true },
            json_fmt: compact_json_fmt(),
            sort_json_keys: true,
        }
    }

    /// Kimi K2.5 tool-call format with `<think>` reasoning tags.
    #[allow(dead_code)]
    fn kimi_k25() -> Self {
        Self {
            model_id: "moonshotai/Kimi-K2.5",
            assistant_stop_suffix: "<|im_end|>",
            tool_call_parser: ParserSelection::Auto,
            reasoning_parser: ParserSelection::Auto,
            thinking_behavior: ThinkingBehavior::Toggleable { default: true },
            json_fmt: spaced_json_fmt(),
            sort_json_keys: false,
        }
    }

    /// SeedOSS with `<seed:think>` / `</seed:think>` reasoning tags.
    fn seed_oss() -> Self {
        Self {
            model_id: "ByteDance-Seed/Seed-OSS-36B-Instruct",
            assistant_stop_suffix: "<seed:eos>",
            tool_call_parser: ParserSelection::Auto,
            reasoning_parser: ParserSelection::Auto,
            thinking_behavior: ThinkingBehavior::Always { value: true },
            json_fmt: compact_json_fmt(),
            sort_json_keys: false,
        }
    }

    /// Step-3.5 with `<think>` / `</think>` reasoning tags and newline trimming.
    fn step3p5() -> Self {
        Self {
            model_id: "stepfun-ai/Step-3.5-Flash",
            assistant_stop_suffix: "<|im_end|>\n",
            tool_call_parser: ParserSelection::Auto,
            reasoning_parser: ParserSelection::Auto,
            thinking_behavior: ThinkingBehavior::Always { value: true },
            json_fmt: compact_json_fmt(),
            sort_json_keys: false,
        }
    }
}

macro_rules! roundtrip_tests {
    ($($case:ident => [$($fixture:ident),* $(,)?]),+ $(,)?) => {
        paste::paste! {
            $(
                $(
                    #[tokio::test]
                    #[file_serial([<hf_ $case>])]
                    async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
                        [<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
                    }
                )*
            )+
        }
    };
}

roundtrip_tests! {
    qwen3 => [reasoning_and_content, tool_call_mix],
    qwen35 => [reasoning_and_content, tool_call_mix],
    minimax_m25 => [reasoning_and_content, tool_call_mix],
    deepseek_v4 => [reasoning_and_content, tool_call_mix],
    glm47 => [reasoning_and_content, tool_call_mix],
    seed_oss => [reasoning_and_content],
    step3p5 => [reasoning_and_content],

    gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call
    kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history
}

/// Run the fixed reasoning+content fixture for one model/parser case.
async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> Result<()> {
    for thinking in case.thinking_behavior.fixtures() {
        run_roundtrip_reasoning_and_content_inner(case.clone(), thinking).await?;
    }
    Ok(())
}

async fn run_roundtrip_reasoning_and_content_inner(
    case: RoundtripCase,
    thinking: Option<bool>,
) -> Result<()> {
    let backends = load_roundtrip_backends(&case).await?;
    let request = roundtrip_request(
        "roundtrip-reasoning-content",
        vec![ChatMessage::text(ChatRole::User, "What is 2 + 2?")],
        Vec::new(),
        thinking,
    );
    let expected_reasoning = "Need compute 2 + 2 directly.";
    let expected_text = "The answer is 4.";
    let effective_thinking = thinking.unwrap_or(case.thinking_behavior.default());

    let assistant = {
        let mut content = Vec::new();
        if effective_thinking {
            content.push(AssistantContentBlock::Reasoning {
                text: expected_reasoning.to_string(),
            });
        }
        content.push(AssistantContentBlock::Text {
            text: expected_text.to_string(),
        });
        AssistantMessage { content }
    };
    let result = run_roundtrip(&case, &backends, &request, assistant).await?;

    assert_eq!(
        result.parsed_message.reasoning().as_deref().map(str::trim),
        effective_thinking.then_some(expected_reasoning)
    );
    assert_eq!(result.parsed_message.text().trim(), expected_text);
    assert_eq!(result.parsed_message.tool_calls().count(), 0);

    assert_eq!(
        result.rerendered_closed_completion,
        result.closed_completion
    );

    Ok(())
}

/// Run the fixed reasoning+multiple-tools fixture for one model/parser case.
async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
    let backends = load_roundtrip_backends(&case).await?;
    let request = roundtrip_request(
        "roundtrip-reasoning-tools",
        vec![ChatMessage::text(
            ChatRole::User,
            "Check Shanghai weather and add 1.0 plus 2.",
        )],
        test_tools(),
        Some(true), // always enable thinking in this fixture
    );
    let expected_reasoning = "Need call the weather and add tools.";
    let expected_text = "I will call the tools.";

    let result = run_roundtrip(
        &case,
        &backends,
        &request,
        AssistantMessage {
            content: vec![
                AssistantContentBlock::Reasoning {
                    text: expected_reasoning.to_string(),
                },
                AssistantContentBlock::Text {
                    text: expected_text.to_string(),
                },
                AssistantContentBlock::ToolCall(AssistantToolCall {
                    id: "functions.get_weather:0".to_string(),
                    name: "get_weather".to_string(),
                    arguments: r#"{"location":"Shanghai"}"#.to_string(),
                }),
                AssistantContentBlock::ToolCall(AssistantToolCall {
                    id: "functions.add:1".to_string(),
                    name: "add".to_string(),
                    // Intentionally use a non-lexical order of keys to verify text-level
                    // fidelity of the roundtrip where JSON formatting remains stable. The
                    // `items` key also exercises templates that call `arguments.items()`.
                    arguments: r#"{"y":1.0,"x":2,"items":["left","right"]}"#.to_string(),
                }),
            ],
        },
    )
    .await?;

    assert_eq!(
        result.parsed_message.reasoning().as_deref().map(str::trim),
        Some(expected_reasoning)
    );
    assert_eq!(result.parsed_message.text().trim(), expected_text);

    let tool_calls = result.parsed_message.tool_calls().collect::<Vec<_>>();
    assert_eq!(
        tool_calls.len(),
        2,
        "parsed message: {:#?}",
        result.parsed_message
    );
    assert_eq!(tool_calls[0].name, "get_weather");
    assert_eq!(
        tool_calls[0].arguments,
        expected_arguments(&case, r#"{"location": "Shanghai"}"#)?,
    );
    assert_eq!(tool_calls[1].name, "add");
    assert_eq!(
        tool_calls[1].arguments,
        expected_arguments(&case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
    );

    assert_eq!(
        result.rerendered_closed_completion,
        result.closed_completion
    );

    Ok(())
}

/// Compact JSON argument formatting used by JSON-native parsers/renderers.
fn compact_json_fmt() -> JsonFmt {
    JsonFmt::new()
}

/// Python `json.dumps`-style compact formatting with a space after commas and
/// colons.
fn spaced_json_fmt() -> JsonFmt {
    JsonFmt::new()
        .comma(", ")
        .expect("literal comma separator is valid JSON")
        .colon(": ")
        .expect("literal colon separator is valid JSON")
}

/// Parse and format expected tool-call arguments from raw JSON text.
/// Pass in a raw JSON string instead of a structured value to ensure the exact precision and
/// formatting of numbers are preserved.
fn expected_arguments(case: &RoundtripCase, raw_json: &str) -> Result<String> {
    let mut value: serde_json::Value =
        serde_json::from_str(raw_json).context("invalid expected tool-call arguments")?;
    if case.sort_json_keys {
        sort_json_value(&mut value);
    }

    case.json_fmt
        .format_to_string(&value)
        .context("failed to format expected tool-call arguments")
}

/// Sort JSON object keys recursively to match templates that render mappings with `dictsort`.
fn sort_json_value(value: &mut serde_json::Value) {
    match value {
        serde_json::Value::Object(map) => {
            for value in map.values_mut() {
                sort_json_value(value);
            }

            let mut entries = std::mem::take(map).into_iter().collect::<Vec<_>>();
            entries.sort_by(|(left, _), (right, _)| left.cmp(right));
            map.extend(entries);
        }
        serde_json::Value::Array(values) => {
            for value in values {
                sort_json_value(value);
            }
        }
        _ => {}
    }
}

/// Load the real model chat/text backend for one roundtrip case.
async fn load_roundtrip_backends(case: &RoundtripCase) -> Result<vllm_chat::LoadedModelBackends> {
    load_model_backends(
        case.model_id,
        LoadModelBackendsOptions {
            renderer: RendererSelection::Auto,
            ..Default::default()
        },
    )
    .await
    .with_context(|| format!("failed to load HF model files for {}", case.model_id))
}

/// Roundtrip artifacts needed for semantic and exact-text assertions.
struct RoundtripResult {
    /// Final assistant message reconstructed by the output processor.
    parsed_message: AssistantMessage,
    /// Assistant-completion suffix cut from rendering the expected assistant as
    /// history.
    closed_completion: String,
    /// Assistant-completion suffix cut after rendering the parsed assistant
    /// back as history.
    rerendered_closed_completion: String,
}

/// Render, parse, and rerender one assistant turn through the production
/// renderer/output-processor boundary.
async fn run_roundtrip(
    case: &RoundtripCase,
    backends: &vllm_chat::LoadedModelBackends,
    request: &ChatRequest,
    assistant: AssistantMessage,
) -> Result<RoundtripResult> {
    let renderer = backends.chat_backend.chat_renderer();
    let (prompt, closed_completion_text) =
        render_closed_completion(renderer.as_ref(), request, &assistant)?;
    let completion_body = closed_completion_text
        .strip_suffix(case.assistant_stop_suffix)
        .with_context(|| {
            format!(
                "closed assistant completion did not end with {:?}: {:?}",
                case.assistant_stop_suffix, closed_completion_text
            )
        })?;

    let parsed_message =
        parse_completion(case, backends, request, &prompt, completion_body).await?;
    let (_, rerendered_closed_completion) =
        render_closed_completion(renderer.as_ref(), request, &parsed_message)?;

    Ok(RoundtripResult {
        parsed_message,
        closed_completion: closed_completion_text,
        rerendered_closed_completion,
    })
}

/// Render `history` as a production prompt and `history + assistant` as closed
/// history, then return the production prompt and assistant-completion suffix.
fn render_closed_completion(
    renderer: &dyn vllm_chat::ChatRenderer,
    base_request: &ChatRequest,
    assistant: &AssistantMessage,
) -> Result<(String, String)> {
    let mut prompt_request = base_request.clone();
    prompt_request.chat_options.generation_prompt_mode = GenerationPromptMode::StartNewAssistant;
    let prompt = render_text(renderer, &prompt_request).context("failed to render prompt")?;

    let mut full_request = base_request.clone();
    full_request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
    full_request.messages.push(ChatMessage::from(assistant.clone()));
    let full = render_text(renderer, &full_request).context("failed to render full prompt")?;

    ensure!(
        full.starts_with(&prompt),
        "full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
    );
    let completion = full[prompt.len()..].to_string();

    Ok((prompt, completion))
}

/// Render one chat request and require a text prompt.
fn render_text(renderer: &dyn vllm_chat::ChatRenderer, request: &ChatRequest) -> Result<String> {
    match renderer.render(request)?.prompt {
        Prompt::Text(text) => Ok(text),
        other => bail!("roundtrip tests expect text prompts, got {other:?}"),
    }
}

/// Feed one rendered assistant completion body into the real output processor
/// and collect its terminal assistant message.
async fn parse_completion(
    case: &RoundtripCase,
    backends: &vllm_chat::LoadedModelBackends,
    base_request: &ChatRequest,
    prompt: &str,
    completion_body: &str,
) -> Result<AssistantMessage> {
    let tokenizer = backends.text_backend.tokenizer();
    let prompt_token_ids = tokenizer
        .encode(prompt, base_request.add_special_tokens)
        .context("failed to encode rendered prompt")?;

    let mut request = base_request.clone();
    let processor = backends.chat_backend.new_chat_output_processor(
        &mut request,
        NewChatOutputProcessorOptions {
            tool_call_parser: &case.tool_call_parser,
            reasoning_parser: &case.reasoning_parser,
        },
    )?;

    let decoded = decoded_completion_stream(prompt_token_ids, completion_body);
    let mut events = processor.process(decoded)?;

    while let Some(event) = events.next().await {
        if let ChatEvent::Done { message, .. } = event? {
            // TODO: currently our parsers are not very strict about preserving or trimming
            // whitespace, so we trim here to avoid roundtrip failures due to
            // insignificant whitespace differences. However, this may hurt token-level
            // fidelity so we should consider improving them.
            return Ok(message.trim());
        }
    }

    bail!("output processor finished without a Done event")
}

/// Build a decoded-text stream from an already-rendered completion body.
///
/// The first event carries real prompt token ids so reasoning parsers can
/// initialize from the same prompt boundary production uses. Completion text is
/// split into small chunks to exercise streaming parser state across marker
/// and JSON boundaries.
fn decoded_completion_stream(
    prompt_token_ids: Vec<u32>,
    completion_body: &str,
) -> Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>> {
    let prompt_token_count = prompt_token_ids.len();
    let mut events = vec![DecodedTextEvent::Start {
        prompt_token_ids: Arc::from(prompt_token_ids.into_boxed_slice()),
        prompt_logprobs: None,
    }];

    let chunks = split_by_chars(completion_body, 7);
    if chunks.is_empty() {
        events.push({
            DecodedTextEvent::TextDelta {
                delta: String::new(),
                token_ids: Vec::new(),
                logprobs: None,
                finished: Some(Finished {
                    usage: vllm_llm::TokenUsage {
                        prompt_token_count: 0,
                        output_token_count: 0,
                        cached_token_count: 0,
                    },
                    finish_reason: FinishReason::stop_eos(),
                    kv_transfer_params: None,
                }),
            }
        });
    } else {
        let last_index = chunks.len() - 1;
        for (index, chunk) in chunks.into_iter().enumerate() {
            let finished = (index == last_index).then(|| Finished {
                usage: vllm_llm::TokenUsage {
                    prompt_token_count,
                    output_token_count: completion_body.chars().count(),
                    cached_token_count: 0,
                },
                finish_reason: FinishReason::stop_eos(),
                kv_transfer_params: None,
            });
            events.push(DecodedTextEvent::TextDelta {
                delta: chunk,
                token_ids: Vec::new(),
                logprobs: None,
                finished,
            });
        }
    }

    stream::iter(events).map(Ok).boxed()
}

/// Split text into chunks containing at most `chunk_chars` Unicode scalar
/// values.
fn split_by_chars(text: &str, chunk_chars: usize) -> Vec<String> {
    let mut chunks = Vec::new();
    let mut start = 0;
    let mut count = 0;

    for (index, _) in text.char_indices() {
        if count == chunk_chars {
            chunks.push(text[start..index].to_string());
            start = index;
            count = 0;
        }
        count += 1;
    }

    if start < text.len() {
        chunks.push(text[start..].to_string());
    }

    chunks
}

/// Build a chat request fixture with parser-enabling tool-choice semantics.
fn roundtrip_request(
    request_id: impl Into<String>,
    messages: Vec<ChatMessage>,
    tools: Vec<ChatTool>,
    thinking: Option<bool>,
) -> ChatRequest {
    let mut request = ChatRequest {
        request_id: request_id.into(),
        messages,
        tool_choice: if tools.is_empty() {
            ChatToolChoice::None
        } else {
            ChatToolChoice::Auto
        },
        tools,
        ..ChatRequest::for_test()
    };

    // Explicitly enable or disable thinking so that rendering and parsing the reasoning block is
    // exercised or skipped in the roundtrip. If unspecified, use the default template behavior.
    if let Some(thinking) = thinking {
        for key in ["thinking", "enable_thinking"] {
            request.chat_options.template_kwargs.insert(key.to_string(), thinking.into());
        }
    }

    request
}

/// Return the function tools used by the multiple-tool-call fixture.
fn test_tools() -> Vec<ChatTool> {
    vec![
        ChatTool {
            name: "get_weather".to_string(),
            description: Some("Get weather for a location".to_string()),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "location": { "type": "string" }
                },
                "required": ["location"]
            }),
            strict: None,
        },
        ChatTool {
            name: "add".to_string(),
            description: Some("Add two integers".to_string()),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "y": { "type": "number" },
                    "x": { "type": "number" },
                    "items": {
                        "type": "array",
                        "items": { "type": "string" }
                    }
                },
                "required": ["y", "x", "items"]
            }),
            strict: None,
        },
    ]
}
