snm-brightdata-client 0.4.0

// src/tools/scrape.rs - ENHANCED VERSION WITH REDIS CACHE SUPPORT
use crate::tool::{Tool, ToolResult, McpContent};
use crate::error::BrightDataError;
use crate::extras::logger::JSON_LOGGER;
use crate::filters::{ResponseFilter, ResponseStrategy};
use crate::services::cache::scrape_cache::get_scrape_cache;
use async_trait::async_trait;
use reqwest::Client;
use serde_json::{json, Value};
use std::env;
use std::time::Duration;
use std::collections::HashMap;
use log::{info, warn, error};

pub struct Scraper;

#[async_trait]
impl Tool for Scraper {
    fn name(&self) -> &str {
        "scrape_website"
    }

    fn description(&self) -> &str {
        "Scrape a webpage using BrightData with intelligent caching and priority-based processing. Supports Web Unlocker with Redis cache for improved performance."
    }

    fn input_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "url": {
                    "type": "string",
                    "description": "The URL to scrape"
                },
                "session_id": {
                    "type": "string",
                    "description": "Session ID for caching and conversation context tracking"
                },
                "data_type": {
                    "type": "string",
                    "enum": ["auto", "article", "product", "news", "contact", "general"],
                    "default": "auto",
                    "description": "Type of content to focus on during extraction"
                },
                "extraction_format": {
                    "type": "string",
                    "enum": ["structured", "markdown", "text", "json"],
                    "default": "structured",
                    "description": "Format for extracted content"
                },
                "clean_content": {
                    "type": "boolean",
                    "default": true,
                    "description": "Remove noise and focus on main content"
                },
                "schema": {
                    "type": "object",
                    "description": "Optional extraction schema for structured data"
                },
                "force_refresh": {
                    "type": "boolean",
                    "default": false,
                    "description": "Force fresh scraping, bypassing cache"
                }
            },
            "required": ["url"]
        })
    }

    async fn execute(&self, parameters: Value) -> Result<ToolResult, BrightDataError> {
        self.execute_internal(parameters).await
    }

    async fn execute_internal(&self, parameters: Value) -> Result<ToolResult, BrightDataError> {
        let url = parameters
            .get("url")
            .and_then(|v| v.as_str())
            .ok_or_else(|| BrightDataError::ToolError("Missing 'url' parameter".into()))?;

        let session_id = parameters
            .get("user_id")
            .and_then(|v| v.as_str())
            .ok_or_else(|| BrightDataError::ToolError("Missing 'user_id' parameter".into()))?;

        let data_type = parameters
            .get("data_type")
            .and_then(|v| v.as_str())
            .unwrap_or("auto");

        let extraction_format = parameters
            .get("extraction_format")
            .and_then(|v| v.as_str())
            .unwrap_or("structured");

        let clean_content = parameters
            .get("clean_content")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);

        let force_refresh = parameters
            .get("force_refresh")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        let schema = parameters.get("schema").cloned();

        let execution_id = self.generate_execution_id();
        
        info!("🌐 Scraping request: '{}' (session: {}, type: {}, format: {})", 
              url, session_id, data_type, extraction_format);
        
        // 🎯 CACHE CHECK - Check Redis cache first (unless force_refresh=true)
        if !force_refresh {
            match self.check_cache_first(url, session_id).await {
                Ok(Some(cached_result)) => {
                    info!("🚀 Cache HIT: Returning cached data for {} in session {}", url, session_id);
                    
                    // Create tool result from cached data
                    let content = cached_result.get("content").and_then(|c| c.as_str()).unwrap_or("");
                    let source_used = "Cache";
                    let method_used = "Redis Cache";
                    
                    let formatted_response = self.create_formatted_scrape_response(
                        url, data_type, extraction_format, content, &execution_id
                    );
                    
                    let tool_result = ToolResult::success_with_raw(
                        vec![McpContent::text(formatted_response)], 
                        cached_result
                    );
                    
                    // Apply filtering only if DEDUCT_DATA=true
                    if self.is_data_reduction_enabled() {
                        return Ok(ResponseStrategy::apply_size_limits(tool_result));
                    } else {
                        return Ok(tool_result);
                    }
                }
                Ok(None) => {
                    info!("💾 Cache MISS: Fetching fresh data for {} in session {}", url, session_id);
                }
                Err(e) => {
                    warn!("🚨 Cache error (continuing with fresh fetch): {}", e);
                }
            }
        } else {
            info!("🔄 Force refresh requested, bypassing cache for {}", url);
        }

        // 🌐 FRESH FETCH - Cache miss or force refresh, fetch from BrightData
        match self.scrape_with_brightdata(url, data_type, extraction_format, clean_content, schema, &execution_id).await {
            Ok(result) => {
                // 🗄️ CACHE STORE - Store successful result in cache
                if let Err(e) = self.store_in_cache(url, session_id, &result).await {
                    warn!("Failed to store result in cache: {}", e);
                }
                
                let content = result.get("content").and_then(|c| c.as_str()).unwrap_or("");
                
                // Create formatted response based on DEDUCT_DATA setting
                let formatted_response = self.create_formatted_scrape_response(
                    url, data_type, extraction_format, content, &execution_id
                );
                
                let tool_result = ToolResult::success_with_raw(
                    vec![McpContent::text(formatted_response)], 
                    result
                );
                
                // Apply filtering only if DEDUCT_DATA=true
                if self.is_data_reduction_enabled() {
                    Ok(ResponseStrategy::apply_size_limits(tool_result))
                } else {
                    Ok(tool_result)
                }
            }
            Err(_e) => {
                // Return empty data for BrightData errors - Anthropic will retry
                warn!("BrightData error for URL '{}', returning empty data for retry", url);
                let empty_response = json!({
                    "url": url,
                    "data_type": data_type,
                    "status": "no_data",
                    "reason": "brightdata_error",
                    "execution_id": execution_id,
                    "session_id": session_id
                });
                
                Ok(ToolResult::success_with_raw(
                    vec![McpContent::text("📊 **No Data Available**\n\nPlease try again with a different URL or check if the website is accessible.".to_string())],
                    empty_response
                ))
            }
        }
    }
}

