use crate::message::{Message, MessageContent};
use crate::utils::safe_truncate;
use anyhow::{anyhow, Result};
use mcp_core::{Content, ResourceContents, Role};
use std::collections::HashSet;
use tracing::{debug, warn};

/// Maximum size for truncated content in characters
const MAX_TRUNCATED_CONTENT_SIZE: usize = 5000;

/// Handles messages that are individually larger than the context limit
/// by truncating their content rather than removing them entirely
fn handle_oversized_messages(
    messages: &[Message],
    token_counts: &[usize],
    context_limit: usize,
    strategy: &dyn TruncationStrategy,
) -> Result<(Vec<Message>, Vec<usize>), anyhow::Error> {
    let mut truncated_messages = Vec::new();
    let mut truncated_token_counts = Vec::new();
    let mut any_truncated = false;

    // Create a basic token counter for re-estimating truncated content
    // Note: This is a rough approximation since we don't have access to the actual tokenizer here
    let estimate_tokens = |text: &str| -> usize {
        // Rough approximation: 1 token per 4 characters for English text
        (text.len() / 4).max(1)
    };

    for (i, (message, &original_tokens)) in messages.iter().zip(token_counts.iter()).enumerate() {
        if original_tokens > context_limit {
            warn!(
                "Message {} has {} tokens, exceeding context limit of {}",
                i, original_tokens, context_limit
            );

            // Try to truncate the message content
            let truncated_message = truncate_message_content(message, MAX_TRUNCATED_CONTENT_SIZE)?;
            let estimated_new_tokens =
                estimate_message_tokens(&truncated_message, &estimate_tokens);

            if estimated_new_tokens > context_limit {
                // Even truncated message is too large, skip it entirely
                warn!("Skipping message {} as even truncated version ({} tokens) exceeds context limit", i, estimated_new_tokens);
                any_truncated = true;
                continue;
            }

            truncated_messages.push(truncated_message);
            truncated_token_counts.push(estimated_new_tokens);
            any_truncated = true;
        } else {
            truncated_messages.push(message.clone());
            truncated_token_counts.push(original_tokens);
        }
    }

    if any_truncated {
        debug!("Truncated large message content, now attempting normal truncation");
        // After content truncation, try normal truncation if still needed
        return truncate_messages(
            &truncated_messages,
            &truncated_token_counts,
            context_limit,
            strategy,
        );
    }

    Ok((truncated_messages, truncated_token_counts))
}

/// Truncates the content within a message while preserving its structure
fn truncate_message_content(message: &Message, max_content_size: usize) -> Result<Message> {
    let mut new_message = message.clone();

    for content in &mut new_message.content {
        match content {
            MessageContent::Text(text_content) => {
                if text_content.text.chars().count() > max_content_size {
                    let truncated = format!(
                        "{}\n\n[... content truncated from {} to {} characters ...]",
                        safe_truncate(&text_content.text, max_content_size),
                        text_content.text.chars().count(),
                        max_content_size
                    );
                    text_content.text = truncated;
                }
            }
            MessageContent::ToolResponse(tool_response) => {
                if let Ok(ref mut result) = tool_response.tool_result {
                    for content_item in result {
                        if let Content::Text(ref mut text_content) = content_item {
                            if text_content.text.chars().count() > max_content_size {
                                let truncated = format!(
                                    "{}\n\n[... tool response truncated from {} to {} characters ...]",
                                    safe_truncate(&text_content.text, max_content_size),
                                    text_content.text.chars().count(),
                                    max_content_size
                                );
                                text_content.text = truncated;
                            }
                        }
                        // Handle Resource content which might contain large text
                        else if let Content::Resource(ref mut resource_content) = content_item {
                            if let ResourceContents::TextResourceContents { text, .. } =
                                &mut resource_content.resource
                            {
                                if text.chars().count() > max_content_size {
                                    let truncated = format!(
                                        "{}\n\n[... resource content truncated from {} to {} characters ...]",
                                        safe_truncate(text, max_content_size),
                                        text.chars().count(),
                                        max_content_size
                                    );
                                    *text = truncated;
                                }
                            }
                        }
                    }
                }
            }
            // Other content types are typically smaller, but we could extend this if needed
            _ => {}
        }
    }

    Ok(new_message)
}

