agent-diva-agent 0.3.0

//! Context builder for assembling prompts

use crate::skills::SkillsLoader;
use agent_diva_core::soul::SoulStateStore;
use agent_diva_providers::Message;
use agent_diva_tools::sanitize::truncate_tool_result;
use std::path::Path;
use std::path::PathBuf;

const DEFAULT_AGENT_NAME: &str = "agent-diva";
const DEFAULT_AGENT_EMOJI: &str = "🐈";
const DEFAULT_AGENT_ROLE: &str = "helpful AI assistant";

/// Runtime controls for soul prompt injection.
#[derive(Debug, Clone)]
pub struct SoulContextSettings {
    pub enabled: bool,
    pub max_chars: usize,
    pub bootstrap_once: bool,
}

impl Default for SoulContextSettings {
    fn default() -> Self {
        Self {
            enabled: true,
            max_chars: 4000,
            bootstrap_once: true,
        }
    }
}

/// Builds the context for LLM requests
pub struct ContextBuilder {
    workspace: PathBuf,
    skills_loader: SkillsLoader,
    soul_settings: SoulContextSettings,
}

impl ContextBuilder {
    /// Create a new context builder
    pub fn new(workspace: PathBuf) -> Self {
        let skills_loader = SkillsLoader::new(&workspace, None);
        Self {
            workspace,
            skills_loader,
            soul_settings: SoulContextSettings::default(),
        }
    }

    /// Create a new context builder with skills
    pub fn with_skills(workspace: PathBuf, builtin_skills_dir: Option<PathBuf>) -> Self {
        let skills_loader = SkillsLoader::new(&workspace, builtin_skills_dir);
        Self {
            workspace,
            skills_loader,
            soul_settings: SoulContextSettings::default(),
        }
    }

    /// Override soul context settings.
    pub fn set_soul_settings(&mut self, settings: SoulContextSettings) {
        self.soul_settings = settings;
    }

    /// Build system prompt from workspace files and memory
    pub fn build_system_prompt(&self) -> String {
        let workspace_path = self.workspace.display();
        let now = chrono::Local::now().format("%Y-%m-%d %H:%M (%A)");
        let identity_header = self.load_identity_header();

        let mut prompt = format!(
            r#"{identity_header}

You have access to tools that allow you to:
- Read, write, and edit files
- Execute shell commands
- Search the web and fetch web pages
- Send messages to users on chat channels
- Schedule reminders and recurring jobs (cron)

## Current Time
{now}

## Workspace
Your workspace is at: {workspace_path}
- Memory files: {workspace_path}/memory/MEMORY.md
- Memory history log: {workspace_path}/memory/HISTORY.md"#
        );

        if self.soul_settings.enabled {
            self.append_soul_sections(&mut prompt);
        }

        // Skills - progressive loading
        // 1) Always-loaded skills (full content)
        let always_skills = self.skills_loader.get_always_skills();
        if !always_skills.is_empty() {
            let always_content = self.skills_loader.load_skills_for_context(&always_skills);
            if !always_content.is_empty() {
                prompt.push_str("\n\n## Active Skills\n");
                prompt.push_str(&always_content);
            }
        }

        // 2) Available skills summary
        let skills_summary = self.skills_loader.build_skills_summary();
        if !skills_summary.is_empty() {
            prompt.push_str("\n\n## Skills\n");
            prompt.push_str(
                "The following skills extend your capabilities. To use a skill, read its SKILL.md file using the read_file tool.\n",
            );
            prompt
                .push_str("Skills with available=\"false\" need dependencies installed first.\n\n");
            prompt.push_str(&skills_summary);
        }

        prompt.push_str(
            r#"

IMPORTANT: When responding to direct questions or conversations, reply directly with your text response.
Only use the 'message' tool when you need to send a message to a specific chat channel.
For normal conversation, just respond with text - do not call the message tool.
When a user asks to create a reminder, timer, or recurring schedule, use the 'cron' tool instead of saying the feature is unavailable.
When the user asks about prior work, project status, recent conclusions, commitments, or user preferences, prefer memory-backed answers over guesses.
If the system already injected recalled memory for the current turn, use that context first.
Use 'memory_recall' for compact answer-oriented recall, 'memory_search' for broader discovery, and 'memory_get' to fetch a full record or source fragment.
Use 'diary_list' and 'diary_read' only to expand diary dates or specific diary details.
Do not treat missing or unmatched memory as established fact.

Always be helpful, accurate, and concise. When using tools, explain what you're doing."#,
        );

        prompt.push_str(&format!(
            "\nWhen remembering something, write to {}/memory/MEMORY.md",
            workspace_path
        ));

        prompt
    }

