use crate::bus::{
    ChatMessage, SYSTEM_CONTEXT_AGENT_PROMPT, SYSTEM_CONTEXT_HISTORY_COMPACTION,
    SYSTEM_CONTEXT_SCHEDULED_PROMPT,
};
use crate::config::LLMProviderConfig;
use crate::providers::{ChatCompletionRequest, LLMProvider, Message, create_provider};
use crate::text::{char_count, take_prefix_chars};

use crate::agent::{AgentError, AgentRuntimeConfig};

const TOKEN_ESTIMATE_SAFETY_MULTIPLIER: f64 = 1.2;
const CJK_CHARS_PER_TOKEN: f64 = 2.0;
const OTHER_CHARS_PER_TOKEN: f64 = 4.0;
const JSON_OVERHEAD_PER_MESSAGE: usize = 50;

/// Check if a character is CJK (Chinese, Japanese, Korean)
fn is_cjk_char(c: char) -> bool {
    matches!(c,
        '\u{4E00}'..='\u{9FFF}' |  // CJK Unified Ideographs
        '\u{3040}'..='\u{309F}' |  // Hiragana
        '\u{30A0}'..='\u{30FF}' |  // Katakana
        '\u{AC00}'..='\u{D7AF}' |  // Korean Hangul
        '\u{3400}'..='\u{4DBF}' |  // CJK Extension A
        '\u{20000}'..='\u{2A6DF}'  // CJK Extension B
    )
}

/// Token estimation using weighted character counting based on language
pub fn estimate_tokens(messages: &[ChatMessage]) -> usize {
    let mut cjk_count = 0usize;
    let mut other_count = 0usize;
    let mut media_refs_count = 0usize;

    for msg in messages {
        // Count content characters
        for c in msg.content.chars() {
            if is_cjk_char(c) {
                cjk_count += 1;
            } else {
                other_count += 1;
            }
        }
        // Count media references
        media_refs_count += msg.media_refs.len();
    }

    // Weighted token calculation: CJK chars need more tokens per character
    let content_tokens = (cjk_count as f64 / CJK_CHARS_PER_TOKEN)
        + (other_count as f64 / OTHER_CHARS_PER_TOKEN);

    // JSON serialization overhead for message structure (fields, brackets, etc.)
    let json_overhead = messages.len() * JSON_OVERHEAD_PER_MESSAGE;

    // Media references add to JSON size (each path is a string in the array)
    let media_overhead = media_refs_count * 20; // Each media ref adds ~20 chars to JSON

    // Apply safety multiplier
    ((content_tokens + json_overhead as f64 + media_overhead as f64)
        * TOKEN_ESTIMATE_SAFETY_MULTIPLIER) as usize
}

/// Configuration for context compression.
#[derive(Debug, Clone)]
pub struct ContextCompressionConfig {
    /// Preserve the latest N complete user turns in full.
    pub retain_last_user_turns: usize,
    /// Maximum characters in summary
    pub summary_max_chars: usize,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct UserTurnRange {
    start: usize,
    end_exclusive: usize,
}

#[derive(Debug, Clone)]
pub struct HistoryCompactionPlan {
    pub preserved_system_messages: Vec<ChatMessage>,
    pub summary_message: ChatMessage,
    pub preserved_messages: Vec<ChatMessage>,
    pub compressed_turns: usize,
    pub preserved_turns: usize,
}

impl Default for ContextCompressionConfig {
    fn default() -> Self {
        Self {
            retain_last_user_turns: 3,
            summary_max_chars: 20_000,
        }
    }
}

/// Context compressor that reduces message history when it exceeds token limits.
#[derive(Clone)]
pub struct ContextCompressor {
    config: ContextCompressionConfig,
    context_window: usize,
    /// Threshold ratio to trigger compression (50% of context window)
    threshold_ratio: f64,
}

impl ContextCompressor {
    #[cfg(test)]
    fn summary_char_budget_for_context_window(context_window: usize) -> usize {
        const SUMMARY_RATIO: f64 = 0.1;
        const CHARS_PER_TOKEN: f64 = 2.5;
        const MIN_SUMMARY_CHARS: usize = 1_500;
        const MAX_SUMMARY_CHARS: usize = 50_000;

        ((context_window as f64 * SUMMARY_RATIO * CHARS_PER_TOKEN) as usize)
            .clamp(MIN_SUMMARY_CHARS, MAX_SUMMARY_CHARS)
    }

