browsing 0.1.5

//! Agent memory — persistent session context for smarter automation
//!
//! Three capabilities:
//! - Progress tracking: visited pages, checked items, task checkpoints
//! - Credential profiles: auto-fill known login info on recognized pages
//! - Site-specific strategy memory: domain-keyed learnings and quirks

use crate::error::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::Path;

/// A single credential profile for a domain
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CredentialProfile {
    /// Domain this profile applies to (e.g., "example.com")
    pub domain: String,
    /// Username / email
    pub username: String,
    /// Password (stored in memory only — no persistent storage for security)
    pub password: String,
    /// Optional: selector for the username field
    pub username_selector: Option<String>,
    /// Optional: selector for the password field
    pub password_selector: Option<String>,
    /// Optional: selector for the submit button
    pub submit_selector: Option<String>,
}

/// A learned strategy for a specific site
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SiteStrategy {
    /// Domain this strategy applies to
    pub domain: String,
    /// Human-readable description of the quirk/workaround
    pub description: String,
    /// Specific selectors or patterns that are known to work
    pub known_selectors: HashMap<String, String>,
    /// Steps to take before the main task on this site
    pub pre_task_actions: Vec<String>,
}

/// Tracks progress through a multi-step task
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ProgressTracker {
    /// Pages that have already been visited and processed
    pub visited_pages: Vec<String>,
    /// Items that have been checked/completed
    pub completed_items: Vec<String>,
    /// Current checkpoint in the task flow
    pub checkpoint: Option<String>,
    /// Number of search result pages already checked
    pub search_pages_checked: u32,
    /// Items found so far (for collection tasks)
    pub collected_items: Vec<String>,
}

impl ProgressTracker {
    /// Mark a page as visited
    pub fn mark_visited(&mut self, url: &str) {
        let normalized = url.to_string();
        if !self.visited_pages.contains(&normalized) {
            self.visited_pages.push(normalized);
        }
    }

    /// Mark an item as completed
    pub fn mark_completed(&mut self, item: &str) {
        let normalized = item.to_string();
        if !self.completed_items.contains(&normalized) {
            self.completed_items.push(normalized);
        }
    }

    /// Collect an item
    pub fn collect(&mut self, item: &str) {
        self.collected_items.push(item.to_string());
    }

    /// Set the current checkpoint
    pub fn set_checkpoint(&mut self, label: &str) {
        self.checkpoint = Some(label.to_string());
    }

    /// Check if a page has been visited
    pub fn is_visited(&self, url: &str) -> bool {
        self.visited_pages.contains(&url.to_string())
    }

    /// Get a progress summary
    pub fn to_prompt_context(&self) -> String {
        let mut lines = vec!["Task Progress:".to_string()];

        if let Some(ref cp) = self.checkpoint {
            lines.push(format!("  Current checkpoint: {}", cp));
        }

        if !self.visited_pages.is_empty() {
            lines.push(format!("  Pages visited: {}", self.visited_pages.len()));
        }

        if !self.completed_items.is_empty() {
            lines.push(format!(
                "  Items completed: {}",
                self.completed_items.len()
            ));
            for item in &self.completed_items {
                lines.push(format!("    - {}", item));
            }
        }

        if !self.collected_items.is_empty() {
            lines.push(format!(
                "  Items collected: {}",
                self.collected_items.len()
            ));
        }

        if self.search_pages_checked > 0 {
            lines.push(format!(
                "  Search pages checked: {}",
                self.search_pages_checked
            ));
        }

        lines.join("\n")
    }
}

/// In-memory agent memory (not persisted to disk for privacy)
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct AgentMemory {
    /// Credential profiles by domain
    pub credentials: HashMap<String, CredentialProfile>,
    /// Site-specific strategies by domain
    pub site_strategies: HashMap<String, SiteStrategy>,
    /// Progress tracker for the current task
    pub progress: ProgressTracker,
}

impl AgentMemory {
    /// Create empty memory
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a credential profile
    pub fn add_credential(&mut self, profile: CredentialProfile) {
        self.credentials.insert(profile.domain.clone(), profile);
    }

    /// Get credential for a domain
    pub fn get_credential(&self, domain: &str) -> Option<&CredentialProfile> {
        self.credentials.get(domain)
    }

    /// Add or update a site strategy
    pub fn add_strategy(&mut self, strategy: SiteStrategy) {
        self.site_strategies
            .insert(strategy.domain.clone(), strategy);
    }

    /// Get strategy for a domain
    pub fn get_strategy(&self, domain: &str) -> Option<&SiteStrategy> {
        self.site_strategies.get(domain)
    }

    /// Save memory to a JSON file (excludes credentials for security)
    pub fn save_to_file(&self, path: impl AsRef<Path>) -> Result<()> {
        let snapshot = MemorySnapshot {
            progress: self.progress.clone(),
            site_strategies: self.site_strategies.clone(),
            saved_at: chrono::Utc::now().to_rfc3339(),
        };
        let json = serde_json::to_string_pretty(&snapshot)?;
        fs::write(path, json)?;
        Ok(())
    }

    /// Load memory from a JSON file
    pub fn load_from_file(path: impl AsRef<Path>) -> Result<Self> {
        let contents = fs::read_to_string(path)?;
        let snapshot: MemorySnapshot = serde_json::from_str(&contents)?;
        Ok(Self {
            credentials: HashMap::new(), // Credentials are never persisted
            site_strategies: snapshot.site_strategies,
            progress: snapshot.progress,
        })
    }

