//! Request handlers for the API endpoints.
//!
//! This module contains the main request handlers and supporting functions
//! for processing chat requests, including both streaming and non-streaming
//! responses. It coordinates between different AI models and handles
//! usage tracking and cost calculations.

use crate::{
    clients::{AnthropicClient, DeepSeekClient},
    config::Config,
    error::{ApiError, Result, SseResponse},
    models::{
        ApiRequest, ApiResponse, ContentBlock, CombinedUsage, DeepSeekUsage, AnthropicUsage,
        ExternalApiResponse, Message, Role, StreamEvent,
    },
};
use axum::{
    extract::State,
    response::{sse::Event, IntoResponse},
    Json,
};
use chrono::Utc;
use futures::StreamExt;
use std::{sync::Arc, collections::HashMap};
use tokio_stream::wrappers::ReceiverStream;

/// Application state shared across request handlers.
///
/// Contains configuration that needs to be accessible
/// to all request handlers.
pub struct AppState {
    pub config: Config,
}

/// Extracts API tokens from request headers.
///
/// # Arguments
///
/// * `headers` - The HTTP headers containing the API tokens
///
/// # Returns
///
/// * `Result<(String, String)>` - A tuple of (DeepSeek token, Anthropic token)
///
/// # Errors
///
/// Returns `ApiError::MissingHeader` if either token is missing
/// Returns `ApiError::BadRequest` if tokens are malformed
fn extract_api_tokens(
    headers: &axum::http::HeaderMap,
) -> Result<(String, String)> {
    let deepseek_token = headers
        .get("X-DeepSeek-API-Token")
        .ok_or_else(|| ApiError::MissingHeader {
            header: "X-DeepSeek-API-Token".to_string()
        })?
        .to_str()
        .map_err(|_| ApiError::BadRequest {
            message: "Invalid DeepSeek API token".to_string()
        })?
        .to_string();

    let anthropic_token = headers
        .get("X-Anthropic-API-Token")
        .ok_or_else(|| ApiError::MissingHeader {
            header: "X-Anthropic-API-Token".to_string()
        })?
        .to_str()
        .map_err(|_| ApiError::BadRequest {
            message: "Invalid Anthropic API token".to_string()
        })?
        .to_string();

    Ok((deepseek_token, anthropic_token))
}

/// Calculates the cost of DeepSeek API usage.
///
/// # Arguments
///
/// * `input_tokens` - Number of input tokens processed
/// * `output_tokens` - Number of output tokens generated
/// * `_reasoning_tokens` - Number of tokens used for reasoning
/// * `cached_tokens` - Number of tokens retrieved from cache
/// * `config` - Configuration containing pricing information
///
/// # Returns
///
/// The total cost in dollars for the API usage
fn calculate_deepseek_cost(
    input_tokens: u32,
    output_tokens: u32,
    _reasoning_tokens: u32,
    cached_tokens: u32,
    config: &Config,
) -> f64 {
    let cache_hit_cost = (cached_tokens as f64 / 1_000_000.0) * config.pricing.deepseek.input_cache_hit_price;
    let cache_miss_cost = ((input_tokens - cached_tokens) as f64 / 1_000_000.0) * config.pricing.deepseek.input_cache_miss_price;
    let output_cost = (output_tokens as f64 / 1_000_000.0) * config.pricing.deepseek.output_price;

    cache_hit_cost + cache_miss_cost + output_cost
}