/// Estimates token count for a message using a simple heuristic
fn estimate_message_tokens(message: &Message, estimate_fn: &dyn Fn(&str) -> usize) -> usize {
    let mut total_tokens = 10; // Base overhead for message structure

    for content in &message.content {
        match content {
            MessageContent::Text(text_content) => {
                total_tokens += estimate_fn(&text_content.text);
            }
            MessageContent::ToolResponse(tool_response) => {
                if let Ok(ref result) = tool_response.tool_result {
                    for content_item in result {
                        match content_item {
                            Content::Text(text_content) => {
                                total_tokens += estimate_fn(&text_content.text);
                            }
                            Content::Resource(resource_content) => {
                                match &resource_content.resource {
                                    ResourceContents::TextResourceContents { text, .. } => {
                                        total_tokens += estimate_fn(text);
                                    }
                                    _ => total_tokens += 5, // Small overhead for other resource types
                                }
                            }
                            _ => total_tokens += 5, // Small overhead for other content types
                        }
                    }
                }
            }
            _ => total_tokens += 5, // Small overhead for other content types
        }
    }

    total_tokens
}

/// Truncates the messages to fit within the model's context window.
/// Mutates the input messages and token counts in place.
/// Returns an error if it's impossible to truncate the messages within the context limit.
/// - messages: The vector of messages in the conversation.
/// - token_counts: A parallel vector containing the token count for each message.
/// - context_limit: The maximum allowed context length in tokens.
/// - strategy: The truncation strategy to use. Only option is OldestFirstTruncation.
pub fn truncate_messages(
    messages: &[Message],
    token_counts: &[usize],
    context_limit: usize,
    strategy: &dyn TruncationStrategy,
) -> Result<(Vec<Message>, Vec<usize>), anyhow::Error> {
    let mut messages = messages.to_owned();
    let mut token_counts = token_counts.to_owned();

    if messages.len() != token_counts.len() {
        return Err(anyhow!(
            "The vector for messages and token_counts must have same length"
        ));
    }

    // Step 1: Calculate total tokens
    let mut total_tokens: usize = token_counts.iter().sum();
    debug!("Total tokens before truncation: {}", total_tokens);

    // Check if any individual message is larger than the context limit
    // First, check for any message that's too large
    let max_message_tokens = token_counts.iter().max().copied().unwrap_or(0);
    if max_message_tokens > context_limit {
        // Try to handle large messages by truncating their content
        debug!(
            "Found oversized message with {} tokens, attempting content truncation",
            max_message_tokens
        );
        return handle_oversized_messages(&messages, &token_counts, context_limit, strategy);
    }

    let min_user_msg_tokens = messages
        .iter()
        .zip(token_counts.iter())
        .filter(|(msg, _)| msg.role == Role::User && msg.has_only_text_content())
        .map(|(_, &tokens)| tokens)
        .min();

    // If there are no valid user messages, or the smallest one is too big for the context
    if min_user_msg_tokens.is_none() || min_user_msg_tokens.unwrap() > context_limit {
        return Err(anyhow!(
            "Not possible to truncate messages within context limit: no suitable user messages found"
        ));
    }

    if total_tokens <= context_limit {
        return Ok((messages, token_counts)); // No truncation needed
    }

    // Step 2: Determine indices to remove based on strategy
    let indices_to_remove =
        strategy.determine_indices_to_remove(&messages, &token_counts, context_limit)?;

    // Circuit breaker: if we can't remove enough messages, fail gracefully
    let tokens_to_remove: usize = indices_to_remove
        .iter()
        .map(|&i| token_counts.get(i).copied().unwrap_or(0))
        .sum();

    if total_tokens - tokens_to_remove > context_limit && !indices_to_remove.is_empty() {
        debug!(
            "Standard truncation insufficient: {} tokens remain after removing {} tokens",
            total_tokens - tokens_to_remove,
            tokens_to_remove
        );
        // Try more aggressive truncation or content truncation
        return handle_oversized_messages(&messages, &token_counts, context_limit, strategy);
    }

    if indices_to_remove.is_empty() && total_tokens > context_limit {
        return Err(anyhow!(
            "Cannot truncate any messages: all messages may be essential or too large individually"
        ));
    }

    // Step 3: Remove the marked messages
    // Vectorize the set and sort in reverse order to avoid shifting indices when removing
    let mut indices_to_remove = indices_to_remove.iter().cloned().collect::<Vec<usize>>();
    indices_to_remove.sort_unstable_by(|a, b| b.cmp(a));

    for &index in &indices_to_remove {
        if index < messages.len() {
            let _ = messages.remove(index);
            let removed_tokens = token_counts.remove(index);
            total_tokens -= removed_tokens;
        }
    }

    // Step 4: Ensure the last message is a user message with TextContent only
    while let Some(last_msg) = messages.last() {
        if last_msg.role != Role::User || !last_msg.has_only_text_content() {
            let _ = messages.pop().ok_or(anyhow!("Failed to pop message"))?;
            let removed_tokens = token_counts
                .pop()
                .ok_or(anyhow!("Failed to pop token count"))?;
            total_tokens -= removed_tokens;
        } else {
            break;
        }
    }

    // Step 5: Check first msg is a User message with TextContent only
    while let Some(first_msg) = messages.first() {
        if first_msg.role != Role::User || !first_msg.has_only_text_content() {
            let _ = messages.remove(0);
            let removed_tokens = token_counts.remove(0);
            total_tokens -= removed_tokens;
        } else {
            break;
        }
    }

    debug!("Total tokens after truncation: {}", total_tokens);

    // Ensure we have at least one message remaining and it's within context limit
    if messages.is_empty() {
        return Err(anyhow!(
            "Unable to preserve any messages within context limit"
        ));
    }

    if total_tokens > context_limit {
        return Err(anyhow!(
            "Unable to truncate messages within context window."
        ));
    }

    debug!("Truncation complete. Total tokens: {}", total_tokens);
    Ok((messages, token_counts))
}

