use std::collections::HashMap;

use anyhow::{bail, Result};
use api_contracts::{
    chat::{
        ChatCompletionRequest, DeltaToolCall, DeltaToolFunctionCall, FinishReason, ResponseFormat,
        Role, ToolCall, ToolChoice, ToolDefinition, ToolFunctionCall, ToolType,
    },
    error::ApiError,
    models::ModelInfo,
};
use axum::http::{HeaderMap, HeaderName, HeaderValue, StatusCode};
use serde_json::{json, Value};
use tower_http::cors::{AllowHeaders, AllowMethods, AllowOrigin, Any, CorsLayer};

const DEFAULT_SYSTEM_PROMPT: &str = concat!(
    "You are a precise coding assistant running locally. ",
    "Prioritize correct code, minimal speculation, explicit assumptions, and actionable fixes. ",
    "When editing or generating code, prefer complete snippets, preserve requested APIs, and explain risks briefly."
);

#[derive(Debug, Clone)]
pub struct ServerConfig {
    pub api_key: Option<String>,
    pub auth_header: String,
    pub auth_scheme: Option<String>,
    pub cors_origins: Vec<String>,
    pub default_system_prompt: String,
    pub enable_tool_calls: bool,
    pub model_catalog: ModelCatalog,
}

impl ServerConfig {
    pub fn from_env(model_name: String) -> Self {
        let aliases = env_csv("LOCALLLM_MODEL_ALIASES");
        Self {
            api_key: std::env::var("LOCALLLM_API_KEY")
                .ok()
                .filter(|v| !v.is_empty()),
            auth_header: std::env::var("LOCALLLM_AUTH_HEADER")
                .ok()
                .filter(|v| !v.trim().is_empty())
                .unwrap_or_else(|| "authorization".to_owned())
                .to_ascii_lowercase(),
            auth_scheme: std::env::var("LOCALLLM_AUTH_SCHEME")
                .ok()
                .map(|v| v.trim().to_owned())
                .filter(|v| !v.is_empty()),
            cors_origins: {
                let origins = env_csv("LOCALLLM_CORS_ORIGINS");
                if origins.is_empty() {
                    vec!["*".to_owned()]
                } else {
                    origins
                }
            },
            default_system_prompt: std::env::var("LOCALLLM_SYSTEM_PROMPT")
                .ok()
                .filter(|v| !v.trim().is_empty())
                .unwrap_or_else(|| DEFAULT_SYSTEM_PROMPT.to_owned()),
            enable_tool_calls: env_bool("LOCALLLM_ENABLE_TOOL_CALLS", true),
            model_catalog: ModelCatalog::new(model_name, aliases),
        }
    }

    pub fn cors_layer(&self) -> CorsLayer {
        let methods = AllowMethods::any();
        let headers = AllowHeaders::any();

        if self.cors_origins.iter().any(|origin| origin == "*") {
            return CorsLayer::new()
                .allow_origin(Any)
                .allow_methods(methods)
                .allow_headers(headers);
        }

        let origins: Vec<HeaderValue> = self
            .cors_origins
            .iter()
            .filter_map(|origin| HeaderValue::from_str(origin).ok())
            .collect();

        CorsLayer::new()
            .allow_origin(AllowOrigin::list(origins))
            .allow_methods(methods)
            .allow_headers(headers)
    }

    pub fn resolve_model(&self, requested: &str) -> Option<&ModelDescriptor> {
        self.model_catalog.resolve(requested)
    }

    pub fn model_list(&self, created: i64) -> Vec<ModelInfo> {
        self.model_catalog.to_model_list(created)
    }
}

#[derive(Debug, Clone)]
pub struct ModelCatalog {
    models: Vec<ModelDescriptor>,
    aliases: HashMap<String, usize>,
}