/// Calculates the cost of Anthropic API usage.
///
/// # Arguments
///
/// * `model` - The specific Claude model used
/// * `input_tokens` - Number of input tokens processed
/// * `output_tokens` - Number of output tokens generated
/// * `cache_write_tokens` - Number of tokens written to cache
/// * `cache_read_tokens` - Number of tokens read from cache
/// * `config` - Configuration containing pricing information
///
/// # Returns
///
/// The total cost in dollars for the API usage
fn calculate_anthropic_cost(
    model: &str,
    input_tokens: u32,
    output_tokens: u32,
    cache_write_tokens: u32,
    cache_read_tokens: u32,
    config: &Config,
) -> f64 {
    let pricing = if model.contains("claude-3-5-sonnet") {
        &config.pricing.anthropic.claude_3_sonnet
    } else if model.contains("claude-3-5-haiku") {
        &config.pricing.anthropic.claude_3_haiku
    } else if model.contains("claude-3-opus") {
        &config.pricing.anthropic.claude_3_opus
    } else {
        &config.pricing.anthropic.claude_3_sonnet // default to sonnet pricing
    };

    let input_cost = (input_tokens as f64 / 1_000_000.0) * pricing.input_price;
    let output_cost = (output_tokens as f64 / 1_000_000.0) * pricing.output_price;
    let cache_write_cost = (cache_write_tokens as f64 / 1_000_000.0) * pricing.cache_write_price;
    let cache_read_cost = (cache_read_tokens as f64 / 1_000_000.0) * pricing.cache_read_price;

    input_cost + output_cost + cache_write_cost + cache_read_cost
}

/// Formats a cost value as a dollar amount string.
///
/// # Arguments
///
/// * `cost` - The cost value to format
///
/// # Returns
///
/// A string representing the cost with 3 decimal places and $ prefix
fn format_cost(cost: f64) -> String {
    format!("${:.3}", cost)
}

/// Main handler for chat requests.
///
/// Routes requests to either streaming or non-streaming handlers
/// based on the request configuration.
///
/// # Arguments
///
/// * `state` - Application state containing configuration
/// * `headers` - HTTP request headers
/// * `request` - The parsed chat request
///
/// # Returns
///
/// * `Result<Response>` - The API response or an error
pub async fn handle_chat(
    state: State<Arc<AppState>>,
    headers: axum::http::HeaderMap,
    Json(request): Json<ApiRequest>,
) -> Result<axum::response::Response> {
    if request.stream {
        let stream_response = chat_stream(state, headers, Json(request)).await?;
        Ok(stream_response.into_response())
    } else {
        let json_response = chat(state, headers, Json(request)).await?;
        Ok(json_response.into_response())
    }
}