    fn format_transcript_entry(message: &ChatMessage) -> String {
        let role = match message.role.as_str() {
            "assistant" => "Assistant",
            "tool" => "Tool",
            _ => message.role.as_str(),
        };
        let name = message
            .tool_name
            .as_ref()
            .map(|n| format!(" ({})", n))
            .unwrap_or_default();
        format!("{}: {}{}", role, message.content, name)
    }

    fn build_transcript(messages: &[ChatMessage]) -> String {
        messages
            .iter()
            .map(Self::format_transcript_entry)
            .collect::<Vec<_>>()
            .join("\n\n")
    }

    fn split_text_chunks(text: &str, max_chars: usize) -> Vec<String> {
        if text.is_empty() {
            return Vec::new();
        }

        let chunk_size = max_chars.max(1);
        let chars: Vec<char> = text.chars().collect();
        chars
            .chunks(chunk_size)
            .map(|chunk| chunk.iter().collect())
            .collect()
    }

    fn chunk_messages_for_summary(messages: &[ChatMessage], max_chars: usize) -> Vec<String> {
        if messages.is_empty() {
            return Vec::new();
        }

        let chunk_limit = max_chars.max(1);
        let mut chunks = Vec::new();
        let mut current = String::new();

        for entry in messages.iter().map(Self::format_transcript_entry) {
            let separator = if current.is_empty() { "" } else { "\n\n" };
            let candidate = format!("{}{}{}", current, separator, entry);

            if !current.is_empty() && char_count(&candidate) > chunk_limit {
                chunks.push(current);
                current = String::new();
            }

            if char_count(&entry) > chunk_limit {
                if !current.is_empty() {
                    chunks.push(current);
                    current = String::new();
                }
                chunks.extend(Self::split_text_chunks(&entry, chunk_limit));
                continue;
            }

            if current.is_empty() {
                current = entry;
            } else {
                current.push_str("\n\n");
                current.push_str(&entry);
            }
        }

        if !current.is_empty() {
            chunks.push(current);
        }

        chunks
    }

    fn build_summary_prompt(transcript: &str, target_chars: usize) -> String {
        format!(
            r#"You are a conversation compaction engine. Summarize the following conversation segment.

PRESERVE:
- Each user question or request in full or as a near-verbatim restatement
- All identifiers (UUIDs, hashes, file paths, URLs)
- Actions taken (tool calls, file operations, commands)
- Key information obtained (results, data, errors)
- Decisions and user preferences
- Current task status

OMIT:
- Reproducing full tool output verbatim unless it is essential
- Repeated greetings or filler

Do not assume tool content was pre-trimmed. You may receive long tool outputs; keep the important results, errors, and artifacts.

Be concise, aim for {} characters or less.

---

{}

"#,
            target_chars, transcript
        )
    }

    async fn summarize_transcript(
        &self,
        provider: &dyn LLMProvider,
        transcript: &str,
        target_chars: usize,
    ) -> Result<String, AgentError> {
        let request = ChatCompletionRequest {
            messages: vec![
                Message::system("You are a helpful assistant."),
                Message::user(Self::build_summary_prompt(transcript, target_chars)),
            ],
            temperature: Some(0.3),
            max_tokens: Some(1000),
            tools: None,
        };

        let response = provider
            .chat(request)
            .await
            .map_err(|e| AgentError::LlmError(e.to_string()))?;
        Ok(response.content)
    }

    async fn summarize_chunked_transcript(
        &self,
        provider: &dyn LLMProvider,
        messages: &[ChatMessage],
        transcript: &str,
    ) -> Result<String, AgentError> {
        let target_chars = self.config.summary_max_chars.max(1);
        let mut layer = Self::chunk_messages_for_summary(messages, target_chars);

        if layer.is_empty() {
            layer.push(transcript.to_string());
        }

        for _ in 0..6 {
            if layer.len() == 1 && char_count(&layer[0]) <= target_chars {
                return self
                    .summarize_transcript(provider, &layer[0], target_chars)
                    .await;
            }

            let per_chunk_target = (target_chars / layer.len().max(1))
                .max(500)
                .min(target_chars);
            let mut summaries = Vec::with_capacity(layer.len());
            for chunk in &layer {
                summaries.push(
                    self.summarize_transcript(provider, chunk, per_chunk_target)
                        .await?,
                );
            }

            if summaries.len() == 1 {
                let summary = summaries.pop().unwrap_or_default();
                if char_count(&summary) <= target_chars {
                    return Ok(summary);
                }
                layer = Self::split_text_chunks(&summary, target_chars);
                continue;
            }

            let merged = summaries.join("\n\n");
            if char_count(&merged) <= target_chars {
                return self
                    .summarize_transcript(provider, &merged, target_chars)
                    .await;
            }

            layer = Self::split_text_chunks(&merged, target_chars);
        }

        Ok(take_prefix_chars(transcript, target_chars))
    }