impl ModelCatalog {
    pub fn new(canonical_id: String, aliases: Vec<String>) -> Self {
        let mut models = vec![ModelDescriptor {
            id: canonical_id.clone(),
            root: canonical_id.clone(),
            parent: None,
        }];

        let mut seen = std::collections::BTreeSet::new();
        seen.insert(canonical_id.to_ascii_lowercase());
        seen.insert("local-model".to_owned());
        seen.insert("default".to_owned());

        for alias in [
            vec!["local-model".to_owned(), "default".to_owned()],
            aliases,
        ]
        .concat()
        {
            let trimmed = alias.trim();
            if trimmed.is_empty() {
                continue;
            }
            let lowered = trimmed.to_ascii_lowercase();
            if !seen.insert(lowered) {
                continue;
            }
            models.push(ModelDescriptor {
                id: trimmed.to_owned(),
                root: canonical_id.clone(),
                parent: Some(canonical_id.clone()),
            });
        }

        let aliases = models
            .iter()
            .enumerate()
            .map(|(idx, model)| (model.id.to_ascii_lowercase(), idx))
            .collect();

        Self { models, aliases }
    }

    pub fn resolve(&self, requested: &str) -> Option<&ModelDescriptor> {
        self.aliases
            .get(&requested.to_ascii_lowercase())
            .and_then(|idx| self.models.get(*idx))
    }
    pub fn to_model_list(&self, created: i64) -> Vec<ModelInfo> {
        self.models
            .iter()
            .map(|model| ModelInfo {
                id: model.id.clone(),
                object: "model",
                created,
                owned_by: "local",
                root: model.root.clone(),
                parent: model.parent.clone(),
                permission: Vec::new(),
            })
            .collect()
    }
}

#[derive(Debug, Clone)]
pub struct ModelDescriptor {
    pub id: String,
    pub root: String,
    pub parent: Option<String>,
}

#[derive(Debug, Clone)]
pub struct PreparedPrompt {
    pub prompt: String,
    pub streaming_mode: StreamingMode,
    /// Extra stop sequences injected by the template (e.g. Generic transcript stops)
    pub extra_stop_sequences: Vec<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamingMode {
    LiveTokens,
    Buffered,
}

#[derive(Debug, Clone)]
pub struct FinalizedAssistant {
    pub content: Option<String>,
    pub tool_calls: Option<Vec<ToolCall>>,
    pub finish_reason: FinishReason,
}

#[derive(Debug, Clone)]
struct PromptMessage {
    role: Role,
    name: Option<String>,
    content: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ChatTemplate {
    Generic,
    Qwen,
    Gemma,
    Llama,
    Kimi,
}

pub fn authorize_request(
    headers: &HeaderMap,
    config: &ServerConfig,
) -> Result<(), (StatusCode, ApiError)> {
    let Some(expected_key) = config.api_key.as_deref() else {
        return Ok(());
    };

    let configured_header = HeaderName::from_bytes(config.auth_header.as_bytes())
        .ok()
        .and_then(|name| headers.get(name))
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty());

    let fallback_bearer = headers
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty());

    let fallback_api_key = headers
        .get("x-api-key")
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty());

    let provided = configured_header.or(fallback_bearer).or(fallback_api_key);
    let Some(provided) = provided else {
        return Err((
            StatusCode::UNAUTHORIZED,
            ApiError::new("Missing API key", "invalid_request_error").with_code("missing_api_key"),
        ));
    };

    let normalized = match config.auth_scheme.as_deref() {
        Some(scheme) => provided
            .strip_prefix(scheme)
            .map(str::trim)
            .unwrap_or(provided),
        None => provided,
    };

    if normalized == expected_key || provided == expected_key {
        Ok(())
    } else {
        Err((
            StatusCode::UNAUTHORIZED,
            ApiError::new("Invalid API key", "invalid_request_error").with_code("invalid_api_key"),
        ))
    }
}

pub fn validate_request(
    req: &ChatCompletionRequest,
    config: &ServerConfig,
) -> Result<(), (StatusCode, ApiError)> {
    if config.resolve_model(&req.model).is_none() {
        return Err((
            StatusCode::NOT_FOUND,
            ApiError::new(
                format!("The model '{}' does not exist", req.model),
                "invalid_request_error",
            )
            .with_code("model_not_found")
            .with_param("model"),
        ));
    }

    if req.messages.is_empty() {
        return Err((
            StatusCode::BAD_REQUEST,
            ApiError::new("messages must not be empty", "invalid_request_error")
                .with_code("empty_messages")
                .with_param("messages"),
        ));
    }

    if req.n != 1 {
        return Err((
            StatusCode::BAD_REQUEST,
            ApiError::new("Only n=1 is currently supported", "invalid_request_error")
                .with_code("unsupported_value")
                .with_param("n"),
        ));
    }

    if !req.tools.is_empty() && !config.enable_tool_calls {
        return Err((
            StatusCode::BAD_REQUEST,
            ApiError::new(
                "Tool calling is disabled on this server",
                "invalid_request_error",
            )
            .with_code("tool_calls_disabled")
            .with_param("tools"),
        ));
    }

    Ok(())
}