    fn append_soul_sections(&self, prompt: &mut String) {
        let sections = [
            ("AGENTS.md", "Agent Rules"),
            ("SOUL.md", "Soul"),
            ("IDENTITY.md", "Identity"),
            ("USER.md", "User Profile"),
        ];

        for (rel, title) in sections {
            if let Some(content) = self.read_soul_file(rel) {
                self.append_section(prompt, title, &content);
            }
        }

        if self.should_include_bootstrap() {
            if let Some(content) = self.read_soul_file("BOOTSTRAP.md") {
                let _ = SoulStateStore::new(&self.workspace).mark_bootstrap_seeded();
                self.append_section(prompt, "Bootstrap", &content);
            }
        }
    }

    fn should_include_bootstrap(&self) -> bool {
        if !self.soul_settings.bootstrap_once {
            return true;
        }
        let store = SoulStateStore::new(&self.workspace);
        !store.is_bootstrap_completed()
    }

    fn read_soul_file(&self, rel: &str) -> Option<String> {
        let path = self.workspace.join(rel);
        read_trimmed_markdown(&path, self.soul_settings.max_chars)
    }

    fn append_section(&self, prompt: &mut String, title: &str, content: &str) {
        prompt.push_str("\n\n## ");
        prompt.push_str(title);
        prompt.push('\n');
        prompt.push_str(content);
    }

    fn load_identity_header(&self) -> String {
        let Some(content) = self.read_soul_file("IDENTITY.md") else {
            return default_identity_header();
        };

        let name = parse_identity_field(&content, &["name", "agent", "assistant"])
            .unwrap_or_else(|| DEFAULT_AGENT_NAME.to_string());
        let emoji = parse_identity_field(&content, &["emoji", "icon", "signature"])
            .unwrap_or_else(|| DEFAULT_AGENT_EMOJI.to_string());
        let role = parse_identity_field(&content, &["role", "nature", "type"])
            .unwrap_or_else(|| DEFAULT_AGENT_ROLE.to_string());
        let voice = parse_identity_field(&content, &["voice", "style", "vibe"]);

        let mut header = format!("# {} {}\n\nYou are {}, a {}.", name, emoji, name, role);
        if let Some(voice) = voice {
            header.push_str(" Preferred communication style: ");
            header.push_str(&voice);
            header.push('.');
        }
        header
    }