    /// Create a new compressor with the given context window size.
    pub fn new(context_window: usize) -> Self {
        Self {
            config: ContextCompressionConfig::default(),
            context_window,
            threshold_ratio: 0.5,
        }
    }

    pub fn from_provider_config(provider_config: &LLMProviderConfig) -> Self {
        Self::from_runtime_config(&AgentRuntimeConfig::from(provider_config.clone()))
    }

    pub fn from_runtime_config(config: &AgentRuntimeConfig) -> Self {
        Self::with_config(
            config.context_window_tokens,
            ContextCompressionConfig {
                summary_max_chars: config.context_summary_char_budget,
                ..ContextCompressionConfig::default()
            },
        )
    }

    /// Create with custom configuration.
    pub fn with_config(context_window: usize, config: ContextCompressionConfig) -> Self {
        Self {
            config,
            context_window,
            threshold_ratio: 0.5,
        }
    }

    /// Get the compression threshold in tokens.
    fn threshold(&self) -> usize {
        (self.context_window as f64 * self.threshold_ratio) as usize
    }

    pub fn should_compress(&self, history: &[ChatMessage]) -> bool {
        estimate_tokens(history) > self.threshold()
    }

    fn user_turn_ranges(&self, history: &[ChatMessage]) -> Vec<UserTurnRange> {
        let user_indices: Vec<usize> = history
            .iter()
            .enumerate()
            .filter(|(_, message)| message.role == "user")
            .map(|(index, _)| index)
            .collect();

        user_indices
            .iter()
            .enumerate()
            .map(|(turn_index, start)| UserTurnRange {
                start: *start,
                end_exclusive: user_indices
                    .get(turn_index + 1)
                    .copied()
                    .unwrap_or(history.len()),
            })
            .collect()
    }