pub fn prepare_prompt(req: &ChatCompletionRequest, config: &ServerConfig) -> PreparedPrompt {
    let template = select_chat_template(req, config);
    let mut messages = Vec::with_capacity(req.messages.len() + 4);
    let has_system = req
        .messages
        .iter()
        .any(|message| matches!(message.role, Role::System | Role::Developer));

    // Default system prompt injection disabled.
    // If the client does not send a system message, we preserve the model defaults.
    if false && !has_system && !config.default_system_prompt.trim().is_empty() {
        messages.push(PromptMessage {
            role: Role::System,
            name: None,
            content: config.default_system_prompt.clone(),
        });
    }

    if let Some(format) = &req.response_format {
        append_response_format_instruction(&mut messages, format);
    }

    if !req.tools.is_empty() {
        append_tool_instructions(&mut messages, &req.tools, req.tool_choice.as_ref());
    }

    for msg in &req.messages {
        let mut content = if let Some(tool_calls) = &msg.tool_calls {
            serde_json::to_string(tool_calls).unwrap_or_else(|_| "[]".to_owned())
        } else {
            msg.text().to_owned()
        };

        if let Some(tool_call_id) = &msg.tool_call_id {
            content.push_str(" [tool_call_id=");
            content.push_str(tool_call_id);
            content.push(']');
        }

        messages.push(PromptMessage {
            role: msg.role.clone(),
            name: msg.name.clone(),
            content,
        });
    }

    let mut prompt = String::with_capacity(1024);
    if template == ChatTemplate::Llama {
        render_llama_prompt(&mut prompt, &messages);
    } else {
        for message in &messages {
            append_chat_message(
                &mut prompt,
                template,
                &message.role,
                message.name.as_deref(),
                &message.content,
            );
        }

        append_assistant_prefix(&mut prompt, template);
    }

    // For Generic template, inject stop sequences to prevent transcript leakage
    // (model may continue generating "User:", "System:", "Assistant:" turns)
    // NOTE: Only use role-prefix stops — avoid stops like "\n1." that cut legitimate numbered lists.
    let extra_stops: Vec<String> = if template == ChatTemplate::Generic && req.tools.is_empty() {
        vec![
            "\nUser:".to_owned(),
            "\nSystem:".to_owned(),
            "\nAssistant:".to_owned(),
            "\n# User:".to_owned(),
            "\n# System:".to_owned(),
            "\n# Assistant:".to_owned(),
            "\n## System:".to_owned(),
            "\n\nUser:".to_owned(),
            "\n\nSystem:".to_owned(),
            "\n\nAssistant:".to_owned(),
            "\nHuman:".to_owned(),
            "\n\nHuman:".to_owned(),
            "\nDeveloper:".to_owned(),
        ]
    } else {
        vec![]
    };

    PreparedPrompt {
        prompt,
        streaming_mode: if req.tools.is_empty() {
            StreamingMode::LiveTokens
        } else {
            StreamingMode::Buffered
        },
        extra_stop_sequences: extra_stops,
    }
}

pub fn finalize_assistant_message(
    raw_output: String,
    req: &ChatCompletionRequest,
) -> FinalizedAssistant {
    if !req.tools.is_empty() {
        if let Some(tool_calls) = parse_tool_calls(&raw_output) {
            return FinalizedAssistant {
                content: None,
                tool_calls: Some(tool_calls),
                finish_reason: FinishReason::ToolCalls,
            };
        }
    }

    FinalizedAssistant {
        content: Some(raw_output.trim().to_owned()),
        tool_calls: None,
        finish_reason: FinishReason::Stop,
    }
}