    /// Build the complete message list for an LLM call
    pub fn build_messages(
        &self,
        history: Vec<agent_diva_core::session::ChatMessage>,
        current_message: String,
        channel: Option<&str>,
        chat_id: Option<&str>,
    ) -> Vec<Message> {
        let mut messages = Vec::new();

        // System prompt
        let mut system_prompt = self.build_system_prompt();
        if let (Some(ch), Some(id)) = (channel, chat_id) {
            system_prompt.push_str(&format!(
                "\n\n## Current Session\nChannel: {}\nChat ID: {}",
                ch, id
            ));
        }
        messages.push(Message::system(system_prompt));

        // History - convert from ChatMessage to Message
        for msg in history {
            let message = match msg.role.as_str() {
                "user" => Message::user(&msg.content),
                "assistant" => {
                    let mut m = Message::assistant(&msg.content);
                    // Restore tool_calls from session history
                    if let Some(ref tc_values) = msg.tool_calls {
                        if let Ok(calls) =
                            serde_json::from_value::<Vec<agent_diva_providers::ToolCallRequest>>(
                                serde_json::Value::Array(tc_values.clone()),
                            )
                        {
                            m.tool_calls = Some(calls);
                        }
                    }
                    if let Some(reasoning) = msg.reasoning_content {
                        m.reasoning_content = Some(reasoning);
                    }
                    if let Some(thinking_blocks) = msg.thinking_blocks {
                        m.thinking_blocks = Some(thinking_blocks);
                    }
                    m
                }
                "tool" => {
                    let tool_call_id = msg.tool_call_id.unwrap_or_default();
                    let mut m = Message::tool(msg.content, tool_call_id);
                    m.name = msg.name;
                    m
                }
                _ => continue,
            };
            messages.push(message);
        }

        // Current message
        messages.push(Message::user(current_message));

        messages
    }

    /// Add a tool result to the message list
    ///
    /// Large tool results are truncated to prevent oversized API requests
    /// that could cause 400 errors from LLM providers.
    pub fn add_tool_result(
        &self,
        messages: &mut Vec<Message>,
        tool_call_id: String,
        _tool_name: String,
        result: String,
    ) {
        // Truncate large tool results to prevent API errors
        let truncated_result = truncate_tool_result(&result);
        messages.push(Message::tool(truncated_result, tool_call_id));
    }

    /// Add an assistant message with optional tool calls
    pub fn add_assistant_message(
        &self,
        messages: &mut Vec<Message>,
        content: Option<String>,
        tool_calls: Option<Vec<agent_diva_providers::ToolCallRequest>>,
        reasoning_content: Option<String>,
        thinking_blocks: Option<Vec<serde_json::Value>>,
    ) {
        let mut msg = Message::assistant(content.unwrap_or_default());
        if let Some(calls) = tool_calls {
            msg.tool_calls = Some(calls);
        }
        if let Some(reasoning) = reasoning_content {
            msg.reasoning_content = Some(reasoning);
        }
        if let Some(blocks) = thinking_blocks {
            msg.thinking_blocks = Some(blocks);
        }
        messages.push(msg);
    }
}

impl Default for ContextBuilder {
    fn default() -> Self {
        Self::new(PathBuf::from("."))
    }
}

fn read_trimmed_markdown(path: &Path, max_chars: usize) -> Option<String> {
    let content = std::fs::read_to_string(path).ok()?;
    let trimmed = content.trim();
    if trimmed.is_empty() {
        return None;
    }

    if trimmed.chars().count() <= max_chars {
        return Some(trimmed.to_string());
    }

    let mut out = String::new();
    for (idx, ch) in trimmed.chars().enumerate() {
        if idx >= max_chars.saturating_sub(3) {
            break;
        }
        out.push(ch);
    }
    out.push_str("...");
    Some(out)
}

fn parse_identity_field(content: &str, keys: &[&str]) -> Option<String> {
    for line in content.lines() {
        let line = line.trim().trim_start_matches(&['-', '*'][..]).trim();
        if line.is_empty() {
            continue;
        }

        // Split on ASCII or full-width colon to avoid manual byte indices.
        let (prefix, value_part) = match line.split_once(':').or_else(|| line.split_once('：')) {
            Some((p, v)) => (p.trim(), v.trim()),
            None => continue,
        };

        for key in keys {
            if prefix.eq_ignore_ascii_case(key) && !value_part.is_empty() {
                return Some(value_part.to_string());
            }
        }
    }
    None
}