/// Handler for non-streaming chat requests.
///
/// Processes the request through both AI models sequentially,
/// combining their responses and tracking usage.
///
/// # Arguments
///
/// * `state` - Application state containing configuration
/// * `headers` - HTTP request headers
/// * `request` - The parsed chat request
///
/// # Returns
///
/// * `Result<Json<ApiResponse>>` - The combined API response or an error
pub(crate) async fn chat(
    State(state): State<Arc<AppState>>,
    headers: axum::http::HeaderMap,
    Json(request): Json<ApiRequest>,
) -> Result<Json<ApiResponse>> {
    // Validate system prompt
    if !request.validate_system_prompt() {
        return Err(ApiError::InvalidSystemPrompt);
    }

    // Extract API tokens
    let (deepseek_token, anthropic_token) = extract_api_tokens(&headers)?;

    // Initialize clients
    let deepseek_client = DeepSeekClient::new(deepseek_token);
    let anthropic_client = AnthropicClient::new(anthropic_token);

    // Get messages with system prompt
    let messages = request.get_messages_with_system();

    // Call DeepSeek API
    let deepseek_response = deepseek_client.chat(messages.clone(), &request.deepseek_config).await?;

    // Store response metadata
    let deepseek_status: u16 = 200;
    let deepseek_headers = HashMap::new(); // Headers not available when using high-level chat method

    // Extract reasoning content and wrap in thinking tags
    let reasoning_content = deepseek_response
        .choices
        .first()
        .and_then(|c| c.message.reasoning_content.as_ref())
        .ok_or_else(|| ApiError::DeepSeekError {
            message: "No reasoning content in response".to_string(),
            type_: "missing_content".to_string(),
            param: None,
            code: None
        })?;

    let thinking_content = format!("<thinking>\n{}\n</thinking>", reasoning_content);

    // Add thinking content to messages for Anthropic
    let mut anthropic_messages = messages;
    anthropic_messages.push(Message {
        role: Role::Assistant,
        content: thinking_content.clone(),
    });

    // Call Anthropic API
    let anthropic_response = anthropic_client.chat(
        anthropic_messages,
        request.get_system_prompt().map(String::from),
        &request.anthropic_config
    ).await?;

    // Store response metadata
    let anthropic_status: u16 = 200;
    let anthropic_headers = HashMap::new(); // Headers not available when using high-level chat method

    // Calculate usage costs
    let deepseek_cost = calculate_deepseek_cost(
        deepseek_response.usage.prompt_tokens,
        deepseek_response.usage.completion_tokens,
        deepseek_response.usage.completion_tokens_details.reasoning_tokens,
        deepseek_response.usage.prompt_tokens_details.cached_tokens,
        &state.config,
    );

    let anthropic_cost = calculate_anthropic_cost(
        &anthropic_response.model,
        anthropic_response.usage.input_tokens,
        anthropic_response.usage.output_tokens,
        anthropic_response.usage.cache_creation_input_tokens,
        anthropic_response.usage.cache_read_input_tokens,
        &state.config,
    );

    // Combine thinking content with Anthropic's response
    let mut content = Vec::new();

    // Add thinking block first
    content.push(ContentBlock::text(thinking_content));

    // Add Anthropic's response blocks
    content.extend(anthropic_response.content.clone().into_iter()
        .map(ContentBlock::from_anthropic));

    // Build response with captured headers
    let response = ApiResponse {
        created: Utc::now(),
        content,
        deepseek_response: request.verbose.then(|| ExternalApiResponse {
            status: deepseek_status,
            headers: deepseek_headers,
            body: serde_json::to_value(&deepseek_response).unwrap_or_default(),
        }),
        anthropic_response: request.verbose.then(|| ExternalApiResponse {
            status: anthropic_status,
            headers: anthropic_headers,
            body: serde_json::to_value(&anthropic_response).unwrap_or_default(),
        }),
        combined_usage: CombinedUsage {
            total_cost: format_cost(deepseek_cost + anthropic_cost),
            deepseek_usage: DeepSeekUsage {
                input_tokens: deepseek_response.usage.prompt_tokens,
                output_tokens: deepseek_response.usage.completion_tokens,
                reasoning_tokens: deepseek_response.usage.completion_tokens_details.reasoning_tokens,
                cached_input_tokens: deepseek_response.usage.prompt_tokens_details.cached_tokens,
                total_tokens: deepseek_response.usage.total_tokens,
                total_cost: format_cost(deepseek_cost),
            },
            anthropic_usage: AnthropicUsage {
                input_tokens: anthropic_response.usage.input_tokens,
                output_tokens: anthropic_response.usage.output_tokens,
                cached_write_tokens: anthropic_response.usage.cache_creation_input_tokens,
                cached_read_tokens: anthropic_response.usage.cache_read_input_tokens,
                total_tokens: anthropic_response.usage.input_tokens + anthropic_response.usage.output_tokens,
                total_cost: format_cost(anthropic_cost),
            },
        },
    };

    Ok(Json(response))
}

/// Handler for streaming chat requests.
///
/// Processes the request through both AI models sequentially,
/// streaming their responses as Server-Sent Events.
///
/// # Arguments
///
/// * `state` - Application state containing configuration
/// * `headers` - HTTP request headers
/// * `request` - The parsed chat request
///
/// # Returns
///
/// * `Result<SseResponse>` - A stream of Server-Sent Events or an error