pub fn tool_calls_to_delta(tool_calls: &[ToolCall]) -> Vec<DeltaToolCall> {
    tool_calls
        .iter()
        .enumerate()
        .map(|(index, tool_call)| DeltaToolCall {
            index: index as u32,
            id: Some(tool_call.id.clone()),
            tool_type: Some(tool_call.tool_type.clone()),
            function: Some(DeltaToolFunctionCall {
                name: Some(tool_call.function.name.clone()),
                arguments: Some(tool_call.function.arguments.clone()),
            }),
        })
        .collect()
}

fn select_chat_template(req: &ChatCompletionRequest, config: &ServerConfig) -> ChatTemplate {
    let mut candidates = vec![req.model.as_str()];
    if let Some(model) = config.resolve_model(&req.model) {
        if !model.id.eq_ignore_ascii_case(req.model.as_str()) {
            candidates.push(model.id.as_str());
        }
        if !model.root.eq_ignore_ascii_case(req.model.as_str())
            && !model.root.eq_ignore_ascii_case(model.id.as_str())
        {
            candidates.push(model.root.as_str());
        }
    }

    for candidate in candidates {
        let normalized = candidate.to_ascii_lowercase();

        if normalized.contains("qwen") || normalized.contains("cooper") || normalized.contains("iron") || normalized.contains("steel") {
            return ChatTemplate::Qwen;
        }
        if normalized.contains("gemma") {
            return ChatTemplate::Gemma;
        }
        if normalized.contains("llama") {
            return ChatTemplate::Llama;
        }
        if normalized.contains("kimi") || normalized.contains("moonshot") {
            return ChatTemplate::Kimi;
        }
    }

    ChatTemplate::Generic
}

fn append_response_format_instruction(messages: &mut Vec<PromptMessage>, format: &ResponseFormat) {
    match format.format_type.as_str() {
        "json_object" => {
            messages.push(PromptMessage {
                role: Role::System,
                name: None,
                content: "Return valid JSON only. Do not wrap it in markdown fences.".to_owned(),
            });
        }
        "json_schema" => {
            let mut content = "Return valid JSON only and follow this schema exactly: ".to_owned();
            if let Some(schema) = &format.json_schema {
                content.push_str(&schema.to_string());
            }
            messages.push(PromptMessage {
                role: Role::System,
                name: None,
                content,
            });
        }
        _ => {}
    }
}

fn append_tool_instructions(
    messages: &mut Vec<PromptMessage>,
    tools: &[ToolDefinition],
    tool_choice: Option<&ToolChoice>,
) {
    let tool_summaries: Vec<Value> = tools
        .iter()
        .map(|tool| {
            json!({
                "name": tool.function.name,
                "description": tool.function.description,
                "parameters": tool.function.parameters,
            })
        })
        .collect();

    let mut available_tools = "Available tools: ".to_owned();
    available_tools
        .push_str(&serde_json::to_string(&tool_summaries).unwrap_or_else(|_| "[]".to_owned()));
    messages.push(PromptMessage {
        role: Role::System,
        name: None,
        content: available_tools,
    });

    let tool_usage = format!(
        "If you need a tool, respond with JSON only using one of these shapes: {} or {}. If no tool is needed, answer normally.",
        r#"{"tool_calls":[{"name":"tool_name","arguments":{}}]}"#,
        r#"{"name":"tool_name","arguments":{}}"#
    );
    messages.push(PromptMessage {
        role: Role::System,
        name: None,
        content: tool_usage,
    });

    match tool_choice {
        Some(ToolChoice::Mode(mode)) if mode == "required" => {
            messages.push(PromptMessage {
                role: Role::System,
                name: None,
                content: "You must call a tool before replying.".to_owned(),
            });
        }
        Some(ToolChoice::Mode(mode)) if mode == "none" => {
            messages.push(PromptMessage {
                role: Role::System,
                name: None,
                content: "Do not call tools in this response.".to_owned(),
            });
        }
        Some(ToolChoice::Specific(choice)) => {
            messages.push(PromptMessage {
                role: Role::System,
                name: None,
                content: format!("You must call the tool named {}.", choice.function.name),
            });
        }
        _ => {}
    }
}