    /// Build a prompt context string with relevant memory for the current domain
    pub fn build_context(&self, current_url: &str, page_intent: Option<&str>) -> String {
        let mut context = String::new();

        // Extract domain from URL
        let domain = extract_domain(current_url).unwrap_or_default();

        // Include site-specific strategy if available
        if let Some(strategy) = self.get_strategy(&domain) {
            context.push_str(&format!(
                "Site-specific knowledge for {}:\n{}",
                domain, strategy.description
            ));
            if !strategy.known_selectors.is_empty() {
                context.push_str("\nKnown working selectors:\n");
                for (name, selector) in &strategy.known_selectors {
                    context.push_str(&format!("  {}: {}\n", name, selector));
                }
            }
            context.push('\n');
        }

        // Include progress context
        let progress = self.progress.to_prompt_context();
        if !progress.is_empty() {
            context.push_str(&progress);
            context.push('\n');
        }

        // Include credential hint if this looks like a login page
        if page_intent == Some("login") {
            if let Some(cred) = self.get_credential(&domain) {
                context.push_str(&format!(
                    "Credential profile available for {} (user: {}). Use input action with known selectors.\n",
                    domain, cred.username
                ));
            }
        }

        context
    }
}

/// Serializable snapshot of memory (credentials excluded for security)
#[derive(Debug, Clone, Serialize, Deserialize)]
struct MemorySnapshot {
    progress: ProgressTracker,
    site_strategies: HashMap<String, SiteStrategy>,
    saved_at: String,
}

/// Extract domain from URL
fn extract_domain(url: &str) -> Option<String> {
    url.split("//").nth(1)?.split('/').next()?.split(':').next().map(|s| s.to_string())
}

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

    #[test]
    fn test_progress_tracker() {
        let mut pt = ProgressTracker::default();
        pt.mark_visited("https://example.com/page1");
        pt.mark_visited("https://example.com/page2");
        pt.mark_completed("Item A");
        pt.collect("Result 1");
        pt.set_checkpoint("Searching");

        assert!(pt.is_visited("https://example.com/page1"));
        assert!(!pt.is_visited("https://example.com/page3"));

        let ctx = pt.to_prompt_context();
        assert!(ctx.contains("Current checkpoint: Searching"));
        assert!(ctx.contains("Pages visited: 2"));
        assert!(ctx.contains("Items completed: 1"));
        assert!(ctx.contains("- Item A"));
    }

    #[test]
    fn test_agent_memory_credentials() {
        let mut memory = AgentMemory::new();
        memory.add_credential(CredentialProfile {
            domain: "example.com".to_string(),
            username: "user@example.com".to_string(),
            password: "secret".to_string(),
            username_selector: Some("#user".to_string()),
            password_selector: Some("#pass".to_string()),
            submit_selector: Some("#login".to_string()),
        });

        let cred = memory.get_credential("example.com").unwrap();
        assert_eq!(cred.username, "user@example.com");

        let ctx = memory.build_context("https://example.com/login", Some("login"));
        assert!(ctx.contains("Credential profile available"));
    }

    #[test]
    fn test_site_strategy() {
        let mut memory = AgentMemory::new();
        let mut selectors = HashMap::new();
        selectors.insert("search_button".to_string(), "button[type=submit]".to_string());

        memory.add_strategy(SiteStrategy {
            domain: "airline.com".to_string(),
            description: "Search button is mislabeled as 'Find'".to_string(),
            known_selectors: selectors,
            pre_task_actions: vec!["click('.cookie-accept')".to_string()],
        });

        let strategy = memory.get_strategy("airline.com").unwrap();
        assert_eq!(strategy.description, "Search button is mislabeled as 'Find'");

        let ctx = memory.build_context("https://airline.com/book", None);
        assert!(ctx.contains("mislabeled"));
        assert!(ctx.contains("search_button"));
    }

    #[test]
    fn test_extract_domain() {
        assert_eq!(
            extract_domain("https://example.com/path"),
            Some("example.com".to_string())
        );
        assert_eq!(
            extract_domain("http://api.example.com:8080/v1"),
            Some("api.example.com".to_string())
        );
    }

    #[test]
    fn test_memory_persistence() {
        let mut memory = AgentMemory::new();
        memory.progress.mark_visited("https://example.com/a");
        memory.progress.set_checkpoint("reviewing");

        let mut selectors = HashMap::new();
        selectors.insert("btn".to_string(), "#go".to_string());
        memory.add_strategy(SiteStrategy {
            domain: "example.com".to_string(),
            description: "Test".to_string(),
            known_selectors: selectors,
            pre_task_actions: vec![],
        });

        let temp = tempfile::NamedTempFile::new().unwrap();
        memory.save_to_file(temp.path()).unwrap();

        let loaded = AgentMemory::load_from_file(temp.path()).unwrap();
        assert!(loaded.progress.is_visited("https://example.com/a"));
        assert_eq!(loaded.progress.checkpoint, Some("reviewing".to_string()));
        assert!(loaded.get_strategy("example.com").is_some());
        // Credentials should never be persisted
        assert!(loaded.credentials.is_empty());
    }
}