pub(crate) async fn chat_stream(
    State(state): State<Arc<AppState>>,
    headers: axum::http::HeaderMap,
    Json(request): Json<ApiRequest>,
) -> Result<SseResponse> {
    if !request.validate_system_prompt() {
        return Err(ApiError::InvalidSystemPrompt);
    }
    let (deepseek_token, anthropic_token) = extract_api_tokens(&headers)?;
    let deepseek_client = DeepSeekClient::new(deepseek_token);
    let anthropic_client = AnthropicClient::new(anthropic_token);
    let messages = request.get_messages_with_system();
    let (tx, rx) = tokio::sync::mpsc::channel(100);
    let tx = Arc::new(tx);
    let config = state.config.clone();
    let request_clone = request.clone();
    tokio::spawn(async move {
        let tx = tx.clone();
        let _ = tx
            .send(Ok(Event::default().event("start").data(
                serde_json::to_string(&StreamEvent::Start {
                    created: Utc::now(),
                })
                .unwrap_or_default(),
            )))
            .await;
        let _ = tx
            .send(Ok(Event::default().event("content").data(
                serde_json::to_string(&StreamEvent::Content {
                    content: vec![ContentBlock {
                        content_type: "text".to_string(),
                        text: "<thinking>\n".to_string(),
                    }],
                })
                .unwrap_or_default(),
            )))
            .await;
        let mut deepseek_usage = None;
        let mut complete_reasoning = String::new();
        let mut deepseek_stream = deepseek_client.chat_stream(messages.clone(), &request_clone.deepseek_config);
        while let Some(chunk) = deepseek_stream.next().await {
            match chunk {
                Ok(response) => {
                    if let Some(choice) = response.choices.first() {
                        if choice.delta.reasoning_content.is_none() {
                            break;
                        }
                        if let Some(reasoning) = &choice.delta.reasoning_content {
                            if !reasoning.is_empty() {
                                let _ = tx
                                    .send(Ok(Event::default().event("content").data(
                                        serde_json::to_string(&StreamEvent::Content {
                                            content: vec![ContentBlock {
                                                content_type: "text_delta".to_string(),
                                                text: reasoning.to_string(),
                                            }],
                                        })
                                        .unwrap_or_default(),
                                    )))
                                    .await;
                                complete_reasoning.push_str(reasoning);
                            }
                        }
                    }
                    if let Some(usage) = response.usage {
                        deepseek_usage = Some(usage);
                    }
                }
                Err(e) => {
                    let _ = tx
                        .send(Ok(Event::default().event("error").data(
                            serde_json::to_string(&StreamEvent::Error {
                                message: e.to_string(),
                                code: 500,
                            })
                            .unwrap_or_default(),
                        )))
                        .await;
                    return;
                }
            }
        }
        let _ = tx
            .send(Ok(Event::default().event("content").data(
                serde_json::to_string(&StreamEvent::Content {
                    content: vec![ContentBlock {
                        content_type: "text".to_string(),
                        text: "\n</thinking>".to_string(),
                    }],
                })
                .unwrap_or_default(),
            )))
            .await;
        let mut anthropic_messages = messages;
        anthropic_messages.push(Message {
            role: Role::Assistant,
            content: format!("<thinking>\n{}\n</thinking>", complete_reasoning),
        });
        let mut anthropic_stream = anthropic_client.chat_stream(
            anthropic_messages,
            request_clone.get_system_prompt().map(String::from),
            &request_clone.anthropic_config,
        );
        while let Some(chunk) = anthropic_stream.next().await {
            match chunk {
                Ok(event) => match event {
                    crate::clients::anthropic::StreamEvent::MessageStart { message } => {
                        if !message.content.is_empty() {
                            let _ = tx
                                .send(Ok(Event::default().event("content").data(
                                    serde_json::to_string(&StreamEvent::Content {
                                        content: message.content.into_iter()
                                            .map(ContentBlock::from_anthropic)
                                            .collect()
                                    })
                                    .unwrap_or_default(),
                                )))
                                .await;
                        }
                    }
                    crate::clients::anthropic::StreamEvent::ContentBlockDelta { delta, .. } => {
                        let _ = tx
                            .send(Ok(Event::default().event("content").data(
                                serde_json::to_string(&StreamEvent::Content {
                                    content: vec![ContentBlock {
                                        content_type: delta.delta_type,
                                        text: delta.text,
                                    }],
                                })
                                .unwrap_or_default(),
                            )))
                            .await;
                    }
                    crate::clients::anthropic::StreamEvent::MessageDelta { usage, .. } => {
                        if let Some(usage) = usage {
                            let anthropic_usage = AnthropicUsage::from_anthropic(usage);
                            let anthropic_cost = calculate_anthropic_cost(
                                "claude-3-5-sonnet-20241022",
                                anthropic_usage.input_tokens,
                                anthropic_usage.output_tokens,
                                anthropic_usage.cached_write_tokens,
                                anthropic_usage.cached_read_tokens,
                                &config,
                            );
                            let (deepseek_usage, deepseek_cost) = if let Some(usage) = deepseek_usage.as_ref() {
                                let cost = calculate_deepseek_cost(
                                    usage.prompt_tokens,
                                    usage.completion_tokens,
                                    usage.completion_tokens_details.reasoning_tokens,
                                    usage.prompt_tokens_details.cached_tokens,
                                    &config,
                                );
                                (DeepSeekUsage {
                                    input_tokens: usage.prompt_tokens,
                                    output_tokens: usage.completion_tokens,
                                    reasoning_tokens: usage.completion_tokens_details.reasoning_tokens,
                                    cached_input_tokens: usage.prompt_tokens_details.cached_tokens,
                                    total_tokens: usage.total_tokens,
                                    total_cost: format_cost(cost),
                                }, cost)
                            } else {
                                (DeepSeekUsage {
                                    input_tokens: 0,
                                    output_tokens: 0,
                                    reasoning_tokens: 0,
                                    cached_input_tokens: 0,
                                    total_tokens: 0,
                                    total_cost: "$0.00".to_string(),
                                }, 0.0)
                            };
                            let _ = tx
                                .send(Ok(Event::default().event("usage").data(
                                    serde_json::to_string(&StreamEvent::Usage {
                                        usage: CombinedUsage {
                                            total_cost: format_cost(deepseek_cost + anthropic_cost),
                                            deepseek_usage,
                                            anthropic_usage: AnthropicUsage {
                                                input_tokens: anthropic_usage.input_tokens,
                                                output_tokens: anthropic_usage.output_tokens,
                                                cached_write_tokens: anthropic_usage.cached_write_tokens,
                                                cached_read_tokens: anthropic_usage.cached_read_tokens,
                                                total_tokens: anthropic_usage.total_tokens,
                                                total_cost: format_cost(anthropic_cost),
                                            },
                                        },
                                    })
                                    .unwrap_or_default(),
                                )))
                                .await;
                        }
                    }
                    _ => {}
                },
                Err(e) => {
                    let _ = tx
                        .send(Ok(Event::default().event("error").data(
                            serde_json::to_string(&StreamEvent::Error {
                                message: e.to_string(),
                                code: 500,
                            })
                            .unwrap_or_default(),
                        )))
                        .await;
                    return;
                }
            }
        }
        let _ = tx
            .send(Ok(Event::default().event("done").data(
                serde_json::to_string(&StreamEvent::Done)
                    .unwrap_or_default(),
            )))
            .await;
    });
    let stream = ReceiverStream::new(rx);
    Ok(SseResponse::new(stream))
}