fn append_chat_message(
    prompt: &mut String,
    template: ChatTemplate,
    role: &Role,
    name: Option<&str>,
    content: &str,
) {
    let content = if let Some(name) = name {
        format!("{} ({})", content, name)
    } else {
        content.to_owned()
    };

    match template {
        ChatTemplate::Generic => {
            prompt.push_str(generic_role_name(&role));
            prompt.push_str(": ");
            prompt.push_str(&content);
            prompt.push('\n');
        }
        ChatTemplate::Qwen => {
            prompt.push_str("<|im_start|>");
            prompt.push_str(qwen_role_name(&role));
            prompt.push('\n');
            prompt.push_str(&content);
            prompt.push_str("<|im_end|>\n");
        }
        ChatTemplate::Gemma => {
            prompt.push_str("<start_of_turn>");
            prompt.push_str(gemma_role_name(&role));
            prompt.push('\n');
            prompt.push_str(&content);
            prompt.push_str("<end_of_turn>\n");
        }
        ChatTemplate::Llama => match role {
            Role::Assistant => {
                prompt.push_str(&content);
                prompt.push('\n');
            }
            _ => {
                prompt.push_str("[INST] ");
                if matches!(role, Role::System | Role::Developer) {
                    prompt.push_str("<<SYS>>\n");
                    prompt.push_str(&content);
                    prompt.push_str("\n<</SYS>>");
                } else {
                    prompt.push_str(&content);
                }
                prompt.push_str(" [/INST]\n");
            }
        },
        ChatTemplate::Kimi => {
            prompt.push_str(kimi_role_marker(&role));
            prompt.push('\n');
            prompt.push_str(&content);
            prompt.push('\n');
        }
    }
}

fn render_llama_prompt(prompt: &mut String, messages: &[PromptMessage]) {
    let mut pending_system = Vec::new();

    for message in messages {
        let content = if let Some(name) = message.name.as_deref() {
            format!("{} ({})", message.content, name)
        } else {
            message.content.clone()
        };

        match message.role {
            Role::System | Role::Developer => pending_system.push(content),
            Role::User | Role::Tool => {
                let mut block = String::new();
                if !pending_system.is_empty() {
                    block.push_str("<<SYS>>\n");
                    block.push_str(&pending_system.join("\n\n"));
                    block.push_str("\n<</SYS>>\n\n");
                    pending_system.clear();
                }
                block.push_str(&content);
                prompt.push_str("[INST] ");
                prompt.push_str(&block);
                prompt.push_str(" [/INST]\n");
            }
            Role::Assistant => {
                prompt.push_str(&content);
                prompt.push('\n');
            }
        }
    }

    if !pending_system.is_empty() {
        prompt.push_str("[INST] <<SYS>>\n");
        prompt.push_str(&pending_system.join("\n\n"));
        prompt.push_str("\n<</SYS>> [/INST]\n");
    }
}

fn append_assistant_prefix(prompt: &mut String, template: ChatTemplate) {
    match template {
        ChatTemplate::Generic => prompt.push_str("Assistant:"),
        ChatTemplate::Qwen => prompt.push_str("<|im_start|>assistant\n"),
        ChatTemplate::Gemma => prompt.push_str("<start_of_turn>model\n"),
        ChatTemplate::Llama => prompt.push_str("[INST] "),
        ChatTemplate::Kimi => prompt.push_str("<|im_assistant|>\n"),
    }
}

fn generic_role_name(role: &Role) -> &'static str {
    match role {
        Role::System => "System",
        Role::Developer => "Developer",
        Role::User => "User",
        Role::Assistant => "Assistant",
        Role::Tool => "Tool",
    }
}

fn qwen_role_name(role: &Role) -> &'static str {
    match role {
        Role::System | Role::Developer => "system",
        Role::User => "user",
        Role::Assistant => "assistant",
        Role::Tool => "tool",
    }
}

fn gemma_role_name(role: &Role) -> &'static str {
    match role {
        Role::Assistant => "model",
        Role::System | Role::Developer | Role::User | Role::Tool => "user",
    }
}

fn kimi_role_marker(role: &Role) -> &'static str {
    match role {
        Role::System | Role::Developer => "<|im_system|>",
        Role::User => "<|im_user|>",
        Role::Assistant => "<|im_assistant|>",
        Role::Tool => "<|im_tool|>",
    }
}