/// Trait representing a truncation strategy
pub trait TruncationStrategy {
    /// Determines the indices of messages to remove to fit within the context limit.
    ///
    /// - `messages`: The list of messages in the conversation.
    /// - `token_counts`: A parallel array containing the token count for each message.
    /// - `context_limit`: The maximum allowed context length in tokens.
    ///
    /// Returns a vector of indices to remove.
    fn determine_indices_to_remove(
        &self,
        messages: &[Message],
        token_counts: &[usize],
        context_limit: usize,
    ) -> Result<HashSet<usize>>;
}

/// Strategy to truncate messages by removing the oldest first
pub struct OldestFirstTruncation;

impl TruncationStrategy for OldestFirstTruncation {
    fn determine_indices_to_remove(
        &self,
        messages: &[Message],
        token_counts: &[usize],
        context_limit: usize,
    ) -> Result<HashSet<usize>> {
        let mut indices_to_remove = HashSet::new();
        let mut total_tokens: usize = token_counts.iter().sum();
        let mut tool_ids_to_remove = HashSet::new();

        for (i, message) in messages.iter().enumerate() {
            if total_tokens <= context_limit {
                break;
            }

            // Remove the message
            indices_to_remove.insert(i);
            total_tokens -= token_counts[i];
            debug!(
                "OldestFirst: Removing message at index {}. Tokens removed: {}",
                i, token_counts[i]
            );

            // If it's a ToolRequest or ToolResponse, mark its pair for removal
            if message.is_tool_call() || message.is_tool_response() {
                message.get_tool_ids().iter().for_each(|id| {
                    tool_ids_to_remove.insert((i, id.to_string()));
                });
            }
        }

        // Now, find and remove paired ToolResponses or ToolRequests
        for (i, message) in messages.iter().enumerate() {
            let message_tool_ids = message.get_tool_ids();
            // Find the other part of the pair - same tool_id but different message index
            for (message_idx, tool_id) in &tool_ids_to_remove {
                if message_idx != &i && message_tool_ids.contains(tool_id.as_str()) {
                    indices_to_remove.insert(i);
                    // No need to check other tool_ids for this message since it's already marked
                    break;
                }
            }
        }

        Ok(indices_to_remove)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::message::Message;
    use anyhow::Result;
    use mcp_core::content::Content;
    use mcp_core::tool::ToolCall;
    use serde_json::json;

    // Helper function to create a user text message with a specified token count
    fn user_text(index: usize, tokens: usize) -> (Message, usize) {
        let content = format!("User message {}", index);
        (Message::user().with_text(content), tokens)
    }

    // Helper function to create an assistant text message with a specified token count
    fn assistant_text(index: usize, tokens: usize) -> (Message, usize) {
        let content = format!("Assistant message {}", index);
        (Message::assistant().with_text(content), tokens)
    }

    // Helper function to create a tool request message with a specified token count
    fn assistant_tool_request(id: &str, tool_call: ToolCall, tokens: usize) -> (Message, usize) {
        (
            Message::assistant().with_tool_request(id, Ok(tool_call)),
            tokens,
        )
    }

    // Helper function to create a tool response message with a specified token count
    fn user_tool_response(id: &str, result: Vec<Content>, tokens: usize) -> (Message, usize) {
        (Message::user().with_tool_response(id, Ok(result)), tokens)
    }

    // Helper function to create a large tool response with massive content
    fn large_tool_response(id: &str, large_text: String, tokens: usize) -> (Message, usize) {
        (
            Message::user().with_tool_response(id, Ok(vec![Content::text(large_text)])),
            tokens,
        )
    }

    // Helper function to create messages with alternating user and assistant
    // text messages of a fixed token count
    fn create_messages_with_counts(
        num_pairs: usize,
        tokens: usize,
        remove_last: bool,
    ) -> (Vec<Message>, Vec<usize>) {
        let mut messages: Vec<Message> = (0..num_pairs)
            .flat_map(|i| {
                vec![
                    user_text(i * 2, tokens).0,
                    assistant_text((i * 2) + 1, tokens).0,
                ]
            })
            .collect();

        if remove_last {
            messages.pop();
        }

        let token_counts = vec![tokens; messages.len()];

        (messages, token_counts)
    }

    #[test]
    fn test_handle_oversized_single_message() -> Result<()> {
        // Create a scenario similar to the real issue: one very large tool response
        let large_content = "A".repeat(50000); // Very large content
        let messages = vec![
            user_text(1, 10).0,
            assistant_tool_request(
                "tool1",
                ToolCall::new("read_file", json!({"path": "large_file.txt"})),
                20,
            )
            .0,
            large_tool_response("tool1", large_content, 100000).0, // Massive tool response
            user_text(2, 10).0,
        ];
        let token_counts = vec![10, 20, 100000, 10]; // One message is huge
        let context_limit = 5000; // Much smaller than the large message

        let result = truncate_messages(
            &messages,
            &token_counts,
            context_limit,
            &OldestFirstTruncation,
        );

        // Should succeed by truncating the large content
        assert!(
            result.is_ok(),
            "Should handle oversized message by content truncation"
        );
        let (truncated_messages, truncated_counts) = result.unwrap();

        // Should have some messages remaining
        assert!(
            !truncated_messages.is_empty(),
            "Should have some messages left"
        );

        // Total should be within limit
        let total_tokens: usize = truncated_counts.iter().sum();
        assert!(
            total_tokens <= context_limit,
            "Total tokens {} should be <= context limit {}",
            total_tokens,
            context_limit
        );

        Ok(())
    }

    #[test]
    fn test_oldest_first_no_truncation() -> Result<()> {
        let (messages, token_counts) = create_messages_with_counts(1, 10, false);
        let context_limit = 25;

        let result = truncate_messages(
            &messages,
            &token_counts,
            context_limit,
            &OldestFirstTruncation,
        )?;

        assert_eq!(result.0, messages);
        assert_eq!(result.1, token_counts);
        Ok(())
    }

    #[test]
    fn test_complex_conversation_with_tools() -> Result<()> {
        // Simulating a real conversation with multiple tool interactions
        let tool_call1 = ToolCall::new("file_read", json!({"path": "/tmp/test.txt"}));
        let tool_call2 = ToolCall::new("database_query", json!({"query": "SELECT * FROM users"}));

        let messages = vec![
            user_text(1, 15).0, // Initial user query
            assistant_tool_request("tool1", tool_call1.clone(), 20).0,
            user_tool_response(
                "tool1",
                vec![Content::text("File contents".to_string())],
                10,
            )
            .0,
            assistant_text(2, 25).0, // Assistant processes file contents
            user_text(3, 10).0,      // User follow-up
            assistant_tool_request("tool2", tool_call2.clone(), 30).0,
            user_tool_response(
                "tool2",
                vec![Content::text("Query results".to_string())],
                20,
            )
            .0,
            assistant_text(4, 35).0, // Assistant analyzes query results
            user_text(5, 5).0,       // Final user confirmation
        ];

        let token_counts = vec![15, 20, 10, 25, 10, 30, 20, 35, 5];
        let context_limit = 100; // Force truncation while preserving some tool interactions

        let result = truncate_messages(
            &messages,
            &token_counts,
            context_limit,
            &OldestFirstTruncation,
        )?;
        let (truncated_messages, truncated_counts) = result;

        // Verify that tool pairs are kept together and the conversation remains coherent
        assert!(truncated_messages.len() >= 3); // At least one complete interaction should remain
        assert!(truncated_messages.last().unwrap().role == Role::User); // Last message should be from user

        // Verify tool pairs are either both present or both removed
        let tool_ids: HashSet<_> = truncated_messages
            .iter()
            .flat_map(|m| m.get_tool_ids())
            .collect();

        // Each tool ID should appear 0 or 2 times (request + response)
        for id in tool_ids {
            let count = truncated_messages
                .iter()
                .flat_map(|m| m.get_tool_ids().into_iter())
                .filter(|&tool_id| tool_id == id)
                .count();
            assert!(count == 0 || count == 2, "Tool pair was split: {}", id);
        }

        // Total should be within limit
        let total_tokens: usize = truncated_counts.iter().sum();
        assert!(total_tokens <= context_limit);

        Ok(())
    }

    #[test]
    fn test_edge_case_context_window() -> Result<()> {
        // Test case where we're exactly at the context limit
        let (messages, token_counts) = create_messages_with_counts(2, 25, false);
        let context_limit = 100; // Exactly matches total tokens

        let result = truncate_messages(
            &messages,
            &token_counts,
            context_limit,
            &OldestFirstTruncation,
        )?;
        let (mut messages, mut token_counts) = result;

        assert_eq!(messages.len(), 4); // No truncation needed
        assert_eq!(token_counts.iter().sum::<usize>(), 100);

        // Now add one more token to force truncation
        messages.push(user_text(5, 1).0);
        token_counts.push(1);

        let result = truncate_messages(
            &messages,
            &token_counts,
            context_limit,
            &OldestFirstTruncation,
        )?;
        let (messages, token_counts) = result;

        assert!(token_counts.iter().sum::<usize>() <= context_limit);
        assert!(messages.last().unwrap().role == Role::User);

        Ok(())
    }

    #[test]
    fn test_multi_tool_chain() -> Result<()> {
        // Simulate a chain of dependent tool calls
        let tool_calls = vec![
            ToolCall::new("git_status", json!({})),
            ToolCall::new("git_diff", json!({"file": "main.rs"})),
            ToolCall::new("git_commit", json!({"message": "Update"})),
        ];

        let mut messages = Vec::new();
        let mut token_counts = Vec::new();

        // Build a chain of related tool calls
        // 30 tokens each round
        for (i, tool_call) in tool_calls.into_iter().enumerate() {
            let id = format!("git_{}", i);
            messages.push(user_text(i, 10).0);
            token_counts.push(10);

            messages.push(assistant_tool_request(&id, tool_call, 15).0);
            token_counts.push(20);
        }

        let context_limit = 50; // Force partial truncation

        let result = truncate_messages(
            &messages,
            &token_counts,
            context_limit,
            &OldestFirstTruncation,
        )?;
        let (truncated_messages, _) = result;

        // Verify that remaining tool chains are complete
        let remaining_tool_ids: HashSet<_> = truncated_messages
            .iter()
            .flat_map(|m| m.get_tool_ids())
            .collect();

        for _id in remaining_tool_ids {
            // Count request/response pairs
            let requests = truncated_messages
                .iter()
                .flat_map(|m| m.get_tool_request_ids().into_iter())
                .count();

            let responses = truncated_messages
                .iter()
                .flat_map(|m| m.get_tool_response_ids().into_iter())
                .count();

            assert_eq!(requests, 1, "Each remaining tool should have one request");
            assert_eq!(responses, 1, "Each remaining tool should have one response");
        }

        Ok(())
    }

    #[test]
    fn test_truncation_with_image_content() -> Result<()> {
        // Create a conversation with image content mixed in
        let messages = vec![
            Message::user().with_image("base64_data", "image/png"), // 50 tokens
            Message::assistant().with_text("I see the image"),      // 10 tokens
            Message::user().with_text("Can you describe it?"),      // 10 tokens
            Message::assistant().with_text("It shows..."),          // 20 tokens
            Message::user().with_text("Thanks!"),                   // 5 tokens
        ];
        let token_counts = vec![50, 10, 10, 20, 5];
        let context_limit = 45; // Force truncation

        let result = truncate_messages(
            &messages,
            &token_counts,
            context_limit,
            &OldestFirstTruncation,
        )?;
        let (messages, token_counts) = result;

        // Verify the conversation still makes sense
        assert!(messages.len() >= 1);
        assert!(messages.last().unwrap().role == Role::User);
        assert!(token_counts.iter().sum::<usize>() <= context_limit);

        Ok(())
    }

    #[test]
    fn test_error_cases() -> Result<()> {
        // Test impossibly small context window
        let (messages, token_counts) = create_messages_with_counts(1, 10, false);
        let result = truncate_messages(
            &messages,
            &token_counts,
            5, // Impossibly small context
            &OldestFirstTruncation,
        );
        assert!(result.is_err());

        // Test unmatched token counts
        let messages = vec![user_text(1, 10).0];
        let token_counts = vec![10, 10]; // Mismatched length
        let result = truncate_messages(&messages, &token_counts, 100, &OldestFirstTruncation);
        assert!(result.is_err());

        Ok(())
    }
}