//! Request models for the API endpoints.
//!
//! This module defines the structures used to represent incoming API requests,
//! including chat messages, configuration options, and request parameters.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Primary request structure for chat API endpoints.
///
/// This structure represents a complete chat request, including messages,
/// system prompts, and configuration options for both DeepSeek and Anthropic APIs.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ApiRequest {
    #[serde(default)]
    pub stream: bool,

    #[serde(default)]
    pub verbose: bool,

    pub system: Option<String>,
    pub messages: Vec<Message>,

    #[serde(default)]
    pub deepseek_config: ApiConfig,

    #[serde(default)]
    pub anthropic_config: ApiConfig,
}

/// A single message in a chat conversation.
///
/// Represents one message in the conversation history, including
/// its role (system, user, or assistant) and content.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Message {
    pub role: Role,
    pub content: String,
}

/// Possible roles for a message in a chat conversation.
///
/// Each message must be associated with one of these roles to
/// properly structure the conversation flow.
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    User,
    Assistant,
}

/// Configuration options for external API requests.
///
/// Contains headers and body parameters that will be passed
/// to the external AI model APIs.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct ApiConfig {
    #[serde(default)]
    pub headers: HashMap<String, String>,

    #[serde(default)]
    pub body: serde_json::Value,
}

impl ApiRequest {
    /// Validates that system prompts are not duplicated.
    ///
    /// Checks that a system prompt is not provided in both the root level
    /// and messages array. The system prompt itself is optional.
    ///
    /// # Returns
    ///
    /// * `bool` - True if system prompt validation passes (no duplicates), false otherwise
    pub fn validate_system_prompt(&self) -> bool {
        let system_in_messages = self.messages.iter().any(|msg| matches!(msg.role, Role::System));

        // Only invalid if system prompt is provided in both places
        !(self.system.is_some() && system_in_messages)
    }