fn default_identity_header() -> String {
    format!(
        "# {} {}\n\nYou are {}, a {}.",
        DEFAULT_AGENT_NAME, DEFAULT_AGENT_EMOJI, DEFAULT_AGENT_NAME, DEFAULT_AGENT_ROLE
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::TempDir;

    #[test]
    fn test_build_system_prompt() {
        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
        let prompt = builder.build_system_prompt();
        assert!(prompt.contains("agent-diva"));
        assert!(prompt.contains("/tmp/test"));
        assert!(prompt.contains("prefer memory-backed answers over guesses"));
        assert!(prompt.contains("'memory_search' for broader discovery"));
        assert!(!prompt.contains("## Long-term Memory"));
    }

    #[test]
    fn test_build_messages() {
        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
        let messages =
            builder.build_messages(vec![], "Hello".to_string(), Some("cli"), Some("test"));
        assert_eq!(messages.len(), 2); // system + user
        assert_eq!(messages[0].role, "system");
        assert_eq!(messages[1].role, "user");
        assert_eq!(messages[1].content, "Hello");
    }

    #[test]
    fn test_build_system_prompt_includes_skills_sections() {
        let workspace = TempDir::new().unwrap();
        let skills_dir = workspace.path().join("skills");
        fs::create_dir_all(skills_dir.join("always-skill")).unwrap();
        fs::write(
            skills_dir.join("always-skill").join("SKILL.md"),
            "---\nname: always-skill\ndescription: Always loaded\nmetadata: '{\"nanobot\":{\"always\":true}}'\n---\n\n# Always skill body\n",
        )
        .unwrap();

        let builder = ContextBuilder::with_skills(workspace.path().to_path_buf(), None);
        let prompt = builder.build_system_prompt();

        assert!(prompt.contains("## Active Skills"));
        assert!(prompt.contains("## Skills"));
        assert!(prompt.contains("<skills>"));
    }

    #[test]
    fn test_add_tool_result() {
        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
        let mut messages = vec![Message::user("test")];
        builder.add_tool_result(
            &mut messages,
            "call_123".to_string(),
            "read_file".to_string(),
            "file content".to_string(),
        );
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[1].role, "tool");
    }

    #[test]
    fn test_add_assistant_message() {
        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
        let mut messages = vec![Message::user("test")];

        // Test with reasoning content
        builder.add_assistant_message(
            &mut messages,
            Some("response".to_string()),
            None,
            Some("reasoning".to_string()),
            None,
        );

        assert_eq!(messages.len(), 2);
        assert_eq!(messages[1].role, "assistant");
        assert_eq!(messages[1].content, "response");
        assert_eq!(messages[1].reasoning_content, Some("reasoning".to_string()));
    }

    #[test]
    fn test_build_system_prompt_includes_soul_sections_in_order() {
        let workspace = TempDir::new().unwrap();
        fs::write(workspace.path().join("AGENTS.md"), "# Repo Rules").unwrap();
        fs::write(workspace.path().join("SOUL.md"), "# Core Traits").unwrap();
        fs::write(workspace.path().join("IDENTITY.md"), "# Identity").unwrap();
        fs::write(workspace.path().join("USER.md"), "# Preferences").unwrap();
        fs::write(workspace.path().join("BOOTSTRAP.md"), "# Bootstrap Steps").unwrap();

        let builder = ContextBuilder::new(workspace.path().to_path_buf());
        let prompt = builder.build_system_prompt();

        let idx_agents = prompt.find("## Agent Rules").unwrap();
        let idx_soul = prompt.find("## Soul").unwrap();
        let idx_identity = prompt.find("## Identity").unwrap();
        let idx_user = prompt.find("## User Profile").unwrap();
        let idx_bootstrap = prompt.find("## Bootstrap").unwrap();

        assert!(idx_agents < idx_soul);
        assert!(idx_soul < idx_identity);
        assert!(idx_identity < idx_user);
        assert!(idx_user < idx_bootstrap);
    }

    #[test]
    fn test_build_system_prompt_skips_bootstrap_when_completed() {
        let workspace = TempDir::new().unwrap();
        fs::write(workspace.path().join("BOOTSTRAP.md"), "# Bootstrap Steps").unwrap();
        let store = SoulStateStore::new(workspace.path());
        let state = agent_diva_core::soul::SoulState {
            bootstrap_completed_at: Some(chrono::Utc::now()),
            ..Default::default()
        };
        store.save(&state).unwrap();

        let builder = ContextBuilder::new(workspace.path().to_path_buf());
        let prompt = builder.build_system_prompt();
        assert!(!prompt.contains("## Bootstrap"));
    }

    #[test]
    fn test_read_trimmed_markdown_respects_char_limit() {
        let temp = TempDir::new().unwrap();
        let path = temp.path().join("SOUL.md");
        fs::write(&path, "abcdefghij").unwrap();

        let got = read_trimmed_markdown(&path, 6).unwrap();
        assert_eq!(got, "abc...");
        assert!(got.chars().count() <= 6);
    }

    #[test]
    fn test_build_system_prompt_uses_identity_file_for_header() {
        let workspace = TempDir::new().unwrap();
        fs::write(
            workspace.path().join("IDENTITY.md"),
            "# Identity\n- Name: Nova\n- Emoji: ✨\n- Role: strategic coding partner\n- Style: concise and direct\n",
        )
        .unwrap();
        let builder = ContextBuilder::new(workspace.path().to_path_buf());
        let prompt = builder.build_system_prompt();
        assert!(prompt.contains("# Nova ✨"));
        assert!(prompt.contains("You are Nova, a strategic coding partner."));
        assert!(prompt.contains("Preferred communication style: concise and direct."));
    }

    #[test]
    fn test_build_system_prompt_identity_header_falls_back_to_default() {
        let workspace = TempDir::new().unwrap();
        let builder = ContextBuilder::new(workspace.path().to_path_buf());
        let prompt = builder.build_system_prompt();
        assert!(prompt.contains("# agent-diva 🐈"));
        assert!(prompt.contains("You are agent-diva, a helpful AI assistant."));
    }

    #[test]
    fn test_build_system_prompt_empty_identity_falls_back_to_default() {
        let workspace = TempDir::new().unwrap();
        fs::write(workspace.path().join("IDENTITY.md"), "   \n").unwrap();
        let builder = ContextBuilder::new(workspace.path().to_path_buf());
        let prompt = builder.build_system_prompt();
        assert!(prompt.contains("# agent-diva 🐈"));
    }

    #[test]
    fn test_build_system_prompt_long_identity_is_trimmed_by_max_chars() {
        let workspace = TempDir::new().unwrap();
        let long_name = "N".repeat(6000);
        fs::write(
            workspace.path().join("IDENTITY.md"),
            format!("- Name: {}\n- Role: helper", long_name),
        )
        .unwrap();
        let mut builder = ContextBuilder::new(workspace.path().to_path_buf());
        builder.set_soul_settings(SoulContextSettings {
            enabled: true,
            max_chars: 120,
            bootstrap_once: true,
        });
        let prompt = builder.build_system_prompt();
        assert!(prompt.contains("You are "));
        assert!(prompt.chars().count() > 120);
    }

    #[test]
    fn test_parse_identity_field_handles_markdown_list() {
        let raw = "- Name: Diva\n- Style: pragmatic";
        assert_eq!(
            parse_identity_field(raw, &["name"]).as_deref(),
            Some("Diva")
        );
        assert_eq!(
            parse_identity_field(raw, &["style"]).as_deref(),
            Some("pragmatic")
        );
    }

    #[test]
    fn test_parse_identity_field_supports_chinese_voice_line() {
        let raw = "- Voice: 简洁、实用、协作";
        assert_eq!(
            parse_identity_field(raw, &["voice"]).as_deref(),
            Some("简洁、实用、协作")
        );
    }
}