impl Scraper {
    /// ENHANCED: Check if data reduction is enabled via DEDUCT_DATA environment variable only
    fn is_data_reduction_enabled(&self) -> bool {
        std::env::var("DEDUCT_DATA")
            .unwrap_or_else(|_| "false".to_string())
            .to_lowercase() == "true"
    }

    /// ENHANCED: Create formatted response with DEDUCT_DATA control
    fn create_formatted_scrape_response(
        &self,
        url: &str,
        data_type: &str,
        extraction_format: &str,
        content: &str,
        execution_id: &str
    ) -> String {
        // If DEDUCT_DATA=false, return full content with basic formatting
        if !self.is_data_reduction_enabled() {
            return format!(
                "📊 **Data Extraction from: {}**\n\n## Full Content\n{}\n\n*Data Type: {} | Format: {} • Execution: {}*",
                url, 
                content,
                data_type, 
                extraction_format,
                execution_id
            );
        }

        // TODO: Add filtered data extraction logic when DEDUCT_DATA=true
        // For now, return full content formatted
        format!(
            "📊 **Data Extraction from: {}**\n\n## Content (TODO: Add Filtering)\n{}\n\n*Data Type: {} | Format: {} • Execution: {}*",
            url, 
            content,
            data_type, 
            extraction_format,
            execution_id
        )
    }

    fn generate_execution_id(&self) -> String {
        format!("scrape_{}", chrono::Utc::now().format("%Y%m%d_%H%M%S%.3f"))
    }