    /// Returns messages with the system prompt in the correct position.
    ///
    /// Ensures the system prompt (if present) is the first message,
    /// followed by the conversation messages in order.
    ///
    /// # Returns
    ///
    /// * `Vec<Message>` - Messages with system prompt correctly positioned
    pub fn get_messages_with_system(&self) -> Vec<Message> {
        let mut messages = Vec::new();

        // Add system message first
        if let Some(system) = &self.system {
            messages.push(Message {
                role: Role::System,
                content: system.clone(),
            });
        }

        // Add remaining messages
        messages.extend(self.messages.iter().filter(|msg| !matches!(msg.role, Role::System)).cloned());

        messages
    }

    /// Retrieves the system prompt if one is present.
    ///
    /// Checks both the root level system field and the messages array
    /// for a system prompt.
    ///
    /// # Returns
    ///
    /// * `Option<&str>` - The system prompt if found, None otherwise
    pub fn get_system_prompt(&self) -> Option<&str> {
        self.system.as_deref().or_else(|| {
            self.messages
                .iter()
                .find(|msg| matches!(msg.role, Role::System))
                .map(|msg| msg.content.as_str())
        })
    }
}

pub(crate) async fn chat_stream (
    State(state): State<Arc<AppState>>,
    headers: axum::http::HeaderMap,
    Json(request): Json<ApiRequest>,
) -> Result<SseResponse> {}

This code segment functions to: Process streaming chat requests by asynchronously handling real-time data flow between two AI services, emit formatted events through a channel, aggregate incremental responses from both services, and calculate combined usage metrics while maintaining a continuous server-sent event stream.