fn parse_tool_calls(raw_output: &str) -> Option<Vec<ToolCall>> {
    let trimmed = raw_output.trim();
    if trimmed.is_empty() {
        return None;
    }

    let value: Value = serde_json::from_str(trimmed).ok()?;
    match value {
        Value::Object(map) => {
            if let Some(tool_calls) = map.get("tool_calls").and_then(Value::as_array) {
                let parsed: Vec<ToolCall> = tool_calls
                    .iter()
                    .enumerate()
                    .filter_map(|(index, item)| parse_tool_call(item, index))
                    .collect();
                if parsed.is_empty() {
                    None
                } else {
                    Some(parsed)
                }
            } else {
                parse_tool_call(&Value::Object(map), 0).map(|tool_call| vec![tool_call])
            }
        }
        _ => None,
    }
}

fn parse_tool_call(value: &Value, index: usize) -> Option<ToolCall> {
    let object = value.as_object()?;
    let function = object.get("function").unwrap_or(value);
    let function_obj = function.as_object()?;
    let name = function_obj.get("name")?.as_str()?.trim();
    if name.is_empty() {
        return None;
    }

    let arguments = function_obj
        .get("arguments")
        .cloned()
        .unwrap_or_else(|| Value::Object(Default::default()));
    let arguments = match arguments {
        Value::String(s) => s,
        other => serde_json::to_string(&other).ok()?,
    };

    Some(ToolCall {
        id: format!("call_{}", index + 1),
        tool_type: ToolType::Function,
        function: ToolFunctionCall {
            name: name.to_owned(),
            arguments,
        },
    })
}

pub fn build_sampler(req: &ChatCompletionRequest) -> inference_core::SamplerConfig {
    let mut sampler = inference_core::SamplerConfig::default();
    sampler.temperature = req.temperature.clamp(0.0, 2.0);
    sampler.top_p = req.top_p.clamp(0.0, 1.0);
    sampler.max_tokens = req.max_tokens.max(1);
    sampler.stop_sequences = req.stop.clone();
    // Slightly higher repeat_penalty (1.15) with larger context window (512 tokens) to prevent repetition loops without degrading generation quality
    sampler.repeat_penalty = if req.tools.is_empty() { 1.15 } else { 1.0 };
    sampler.repeat_last_n = 512;
    sampler.seed = req.seed.unwrap_or(0);
    sampler
}

fn env_csv(key: &str) -> Vec<String> {
    std::env::var(key)
        .ok()
        .map(|value| {
            value
                .split(',')
                .map(str::trim)
                .filter(|part| !part.is_empty())
                .map(ToOwned::to_owned)
                .collect()
        })
        .unwrap_or_default()
}

fn env_bool(key: &str, default: bool) -> bool {
    std::env::var(key)
        .ok()
        .and_then(|value| match value.trim().to_ascii_lowercase().as_str() {
            "1" | "true" | "yes" | "on" => Some(true),
            "0" | "false" | "no" | "off" => Some(false),
            _ => None,
        })
        .unwrap_or(default)
}