    // 🎯 ADDED: Check Redis cache first
    async fn check_cache_first(
        &self,
        url: &str,
        session_id: &str,
    ) -> Result<Option<Value>, BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.get_cached_scrape_data(session_id, url).await
    }

    // 🗄️ ADDED: Store successful result in Redis cache
    async fn store_in_cache(
        &self,
        url: &str,
        session_id: &str,
        data: &Value,
    ) -> Result<(), BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.cache_scrape_data(session_id, url, data.clone()).await
    }

    // 🔍 ADDED: Get all cached URLs for session (useful for finding related content)
    pub async fn get_session_cached_urls(&self, session_id: &str) -> Result<Vec<String>, BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.get_session_scrape_urls(session_id).await
    }

    // 🔍 ADDED: Get cached URLs by domain
    pub async fn get_cached_urls_by_domain(
        &self,
        session_id: &str,
        domain: &str,
    ) -> Result<Vec<String>, BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.get_cached_urls_by_domain(session_id, domain).await
    }

    // 🗑️ ADDED: Clear cache for specific URL
    pub async fn clear_url_cache(
        &self,
        url: &str,
        session_id: &str,
    ) -> Result<(), BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.clear_scrape_url_cache(session_id, url).await
    }

    // 🗑️ ADDED: Clear entire session cache
    pub async fn clear_session_cache(&self, session_id: &str) -> Result<u32, BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.clear_session_scrape_cache(session_id).await
    }

    // 📊 ADDED: Get cache statistics
    pub async fn get_cache_stats(&self) -> Result<Value, BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.get_scrape_cache_stats().await
    }

    // 📊 ADDED: Get cache summary for session
    pub async fn get_cache_summary(&self, session_id: &str) -> Result<Value, BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.get_cache_summary(session_id).await
    }

    // 🏥 ADDED: Enhanced connectivity test including cache
    pub async fn test_connectivity_with_cache(&self) -> Result<String, BrightDataError> {
        let mut results = Vec::new();
        
        // Test cache connectivity
        info!("🧪 Testing Redis Cache...");
        match get_scrape_cache().await {
            Ok(cache_service) => {
                match cache_service.health_check().await {
                    Ok(_) => results.push("✅ Redis Cache: SUCCESS".to_string()),
                    Err(e) => results.push(format!("❌ Redis Cache: FAILED - {}", e)),
                }
            }
            Err(e) => results.push(format!("❌ Redis Cache: FAILED - {}", e)),
        }
        
        // Test existing connectivity
        let api_test = self.test_connectivity().await?;
        results.push(api_test);
        
        Ok(format!("🔍 Enhanced Connectivity Test Results:\n{}", results.join("\n")))
    }

    /// ENHANCED: Extract data with BrightData using proxy method (similar to forex.rs)
    async fn scrape_with_brightdata(
        &self,
        url: &str,
        data_type: &str,
        extraction_format: &str,
        clean_content: bool,
        schema: Option<Value>,
        execution_id: &str,
    ) -> Result<Value, BrightDataError> {
        let max_retries = env::var("MAX_RETRIES")
            .ok()
            .and_then(|s| s.parse::<u32>().ok())
            .unwrap_or(1);
        
        let mut last_error = None;
        
        // Get proxy configuration from environment
        let proxy_host = env::var("BRIGHTDATA_PROXY_HOST")
            .map_err(|_| BrightDataError::ToolError("Missing BRIGHTDATA_PROXY_HOST environment variable".into()))?;
        let proxy_port = env::var("BRIGHTDATA_PROXY_PORT")
            .map_err(|_| BrightDataError::ToolError("Missing BRIGHTDATA_PROXY_PORT environment variable".into()))?;
        let proxy_username = env::var("BRIGHTDATA_PROXY_USERNAME")
            .map_err(|_| BrightDataError::ToolError("Missing BRIGHTDATA_PROXY_USERNAME environment variable".into()))?;
        let proxy_password = env::var("BRIGHTDATA_PROXY_PASSWORD")
            .map_err(|_| BrightDataError::ToolError("Missing BRIGHTDATA_PROXY_PASSWORD environment variable".into()))?;

        let proxy_url = format!("http://{}:{}@{}:{}", proxy_username, proxy_password, proxy_host, proxy_port);
        
        for retry_attempt in 0..max_retries {
            let start_time = std::time::Instant::now();
            let attempt_id = format!("{}_proxy_r{}", execution_id, retry_attempt);
            
            info!("🔄 Proxy Scrape: Fetching from {} via proxy (execution: {}, retry: {}/{})", 
                  url, attempt_id, retry_attempt + 1, max_retries);
            
            if retry_attempt == 0 {
                info!("📤 Proxy Scrape Request:");
                info!("   Proxy: {}:{}@{}:{}", proxy_username, "***", proxy_host, proxy_port);
                info!("   Target: {}", url);
                info!("   Data Type: {}", data_type);
                info!("   Extraction Format: {}", extraction_format);
            }

            // Create client with proxy configuration
            let proxy = reqwest::Proxy::all(&proxy_url)
                .map_err(|e| BrightDataError::ToolError(format!("Failed to create proxy: {}", e)))?;

            let client = Client::builder()
                .proxy(proxy)
                .timeout(Duration::from_secs(120))
                .danger_accept_invalid_certs(true) // Often needed for proxy connections
                .build()
                .map_err(|e| BrightDataError::ToolError(format!("Failed to create proxy client: {}", e)))?;

            let response = client
                .get(url)
                .header("x-unblock-data-format", "markdown")
                .send()
                .await
                .map_err(|e| BrightDataError::ToolError(format!("Proxy scrape request failed to {}: {}", url, e)))?;

            let duration = start_time.elapsed();
            let status = response.status().as_u16();
            let response_headers: HashMap<String, String> = response
                .headers()
                .iter()
                .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string()))
                .collect();

            info!("📥 Proxy Scrape Response (retry {}):", retry_attempt + 1);
            info!("   Status: {}", status);
            info!("   Duration: {}ms", duration.as_millis());

            let response_text = response.text().await
                .map_err(|e| BrightDataError::ToolError(format!("Failed to read proxy scrape response body from {}: {}", url, e)))?;

            // Handle server errors with retry
            if matches!(status, 502 | 503 | 504) && retry_attempt < max_retries - 1 {
                let wait_time = Duration::from_millis(1000 + (retry_attempt as u64 * 1000));
                warn!("⏳ Proxy Scrape: Server error {}, waiting {}ms before retry...", status, wait_time.as_millis());
                tokio::time::sleep(wait_time).await;
                last_error = Some(BrightDataError::ToolError(format!("Proxy scrape server error: {}", status)));
                continue;
            }
            
            if !(200..300).contains(&status) {
                let error_msg = format!("Proxy Scrape: {} returned HTTP {}: {}", url, status, 
                                      &response_text[..response_text.len().min(200)]);
                
                warn!("Proxy scrape HTTP error: {}", error_msg);
                last_error = Some(BrightDataError::ToolError(error_msg));
                
                if retry_attempt == max_retries - 1 {
                    return Err(last_error.unwrap());
                }
                continue;
            }

            // SUCCESS - Process response
            let raw_content = response_text;

            // Print what came from BrightData Proxy
            println!("################################################################################################################");
            println!("BRIGHTDATA PROXY RAW RESPONSE FROM: {}", url);
            println!("PROXY: {}:{}", proxy_host, proxy_port);
            println!("EXECUTION: {}", execution_id);
            println!("DATA TYPE: {}", data_type);
            println!("EXTRACTION FORMAT: {}", extraction_format);
            println!("CONTENT LENGTH: {} bytes", raw_content.len());
            println!("################################################################################################################");
            println!("{}", raw_content);
            println!("################################################################################################################");
            println!("END OF BRIGHTDATA PROXY RESPONSE");
            println!("################################################################################################################");

            // Apply filters only if DEDUCT_DATA=true
            if self.is_data_reduction_enabled() {
                if ResponseFilter::is_error_page(&raw_content) {
                    return Err(BrightDataError::ToolError("Extraction returned error page".into()));
                } else if ResponseStrategy::should_try_next_source(&raw_content) {
                    return Err(BrightDataError::ToolError("Content quality too low".into()));
                }
            }

            // Print what will be sent to Anthropic
            println!("--------------------------------------------------------------------------");
            println!("SENDING TO ANTHROPIC FROM SCRAPE TOOL (PROXY):");
            println!("URL: {}", url);
            println!("DATA TYPE: {}", data_type);
            println!("EXTRACTION FORMAT: {}", extraction_format);
            println!("DATA REDUCTION ENABLED: {}", self.is_data_reduction_enabled());
            println!("CONTENT LENGTH: {} bytes", raw_content.len());
            println!("--------------------------------------------------------------------------");
            println!("{}", raw_content);
            println!("--------------------------------------------------------------------------");
            println!("END OF CONTENT SENT TO ANTHROPIC");
            println!("--------------------------------------------------------------------------");

            // Return enhanced result with additional metadata
            return Ok(json!({
                "content": raw_content,
                "metadata": {
                    "url": url,
                    "proxy_host": proxy_host,
                    "proxy_port": proxy_port,
                    "execution_id": execution_id,
                    "data_type": data_type,
                    "extraction_format": extraction_format,
                    "clean_content": clean_content,
                    "data_format": "markdown",
                    "data_reduction_enabled": self.is_data_reduction_enabled(),
                    "status_code": status,
                    "content_size_bytes": raw_content.len(),
                    "duration_ms": duration.as_millis(),
                    "timestamp": chrono::Utc::now().to_rfc3339(),
                    "retry_attempts": retry_attempt + 1,
                    "max_retries": max_retries,
                    "method": "BrightData Proxy"
                },
                "success": true
            }));
        }

        Err(last_error.unwrap_or_else(|| BrightDataError::ToolError("Proxy Scrape: All retry attempts failed".into())))
    }

    /// Test BrightData connectivity
    pub async fn test_connectivity(&self) -> Result<String, BrightDataError> {
        let test_url = "https://httpbin.org/json";
        let mut results = Vec::new();
        
        // Test BrightData API
        info!("🧪 Testing BrightData Web Unlocker...");
        match self.scrape_with_brightdata(
            test_url, "auto", "structured", true, None, "connectivity_test"
        ).await {
            Ok(_) => {
                results.push("✅ BrightData Web Unlocker: SUCCESS".to_string());
            }
            Err(e) => {
                results.push(format!("❌ BrightData Web Unlocker: FAILED - {}", e));
            }
        }
        
        Ok(format!("🔍 Connectivity Test Results:\n{}", results.join("\n")))
    }

    /// Check if URL is cached
    pub async fn is_url_cached(&self, session_id: &str, url: &str) -> Result<bool, BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.is_url_cached(session_id, url).await
    }

    /// Batch cache multiple URLs (useful for bulk operations)
    pub async fn batch_cache_urls(
        &self,
        session_id: &str,
        url_data: Vec<(String, Value)>, // Vec<(url, data)>
    ) -> Result<Vec<String>, BrightDataError> {
        let cache_service = get_scrape_cache().await?;
        cache_service.batch_cache_scrape_data(session_id, url_data).await
    }
}