```rust
[TASK_BEGIN]
pub(crate) async fn chat_stream(
    State(state): State<Arc<AppState>>,
    headers: axum::http::HeaderMap,
    Json(request): Json<ApiRequest>,
) -> Result<SseResponse> {
    let (tx, rx) = tokio::sync::mpsc::channel(128);

    tokio::spawn(async move {
        // Extract API tokens
        let (deepseek_token, anthropic_token) = match extract_api_tokens(&headers) {
            Ok(tokens) => tokens,
            Err(e) => {
                let _ = tx.send(Err(e)).await;
                return;
            }
        };

        // Initialize clients
        let deepseek_client = DeepSeekClient::new(deepseek_token);
        let anthropic_client = AnthropicClient::new(anthropic_token);
        let messages = request.get_messages_with_system();

        // Process DeepSeek stream
        let mut deepseek_stream = match deepseek_client.chat_stream(messages.clone(), &request.deepseek_config).await {
            Ok(stream) => stream,
            Err(e) => {
                let _ = tx.send(Err(e)).await;
                return;
            }
        };

        let mut reasoning_content = String::new();
        let mut deepseek_usage = None;

        while let Some(chunk) = deepseek_stream.next().await {
            match chunk {
                Ok(chunk) => {
                    if let Some(reasoning) = chunk.choices[0].delta.reasoning_content.as_deref() {
                        reasoning_content.push_str(reasoning);
                    }
                    deepseek_usage = chunk.usage;
                }
                Err(e) => {
                    let _ = tx.send(Err(e)).await;
                    return;
                }
            }
        }

        // Send thinking block
        let thinking_content = format!("<thinking>\n{}\n</thinking>", reasoning_content);
        if let Err(_) = tx.send(Ok(Event::default().data(thinking_content.clone()))).await {
            return;
        }

        // Prepare messages for Anthropic
        let mut anthropic_messages = messages;
        anthropic_messages.push(Message {
            role: Role::Assistant,
            content: thinking_content,
        });

        // Process Anthropic stream
        let mut anthropic_stream = match anthropic_client.chat_stream(
            anthropic_messages,
            request.get_system_prompt().map(String::from),
            &request.anthropic_config
        ).await {
            Ok(stream) => stream,
            Err(e) => {
                let _ = tx.send(Err(e)).await;
                return;
            }
        };

        let mut anthropic_usage = None;
        let mut anthropic_model = None;

        while let Some(chunk) = anthropic_stream.next().await {
            match chunk {
                Ok(chunk) => {
                    // Capture model from first chunk
                    if anthropic_model.is_none() {
                        anthropic_model = Some(chunk.model.clone());
                    }

                    // Send content delta
                    let event: StreamEvent = chunk.into();
                    match Event::default().json_data(event) {
                        Ok(event) => {
                            if let Err(_) = tx.send(Ok(event)).await {
                                return;
                            }
                        }
                        Err(e) => {
                            let _ = tx.send(Err(ApiError::InternalError {
                                message: e.to_string()
                            })).await;
                            return;
                        }
                    }

                    // Capture final usage
                    anthropic_usage = chunk.usage;
                }
                Err(e) => {
                    let _ = tx.send(Err(e)).await;
                    return;
                }
            }
        }

        // Calculate combined usage
        let (deepseek_cost, anthropic_cost) = match (deepseek_usage, anthropic_usage, anthropic_model) {
            (Some(ds_usage), Some(an_usage), Some(model)) => {
                let ds_cost = calculate_deepseek_cost(
                    ds_usage.prompt_tokens,
                    ds_usage.completion_tokens,
                    ds_usage.completion_tokens_details.reasoning_tokens,
                    ds_usage.prompt_tokens_details.cached_tokens,
                    &state.config,
                );

                let an_cost = calculate_anthropic_cost(
                    &model,
                    an_usage.input_tokens,
                    an_usage.output_tokens,
                    an_usage.cache_creation_input_tokens,
                    an_usage.cache_read_input_tokens,
                    &state.config,
                );

                (ds_cost, an_cost)
            }
            _ => {
                let _ = tx.send(Err(ApiError::InternalError {
                    message: "Missing usage data in stream".to_string()
                })).await;
                return;
            }
        };

        // Send final usage event
        let combined_usage = CombinedUsage {
            total_cost: format_cost(deepseek_cost + anthropic_cost),
            deepseek_usage: DeepSeekUsage {
                input_tokens: deepseek_usage.unwrap().prompt_tokens,
                output_tokens: deepseek_usage.unwrap().completion_tokens,
                reasoning_tokens: deepseek_usage.unwrap().completion_tokens_details.reasoning_tokens,
                cached_input_tokens: deepseek_usage.unwrap().prompt_tokens_details.cached_tokens,
                total_tokens: deepseek_usage.unwrap().total_tokens,
                total_cost: format_cost(deepseek_cost),
            },
            anthropic_usage: AnthropicUsage {
                input_tokens: anthropic_usage.unwrap().input_tokens,
                output_tokens: anthropic_usage.unwrap().output_tokens,
                cached_write_tokens: anthropic_usage.unwrap().cache_creation_input_tokens,
                cached_read_tokens: anthropic_usage.unwrap().cache_read_input_tokens,
                total_tokens: anthropic_usage.unwrap().input_tokens + anthropic_usage.unwrap().output_tokens,
                total_cost: format_cost(anthropic_cost),
            },
        };

        if let Err(_) = tx.send(Ok(Event::default().json_data(StreamEvent::Usage(combined_usage)).unwrap())).await {
            return;
        }
    });

    Ok(ReceiverStream::new(rx))
}
[TASK_END]
```

