use std::sync::Arc;

use crate::agent::system_prompt::build_runtime_context;
use crate::bus::ChatMessage;
use crate::memory::{MemoryCategory, MemoryManager};
use crate::work::WorkManager;

/// Immutable context used to assemble provider input for both the initial call
/// and context-overflow recovery.
pub(super) struct TurnRuntimeContext {
    system_prompt: String,
    runtime_context: String,
    memory_tokens: usize,
    active_plan_tokens: usize,
}

impl TurnRuntimeContext {
    pub(super) fn assemble(&self, mut history: Vec<ChatMessage>) -> Vec<ChatMessage> {
        history.insert(0, ChatMessage::system(self.system_prompt.clone()));
        if let Some(last_user) = history
            .iter_mut()
            .rev()
            .find(|message| message.role == "user")
        {
            append_runtime_context(last_user, &self.runtime_context);
        }
        history
    }

    pub(super) fn budget_hints(&self) -> (usize, usize) {
        (self.memory_tokens, self.active_plan_tokens)
    }
}

/// Builds runtime-only context outside the Session lock. Durable history
/// projection and compaction are orchestrated by SessionManager after these
/// variable-size sources are known, so the budget covers the complete request.
pub(super) async fn prepare_turn_runtime(
    memory_manager: Arc<MemoryManager>,
    work_manager: Arc<WorkManager>,
    session_id: &str,
    query: &str,
    system_prompt: String,
) -> TurnRuntimeContext {
    let memory_future = memory_manager.recall(query, 5, Some(MemoryCategory::Knowledge), None);
    let work_future = work_manager.active_plan(session_id);
    let (memory_result, work_result) = tokio::join!(memory_future, work_future);

    let memory_context = match memory_result {
        Ok(entries) if !entries.is_empty() => Some(
            entries
                .iter()
                .map(|entry| format!("- {}: {}", entry.key, entry.content))
                .collect::<Vec<_>>()
                .join("\n"),
        ),
        Err(error) => {
            tracing::warn!(error = %error, "Failed to fetch memory context");
            None
        }
        _ => None,
    };
    let work_context = match work_result {
        Ok(Some(plan)) => Some(plan.compact_context()),
        Ok(None) => None,
        Err(error) => {
            tracing::warn!(error = %error, "Failed to load active task plan");
            None
        }
    };
    TurnRuntimeContext {
        system_prompt,
        memory_tokens: memory_context
            .as_deref()
            .map(crate::agent::context_compaction::estimate_text_tokens)
            .unwrap_or_default(),
        active_plan_tokens: work_context
            .as_deref()
            .map(crate::agent::context_compaction::estimate_text_tokens)
            .unwrap_or_default(),
        runtime_context: build_runtime_context(
            Some(session_id),
            memory_context.as_deref(),
            work_context.as_deref(),
        ),
    }
}

fn append_runtime_context(message: &mut ChatMessage, runtime_context: &str) {
    if runtime_context.trim().is_empty() {
        return;
    }
    if message.content.trim().is_empty() {
        message.content = runtime_context.to_string();
    } else {
        message.content = format!("{}\n\n{}", message.content, runtime_context);
    }
}

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

    #[test]
    fn runtime_context_is_added_only_to_latest_user_message() {
        let runtime = TurnRuntimeContext {
            system_prompt: "system".to_string(),
            runtime_context: "runtime".to_string(),
            memory_tokens: 0,
            active_plan_tokens: 0,
        };
        let messages = runtime.assemble(vec![
            ChatMessage::user("old"),
            ChatMessage::assistant("answer"),
            ChatMessage::user("new"),
        ]);

        assert_eq!(messages[0].role, "system");
        assert_eq!(messages[1].content, "old");
        assert_eq!(messages[3].content, "new\n\nruntime");
    }

    #[test]
    fn overflow_reassembly_does_not_duplicate_runtime_context() {
        let runtime = TurnRuntimeContext {
            system_prompt: "system".to_string(),
            runtime_context: "runtime".to_string(),
            memory_tokens: 0,
            active_plan_tokens: 0,
        };

        let first = runtime.assemble(vec![ChatMessage::user("question")]);
        let recovered = runtime.assemble(vec![ChatMessage::user("question")]);

        assert_eq!(first.len(), recovered.len());
        assert_eq!(first[0].content, recovered[0].content);
        assert_eq!(first[1].content, recovered[1].content);
        assert_eq!(recovered.len(), 2);
    }
}
