Financial Compliance Agent

use mofa_sdk::plugins::rhai_runtime::{RhaiPlugin, RhaiPluginConfig};
use mofa_sdk::plugins::{AgentPlugin, PluginContext};
use notify::{
    Config, EventKind, RecommendedWatcher, RecursiveMode, Watcher
};
use std::path::Path;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::time::Duration;
use tracing::{info, warn};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    info!("=== Financial Compliance Agent Demo ===");

    // Define rule file path
    let plugin_dir = Path::new("./compliance_rules");
    let rules_file = plugin_dir.join("financial_rules.rhai");

    // Create initial rules file
    create_initial_rules_file(plugin_dir, &rules_file).await?;

    // Create plugin context
    let ctx = PluginContext::new("financial_compliance_agent");

    // Create and load compliance plugin
    let config = RhaiPluginConfig::new_file(
        "financial_compliance",
        &rules_file
    );
    let mut rhai_plugin = RhaiPlugin::new(config).await?;
    rhai_plugin.load(&ctx).await?;
    rhai_plugin.init_plugin().await?;
    rhai_plugin.start().await?;

    let plugin_arc = Arc::new(Mutex::new(rhai_plugin));

    // Initialize file watcher
    let (tx, mut rx) = tokio::sync::mpsc::channel(10);

    let mut watcher = RecommendedWatcher::new(
        move |res| {
            let _ = tx.blocking_send(res);
        },
        Config::default()
    )?;

    watcher.watch(&rules_file, RecursiveMode::NonRecursive)?;
    info!("Started monitoring rule file: {:?}", rules_file);

    // Clone for file change handler
    let plugin_arc_clone = Arc::clone(&plugin_arc);

    // Start file change handler
    let watcher_task = tokio::spawn(async move {
        loop {
            match rx.recv().await {
                Some(Ok(event)) => {
                    if matches!(event.kind, EventKind::Modify(_)) {
                        info!("Rule file changed, reloading...");
                        let mut plugin = plugin_arc_clone.lock().await;
                        match plugin.reload().await {
                            Ok(_) => info!("Rule plugin reloaded"),
                            Err(e) => warn!("Reload failed: {}", e),
                        }
                    }
                },
                _ => break,
            }
        }
    });

    // Test initial compliance rules
    info!("\n--- Testing Initial Compliance Rules ---");

    let test_contents = vec![
        "We offer 10% annual return products",     // Compliant
        "We offer 10% return products",            // Non-compliant
        "Guaranteed principal and returns",        // Non-compliant
        "Crypto investment services",              // Non-compliant
        "High-risk investment projects",           // Non-compliant
    ];

    let mut plugin = plugin_arc.lock().await;
    for content in test_contents {
        match plugin.execute(content.to_string()).await {
            Ok(result) => {
                info!("Content: {}", content);
                info!("Result: {}", result);
            },
            Err(e) => {
                warn!("Check failed: {}", e);
            }
        }
        println!();
    }

    // Demonstrate rule updates
    info!("\n--- Demonstrating Dynamic Rule Updates ---");
    info!("Modify file: {:?}", rules_file);
    info!("Suggested: Add 'forex margin' to sensitive words");
    info!("\nWaiting 30 seconds for manual modification...\n");

    tokio::time::sleep(Duration::from_secs(30)).await;

    // Test updated rules
    info!("\n--- Testing Updated Rules ---");
    let updated_tests = vec![
        "We offer 10% annual return products",
        "We provide forex margin trading",
    ];

    for content in updated_tests {
        match plugin.execute(content.to_string()).await {
            Ok(result) => {
                info!("Content: {}", content);
                info!("Result: {}", result);
            },
            Err(e) => {
                warn!("Check failed: {}", e);
            }
        }
        println!();
    }

    // Cleanup
    plugin.stop().await?;
    plugin.unload().await?;
    watcher_task.abort();

    tokio::fs::remove_file(rules_file).await?;
    tokio::fs::remove_dir(plugin_dir).await?;

    info!("\n=== Financial Compliance Agent Demo Complete ===");
    Ok(())
}

async fn create_initial_rules_file(
    plugin_dir: &Path,
    rules_file: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    tokio::fs::create_dir_all(plugin_dir).await?;

    let initial_rules = r#"
// Financial Compliance Rules Engine
// @name: FinancialComplianceRules
// @version: 1.0.0

fn execute(content) {
    // Rule 1: Forbidden sensitive products
    let sensitive_products = [
        "high-risk", "illegal fundraising", 
        "cryptocurrency", "pyramid scheme"
    ];
    for product in sensitive_products {
        if content.contains(product) {
            return #{
                compliant: false,
                reason: "Contains forbidden product: " + product,
                rule_id: "FIN-001",
                content: content
            };
        }
    }

    // Rule 2: Interest rate disclosure
    let interest_keywords = ["rate", "annual", "monthly", "return"];
    let has_interest = interest_keywords.some(|k| content.contains(k));
    let has_period = content.contains("annual") || content.contains("monthly");

    if has_interest && !has_period {
        return #{
            compliant: false,
            reason: "Interest rate must specify annual/monthly",
            rule_id: "FIN-002",
            content: content
        };
    }

    // Rule 3: No guarantee promises
    let guarantee_words = ["guaranteed", "principal", "risk-free"];
    for word in guarantee_words {
        if content.contains(word) {
            return #{
                compliant: false,
                reason: "Cannot promise guarantees",
                rule_id: "FIN-003",
                content: content
            };
        }
    }

    // All rules passed
    return #{
        compliant: true,
        reason: "Compliance check passed",
        rule_id: "PASS",
        content: content
    };
}
"#;

    tokio::fs::write(rules_file, initial_rules).await?;
    info!("Created initial compliance rules: {:?}", rules_file);
    Ok(())
}