pub(crate) async fn chat_stream(
    State(state): State<Arc<AppState>>,
    headers: axum::http::HeaderMap,
    Json(request): Json<ApiRequest>,
) -> Result<SseResponse> {
    let (tx, rx) = tokio::sync::mpsc::channel(128);
    tokio::spawn(async move {
        let (deepseek_token, anthropic_token) = match extract_api_tokens(&headers) {
            Ok(tokens) => tokens,
            Err(e) => {
                let _ = tx.send(Err(e)).await;
                return;
            }
        };
        let deepseek_client = DeepSeekClient::new(deepseek_token);
        let anthropic_client = AnthropicClient::new(anthropic_token);
        let messages = request.get_messages_with_system();
        let mut deepseek_stream = match deepseek_client.chat_stream(messages.clone(), &request.deepseek_config).await {
            Ok(stream) => stream,
            Err(e) => {
                let _ = tx.send(Err(e)).await;
                return;
            }
        };
        let mut reasoning_content = String::new();
        let mut deepseek_usage = None;
        while let Some(chunk) = deepseek_stream.next().await {
            match chunk {
                Ok(chunk) => {
                    if let Some(reasoning) = chunk.choices[0].delta.reasoning_content.as_deref() {
                        reasoning_content.push_str(reasoning);
                    }
                    deepseek_usage = chunk.usage;
                }
                Err(e) => {
                    let _ = tx.send(Err(e)).await;
                    return;
                }
            }
        }
        let thinking_content = format!("<thinking>\n{}\n</thinking>", reasoning_content);
        if let Err(_) = tx.send(Ok(Event::default().data(thinking_content.clone()))).await {
            return;
        }
        let mut anthropic_messages = messages;
        anthropic_messages.push(Message {
            role: Role::Assistant,
            content: thinking_content,
        });
        let mut anthropic_stream = match anthropic_client.chat_stream(
            anthropic_messages,
            request.get_system_prompt().map(String::from),
            &request.anthropic_config
        ).await {
            Ok(stream) => stream,
            Err(e) => {
                let _ = tx.send(Err(e)).await;
                return;
            }
        };
        let mut anthropic_usage = None;
        let mut anthropic_model = None;
        while let Some(chunk) = anthropic_stream.next().await {
            match chunk {
                Ok(chunk) => {
                    if anthropic_model.is_none() {
                        anthropic_model = Some(chunk.model.clone());
                    }
                    let event: StreamEvent = chunk.into();
                    match Event::default().json_data(event) {
                        Ok(event) => {
                            if let Err(_) = tx.send(Ok(event)).await {
                                return;
                            }
                        }
                        Err(e) => {
                            let _ = tx.send(Err(ApiError::InternalError {
                                message: e.to_string()
                            })).await;
                            return;
                        }
                    }
                    anthropic_usage = chunk.usage;
                }
                Err(e) => {
                    let _ = tx.send(Err(e)).await;
                    return;
                }
            }
        }
        let (deepseek_cost, anthropic_cost) = match (deepseek_usage, anthropic_usage, anthropic_model) {
            (Some(ds_usage), Some(an_usage), Some(model)) => {
                let ds_cost = calculate_deepseek_cost(
                    ds_usage.prompt_tokens,
                    ds_usage.completion_tokens,
                    ds_usage.completion_tokens_details.reasoning_tokens,
                    ds_usage.prompt_tokens_details.cached_tokens,
                    &state.config,
                );
                let an_cost = calculate_anthropic_cost(
                    &model,
                    an_usage.input_tokens,
                    an_usage.output_tokens,
                    an_usage.cache_creation_input_tokens,
                    an_usage.cache_read_input_tokens,
                    &state.config,
                );
                (ds_cost, an_cost)
            }
            _ => {
                let _ = tx.send(Err(ApiError::InternalError {
                    message: "Missing usage data in stream".to_string()
                })).await;
                return;
            }
        };
        let combined_usage = CombinedUsage {
            total_cost: format_cost(deepseek_cost + anthropic_cost),
            deepseek_usage: DeepSeekUsage {
                input_tokens: deepseek_usage.unwrap().prompt_tokens,
                output_tokens: deepseek_usage.unwrap().completion_tokens,
                reasoning_tokens: deepseek_usage.unwrap().completion_tokens_details.reasoning_tokens,
                cached_input_tokens: deepseek_usage.unwrap().prompt_tokens_details.cached_tokens,
                total_tokens: deepseek_usage.unwrap().total_tokens,
                total_cost: format_cost(deepseek_cost),
            },
            anthropic_usage: AnthropicUsage {
                input_tokens: anthropic_usage.unwrap().input_tokens,
                output_tokens: anthropic_usage.unwrap().output_tokens,
                cached_write_tokens: anthropic_usage.unwrap().cache_creation_input_tokens,
                cached_read_tokens: anthropic_usage.unwrap().cache_read_input_tokens,
                total_tokens: anthropic_usage.unwrap().input_tokens + anthropic_usage.unwrap().output_tokens,
                total_cost: format_cost(anthropic_cost),
            },
        };
        if let Err(_) = tx.send(Ok(Event::default().json_data(StreamEvent::Usage(combined_usage)).unwrap())).await {
            return;
        }
    });
    Ok(ReceiverStream::new(rx))
}