    fn should_preserve_system_message(&self, message: &ChatMessage) -> bool {
        message.role == "system"
            && (message.has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT)
                || message.has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT))
    }

    fn split_prefix_messages(
        &self,
        history: &[ChatMessage],
    ) -> (Vec<ChatMessage>, Vec<ChatMessage>) {
        let preserved_system_messages = history
            .iter()
            .filter(|message| self.should_preserve_system_message(message))
            .cloned()
            .collect();
        let summary_source = history
            .iter()
            .filter(|message| !self.should_preserve_system_message(message))
            .cloned()
            .collect();

        (preserved_system_messages, summary_source)
    }

    pub async fn build_compaction_plan(
        &self,
        history: &[ChatMessage],
        provider_config: &LLMProviderConfig,
    ) -> Result<Option<HistoryCompactionPlan>, AgentError> {
        if !self.should_compress(history) {
            return Ok(None);
        }

        let turn_ranges = self.user_turn_ranges(history);
        if turn_ranges.len() <= self.config.retain_last_user_turns {
            return Ok(None);
        }

        let preserved_turn_start =
            turn_ranges[turn_ranges.len() - self.config.retain_last_user_turns].start;
        if preserved_turn_start == 0 {
            return Ok(None);
        }

        let (preserved_system_messages, summary_source) =
            self.split_prefix_messages(&history[..preserved_turn_start]);

        let summary = self
            .summarize_segment(&summary_source, provider_config)
            .await?;

        Ok(Some(HistoryCompactionPlan {
            preserved_system_messages,
            summary_message: ChatMessage::system_with_context(
                format!("[Compressed History]\n\n{}", summary),
                Some(SYSTEM_CONTEXT_HISTORY_COMPACTION.to_string()),
            ),
            preserved_messages: history[preserved_turn_start..].to_vec(),
            compressed_turns: turn_ranges.len() - self.config.retain_last_user_turns,
            preserved_turns: self.config.retain_last_user_turns,
        }))
    }

    /// Main entry point - compresses history if over threshold.
    pub async fn compress_if_needed(
        &self,
        history: Vec<ChatMessage>,
        provider_config: &LLMProviderConfig,
    ) -> Result<Vec<ChatMessage>, AgentError> {
        let tokens = estimate_tokens(&history);
        if tokens <= self.threshold() {
            #[cfg(debug_assertions)]
            tracing::info!(
                tokens = tokens,
                threshold = self.threshold(),
                msg_count = history.len(),
                "Context compression not needed"
            );
            return Ok(history);
        }

        tracing::info!(
            tokens = tokens,
            threshold = self.threshold(),
            msg_count = history.len(),
            "Starting context compression"
        );

        let current_history = match self
            .build_compaction_plan(&history, provider_config)
            .await?
        {
            Some(plan) => {
                let mut compressed = Vec::with_capacity(
                    plan.preserved_system_messages.len() + plan.preserved_messages.len() + 1,
                );
                compressed.extend(plan.preserved_system_messages);
                compressed.push(plan.summary_message);
                compressed.extend(plan.preserved_messages);
                compressed
            }
            None => history,
        };

        tracing::info!(
            final_tokens = estimate_tokens(&current_history),
            final_msg_count = current_history.len(),
            "Context compression completed"
        );

        Ok(current_history)
    }

    /// Summarize a segment of messages using LLM.
    async fn summarize_segment(
        &self,
        messages: &[ChatMessage],
        provider_config: &LLMProviderConfig,
    ) -> Result<String, AgentError> {
        if messages.is_empty() {
            return Ok(String::new());
        }

        let runtime_config = AgentRuntimeConfig::from(provider_config.clone());
        let provider = create_provider(runtime_config.provider)
            .map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
        let transcript = Self::build_transcript(messages);

        let result = if char_count(&transcript) <= self.config.summary_max_chars {
            self.summarize_transcript(
                provider.as_ref(),
                &transcript,
                self.config.summary_max_chars,
            )
            .await
        } else {
            self.summarize_chunked_transcript(provider.as_ref(), messages, &transcript)
                .await
        };

        match result {
            Ok(summary) => Ok(summary),
            Err(e) => {
                tracing::warn!(error = %e, "LLM summarization failed, using truncated transcript");
                Ok(take_prefix_chars(
                    &transcript,
                    self.config.summary_max_chars,
                ))
            }
        }
    }
}

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

    #[test]
    fn test_estimate_tokens() {
        let messages = vec![
            ChatMessage::user("Hello"),
            ChatMessage::assistant("Hi there!"),
            ChatMessage::user("How are you?"),
        ];

        let tokens = estimate_tokens(&messages);
        // Content: "Hello" (5) + "Hi there!" (9) + "How are you?" (12) = 26 chars
        // English: 26 / 4 = 6.5 tokens for content
        // JSON overhead: 3 * 50 = 150
        // Total before multiplier: 156.5
        // After 1.2x: ~188 tokens
        assert!(
            tokens > 100 && tokens < 300,
            "Expected ~150-250 tokens for English content, got {}",
            tokens
        );
    }

    #[test]
    fn test_estimate_tokens_chinese_content() {
        let messages = vec![
            ChatMessage::user("你好，这是一个中文测试"),
            ChatMessage::assistant("这是一个中文回复消息"),
        ];

        let tokens = estimate_tokens(&messages);
        // Content: ~20 CJK chars, CJK uses 2 chars/token = ~10 tokens for content
        // JSON overhead: 2 * 50 = 100
        // Total before multiplier: ~110
        // After 1.2x: ~132 tokens
        assert!(
            tokens > 80 && tokens < 200,
            "Expected ~100-180 tokens for Chinese content, got {}",
            tokens
        );
    }

    #[test]
    fn test_estimate_tokens_mixed_content() {
        let messages = vec![
            ChatMessage::user("Hello 世界 this is 测试"),
        ];

        let tokens = estimate_tokens(&messages);
        // Content: 18 English chars + 4 CJK chars
        // English: 18 / 4 = 4.5, CJK: 4 / 2 = 2, content tokens = 6.5
        // JSON overhead: 1 * 50 = 50
        // Total before multiplier: 56.5
        // After 1.2x: ~68 tokens
        assert!(
            tokens > 40 && tokens < 120,
            "Expected ~50-100 tokens for mixed content, got {}",
            tokens
        );
    }

    #[test]
    fn test_chinese_tokens_higher_than_english() {
        // Use more characters to make the content difference significant
        // compared to JSON overhead (50 tokens per message)
        let english = vec![ChatMessage::user(&"abcdefghij".repeat(20))]; // 200 English chars
        let chinese = vec![ChatMessage::user(&"这是一个测试消息字".repeat(20))]; // 200 CJK chars (10 chars * 20)

        let english_tokens = estimate_tokens(&english);
        let chinese_tokens = estimate_tokens(&chinese);

        // 200 English chars: 200/4 = 50 content tokens
        // 200 CJK chars: 200/2 = 100 content tokens
        // With same JSON overhead, Chinese should use ~1.5x tokens
        assert!(
            chinese_tokens > english_tokens * 130 / 100, // At least 1.3x
            "Chinese (200 chars) should use more tokens than English (200 chars): EN={} CN={}",
            english_tokens,
            chinese_tokens
        );
    }

    #[test]
    fn test_estimate_tokens_includes_image_media_refs() {
        let temp_dir = tempfile::tempdir().unwrap();
        let image_path = temp_dir.path().join("demo.jpg");
        std::fs::write(&image_path, vec![0_u8; 12_000]).unwrap();

        let plain = vec![ChatMessage::user("hello")];
        let with_image = vec![ChatMessage::user_with_media(
            "hello",
            vec![image_path.to_string_lossy().to_string()],
        )];

        assert!(estimate_tokens(&with_image) > estimate_tokens(&plain));
    }

    #[test]
    fn test_should_compress() {
        let compressor = ContextCompressor::new(20);
        // Need more content to trigger compression with new weighted calculation
        // 200 English chars / 4 = 50 tokens, plus overhead
        let messages = vec![ChatMessage::user(&"x".repeat(400))];
        assert!(compressor.should_compress(&messages));
    }

    #[test]
    fn test_user_turn_ranges_follow_user_boundaries() {
        let compressor = ContextCompressor::new(100_000);
        let history = vec![
            ChatMessage::system("system"),
            ChatMessage::user("u1"),
            ChatMessage::assistant("a1"),
            ChatMessage::tool("call-1", "bash", "t1"),
            ChatMessage::user("u2"),
            ChatMessage::assistant("a2"),
            ChatMessage::user("u3"),
        ];

        let turns = compressor.user_turn_ranges(&history);
        assert_eq!(
            turns,
            vec![
                UserTurnRange {
                    start: 1,
                    end_exclusive: 4
                },
                UserTurnRange {
                    start: 4,
                    end_exclusive: 6
                },
                UserTurnRange {
                    start: 6,
                    end_exclusive: 7
                },
            ]
        );
    }

    #[test]
    fn test_split_prefix_messages_preserves_key_system_messages() {
        let compressor = ContextCompressor::new(50);
        let prefix = vec![
            ChatMessage::system_with_context(
                "agent prompt",
                Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()),
            ),
            ChatMessage::user("u1"),
            ChatMessage::assistant("a1"),
            ChatMessage::system_with_context(
                "scheduled prompt",
                Some(SYSTEM_CONTEXT_SCHEDULED_PROMPT.to_string()),
            ),
        ];

        let (preserved_system_messages, summary_source) = compressor.split_prefix_messages(&prefix);
        assert_eq!(preserved_system_messages.len(), 2);
        assert_eq!(summary_source.len(), 2);
        assert!(preserved_system_messages[0].has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT));
        assert!(preserved_system_messages[1].has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT));
    }

    #[test]
    fn test_threshold() {
        let compressor = ContextCompressor::new(128_000);
        assert_eq!(compressor.threshold(), 64_000);
    }

    #[test]
    fn test_summary_char_budget_for_context_window_scales_and_clamps() {
        assert_eq!(
            ContextCompressor::summary_char_budget_for_context_window(4_096),
            1_500
        );
        assert_eq!(
            ContextCompressor::summary_char_budget_for_context_window(65_536),
            16_384
        );
        assert_eq!(
            ContextCompressor::summary_char_budget_for_context_window(128_000),
            32_000
        );
        assert_eq!(
            ContextCompressor::summary_char_budget_for_context_window(400_000),
            50_000
        );
    }

    #[test]
    fn test_chunk_messages_for_summary_keeps_message_boundaries_when_possible() {
        let messages = vec![
            ChatMessage::user("alpha"),
            ChatMessage::assistant("beta"),
            ChatMessage::user("gamma"),
        ];

        let chunks = ContextCompressor::chunk_messages_for_summary(&messages, 30);

        assert_eq!(chunks.len(), 2);
        assert!(chunks.iter().all(|chunk| char_count(chunk) <= 30));
        assert_eq!(chunks[0], "user: alpha\n\nAssistant: beta");
        assert_eq!(chunks[1], "user: gamma");
    }

    #[test]
    fn test_chunk_messages_for_summary_splits_oversized_message() {
        let messages = vec![ChatMessage::user(&"x".repeat(25))];

        let chunks = ContextCompressor::chunk_messages_for_summary(&messages, 10);

        assert!(chunks.len() > 1);
        assert!(chunks.iter().all(|chunk| char_count(chunk) <= 10));
        assert_eq!(chunks.concat(), "user: xxxxxxxxxxxxxxxxxxxxxxxxx");
    }
}