pub fn ensure_supported_mode(req: &ChatCompletionRequest) -> Result<()> {
    if matches!(req.tool_choice, Some(ToolChoice::Mode(ref mode)) if mode == "none") {
        return Ok(());
    }

    if let Some(format) = &req.response_format {
        if !matches!(
            format.format_type.as_str(),
            "json_object" | "json_schema" | "text"
        ) {
            bail!("unsupported response_format type: {}", format.format_type);
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{
        authorize_request, finalize_assistant_message, prepare_prompt, validate_request,
        FinalizedAssistant, ModelCatalog, ServerConfig, StreamingMode,
    };
    use api_contracts::chat::ChatCompletionRequest;
    use axum::http::{HeaderMap, HeaderValue, StatusCode};

    fn base_config() -> ServerConfig {
        ServerConfig {
            api_key: None,
            auth_header: "authorization".to_owned(),
            auth_scheme: Some("Bearer".to_owned()),
            cors_origins: vec!["*".to_owned()],
            default_system_prompt: "Default coding system prompt".to_owned(),
            enable_tool_calls: true,
            model_catalog: ModelCatalog::new(
                "real-model".to_owned(),
                vec!["alias-model".to_owned()],
            ),
        }
    }

    fn parse_request(json: &str) -> ChatCompletionRequest {
        serde_json::from_str(json).unwrap()
    }

    #[test]
    fn resolves_model_aliases_case_insensitively() {
        let catalog = ModelCatalog::new("real-model".to_owned(), vec!["Alias-Model".to_owned()]);

        assert_eq!(catalog.resolve("alias-model").unwrap().root, "real-model");
        assert_eq!(catalog.resolve("REAL-MODEL").unwrap().id, "real-model");
        assert!(catalog.resolve("missing").is_none());
    }

    #[test]
    fn validates_model_and_n() {
        let config = base_config();
        let request = parse_request(
            r#"{
                "model": "missing-model",
                "messages": [{"role":"user","content":"hi"}]
            }"#,
        );

        let err = validate_request(&request, &config).unwrap_err();
        assert_eq!(err.0, StatusCode::NOT_FOUND);

        let request = parse_request(
            r#"{
                "model": "alias-model",
                "n": 2,
                "messages": [{"role":"user","content":"hi"}]
            }"#,
        );
        let err = validate_request(&request, &config).unwrap_err();
        assert_eq!(err.0, StatusCode::BAD_REQUEST);
    }

    #[test]
    fn authorizes_bearer_and_x_api_key_headers() {
        let mut config = base_config();
        config.api_key = Some("secret".to_owned());

        let mut headers = HeaderMap::new();
        headers.insert("authorization", HeaderValue::from_static("Bearer secret"));
        assert!(authorize_request(&headers, &config).is_ok());

        let mut headers = HeaderMap::new();
        headers.insert("x-api-key", HeaderValue::from_static("secret"));
        assert!(authorize_request(&headers, &config).is_ok());
    }

    #[test]
    fn adds_default_system_prompt_and_tool_instructions() {
        let config = base_config();
        let request = parse_request(
            r#"{
                "model": "alias-model",
                "messages": [{"role":"user","content":"List files"}],
                "tools": [{
                    "type": "function",
                    "function": {
                        "name": "list_files",
                        "description": "List project files",
                        "parameters": {"type": "object"}
                    }
                }],
                "tool_choice": "required"
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert_eq!(prepared.streaming_mode, StreamingMode::Buffered);
        assert!(prepared.prompt.contains("list_files"));
        assert!(prepared.prompt.contains("must call a tool"));
    }

    #[test]
    fn extracts_tool_calls_from_model_output() {
        let request = parse_request(
            r#"{
                "model": "alias-model",
                "messages": [{"role":"user","content":"List files"}],
                "tools": [{
                    "type": "function",
                    "function": {
                        "name": "list_files",
                        "parameters": {"type": "object"}
                    }
                }]
            }"#,
        );

        let finalized = finalize_assistant_message(
            r#"{"tool_calls":[{"name":"list_files","arguments":{"path":"src"}}]}"#.to_owned(),
            &request,
        );

        assert!(matches!(
            finalized.finish_reason,
            api_contracts::chat::FinishReason::ToolCalls
        ));
        assert_eq!(finalized.content, None);
        let tool_calls = finalized.tool_calls.unwrap();
        assert_eq!(tool_calls[0].function.name, "list_files");
        assert_eq!(tool_calls[0].function.arguments, "{\"path\":\"src\"}");
    }

    #[test]
    fn preserves_plain_text_when_no_tool_call_is_present() {
        let request = parse_request(
            r#"{
                "model": "alias-model",
                "messages": [{"role":"user","content":"Explain borrow checker"}]
            }"#,
        );

        let finalized =
            finalize_assistant_message("  Use references and lifetimes.  ".to_owned(), &request);
        assert!(matches!(
            finalized,
            FinalizedAssistant {
                content: Some(_),
                tool_calls: None,
                ..
            }
        ));
        assert_eq!(
            finalized.content.as_deref(),
            Some("Use references and lifetimes.")
        );
    }

    #[test]
    fn generic_template_injects_markdown_transcript_stop_sequences() {
        let config = base_config();
        let request = parse_request(
            r#"{
                "model": "alias-model",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared.extra_stop_sequences.contains(&"\n# System:".to_owned()));
        assert!(prepared.extra_stop_sequences.contains(&"\n# User:".to_owned()));
        assert!(prepared.extra_stop_sequences.contains(&"\n# Assistant:".to_owned()));
        assert!(prepared.extra_stop_sequences.contains(&"\n## System:".to_owned()));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_qwen_family_uses_im_tokens() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new("Qwen2.5-Coder-7B-Instruct".to_owned(), vec![]),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "Qwen2.5-Coder-7B-Instruct",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared.prompt.contains("<|im_start|>"));
        assert!(prepared.prompt.contains("<|im_end|>"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_gemma_family_uses_turn_tokens() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new("gemma-2-9b-it".to_owned(), vec![]),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "gemma-2-9b-it",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared.prompt.contains("<start_of_turn>user"));
        assert!(prepared.prompt.contains("<start_of_turn>model"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_llama_family_uses_inst_blocks() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new("Meta-Llama-3-8B-Instruct".to_owned(), vec![]),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "Meta-Llama-3-8B-Instruct",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared.prompt.contains("[INST]"));
        assert!(prepared.prompt.contains("[/INST]"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_llama_family_combines_system_and_user_in_first_inst_block() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new("Meta-Llama-3-8B-Instruct".to_owned(), vec![]),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "Meta-Llama-3-8B-Instruct",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared.prompt.contains("[INST] <<SYS>>"));
        assert!(prepared.prompt.contains("hello [/INST]"));
        assert_eq!(prepared.prompt.matches("[INST]").count(), 1);
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_llama_family_does_not_end_with_open_inst_prefix() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new("Meta-Llama-3-8B-Instruct".to_owned(), vec![]),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "Meta-Llama-3-8B-Instruct",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(!prepared.prompt.ends_with("[INST] "));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_kimi_family_uses_moonshot_markers() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new("Kimi-K2-Instruct".to_owned(), vec![]),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "Kimi-K2-Instruct",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared.prompt.contains("<|im_user|>"));
        assert!(prepared.prompt.contains("<|im_assistant|>"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_unknown_models_keep_generic_fallback() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new("my-custom-model".to_owned(), vec![]),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "my-custom-model",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared
            .prompt
            .contains("System: Default coding system prompt"));
        assert!(prepared.prompt.contains("User: hello"));
        assert!(prepared.prompt.ends_with("Assistant:"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_alias_family_can_override_generic_canonical_id() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new(
                "my-local-model".to_owned(),
                vec!["qwen-coder".to_owned()],
            ),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "qwen-coder",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared.prompt.contains("<|im_start|>system\n"));
        assert!(prepared.prompt.ends_with("<|im_start|>assistant\n"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_generic_alias_stays_generic_without_family_hint() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new(
                "my-local-model".to_owned(),
                vec!["assistant-prod".to_owned()],
            ),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "assistant-prod",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared
            .prompt
            .contains("System: Default coding system prompt"));
        assert!(prepared.prompt.contains("User: hello"));
        assert!(prepared.prompt.ends_with("Assistant:"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_moonshot_family_uses_kimi_markers() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new("moonshotai/kimi-k2-instruct".to_owned(), vec![]),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "moonshot-v1-8k",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared.prompt.contains("<|im_system|>"));
        assert!(prepared.prompt.contains("<|im_user|>"));
        assert!(prepared.prompt.contains("<|im_assistant|>"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn dynamic_template_moonshot_like_name_without_family_hint_stays_generic() {
        let config = ServerConfig {
            model_catalog: ModelCatalog::new("moonshine-local".to_owned(), vec![]),
            ..base_config()
        };
        let request = parse_request(
            r#"{
                "model": "moonshine-local",
                "messages": [{"role":"user","content":"hello"}]
            }"#,
        );

        let prepared = prepare_prompt(&request, &config);
        assert!(prepared
            .prompt
            .contains("System: Default coding system prompt"));
        assert!(prepared.prompt.contains("User: hello"));
        assert!(prepared.prompt.ends_with("Assistant:"));
    }
}