raxit-core 0.1.2

////! CrewAI framework extractor
//!
//! Extracts agents, tasks, and crews from CrewAI decorator-based code.
//!
//! CrewAI uses decorators (@agent, @task, @crew) on functions that return
//! Agent(), Task(), and Crew() instances.

use super::ExtractedAssets;
use crate::error::Result;
use crate::schema::{Agent, SourceLocation, Tool};
use std::collections::HashMap;
use std::path::Path;

/// Extract assets from a CrewAI file
pub fn extract(path: &Path) -> Result<ExtractedAssets> {
    // Read source code for pattern matching
    let _source_code = std::fs::read_to_string(path)?;

    // Parse the file using tree-sitter
    let nodes = crate::ast::parse_python_file(path)?;

    let mut assets = ExtractedAssets::default();
    let mut agent_functions = HashMap::new(); // function_name -> agent_id
    let mut task_functions = HashMap::new(); // function_name -> (task info)
    let mut crew_info = Vec::new(); // Store crew information

    // First pass: Extract @agent decorated functions
    for node in &nodes {
        if node.kind == "decorated_definition" && node.text.contains("@agent") {
            if let Some(func_name) = extract_function_name(&node.text) {
                // Extract Agent() call details from the function body
                if let Some(agent) = extract_agent_from_decorated_function(node, path, &func_name) {
                    agent_functions.insert(func_name.clone(), agent.id.clone());
                    assets.agents.push(agent);
                }
            }
        }
    }

    // Second pass: Extract @task decorated functions
    for node in &nodes {
        if node.kind == "decorated_definition" && node.text.contains("@task") {
            if let Some(func_name) = extract_function_name(&node.text) {
                if let Some(task_info) =
                    extract_task_from_decorated_function(node, path, &agent_functions)
                {
                    task_functions.insert(func_name.clone(), task_info);
                }
            }
        }
    }

    // Third pass: Extract @crew decorated functions
    for node in &nodes {
        if node.kind == "decorated_definition" && node.text.contains("@crew") {
            if let Some(crew) = extract_crew_from_decorated_function(node, path) {
                crew_info.push(crew);
            }
        }
    }

    // Fourth pass: Extract tools from module-level assignments
    // Pattern: tool_name = SomeToolClass()
    for node in &nodes {
        if node.kind == "assignment" && contains_tool_pattern(&node.text) {
            if let Some(tool) = extract_tool_from_assignment(node, path) {
                assets.tools.push(tool);
            }
        }
    }

    Ok(assets)
}

fn extract_agent_from_decorated_function(
    node: &crate::ast::AstNode,
    path: &Path,
    func_name: &str,
) -> Option<Agent> {
    let id = format!("crewai_agent_{}", node.start_line);

    // Extract Agent() parameters
    let role = extract_parameter_value(&node.text, "role");
    let goal = extract_parameter_value(&node.text, "goal");
    let backstory = extract_parameter_value(&node.text, "backstory");
    let tools = extract_tools_list(&node.text);

    // Use role as the agent name, fallback to function name
    let name = role.clone().unwrap_or_else(|| func_name.to_string());

    // Build system prompt from role, goal, and backstory
    let system_prompt = build_system_prompt(&role, &goal, &backstory);

    Some(Agent {
        id,
        name,
        location: SourceLocation {
            file: path.to_string_lossy().to_string(),
            line: node.start_line,
            end_line: Some(node.end_line),
            function: Some(func_name.to_string()),
        },
        model_id: None, // CrewAI doesn't explicitly specify model in Agent()
        tool_ids: tools,
        memory_id: None,
        system_prompt,
        result_type: Some("CrewAI::Agent".to_string()),
        deps_type: None,
    })
}

#[allow(dead_code)]
struct TaskInfo {
    description: Option<String>,
    agent_id: Option<String>,
    expected_output: Option<String>,
}

fn extract_task_from_decorated_function(
    node: &crate::ast::AstNode,
    _path: &Path,
    agent_functions: &HashMap<String, String>,
) -> Option<TaskInfo> {
    let description = extract_parameter_value(&node.text, "description");
    let agent_func = extract_parameter_value(&node.text, "agent");
    let expected_output = extract_parameter_value(&node.text, "expected_output");

    // Resolve agent function name to agent ID
    let agent_id = agent_func.and_then(|func| agent_functions.get(&func).cloned());

    Some(TaskInfo {
        description,
        agent_id,
        expected_output,
    })
}

#[allow(dead_code)]
struct CrewInfo {
    agents: Vec<String>,
    tasks: Vec<String>,
    process: Option<String>,
}

fn extract_crew_from_decorated_function(
    node: &crate::ast::AstNode,
    _path: &Path,
) -> Option<CrewInfo> {
    let agents = extract_list_parameter(&node.text, "agents");
    let tasks = extract_list_parameter(&node.text, "tasks");
    let process = extract_parameter_value(&node.text, "process");

    Some(CrewInfo {
        agents,
        tasks,
        process,
    })
}

fn extract_tool_from_assignment(node: &crate::ast::AstNode, path: &Path) -> Option<Tool> {
    // Extract variable name (left side of =)
    let var_name = extract_variable_name(&node.text)?;

    // Extract tool class name (right side)
    let tool_class = extract_tool_class(&node.text)?;

    let id = format!("crewai_tool_{}", node.start_line);

    Some(Tool {
        id,
        name: var_name.clone(),
        location: SourceLocation {
            file: path.to_string_lossy().to_string(),
            line: node.start_line,
            end_line: Some(node.end_line),
            function: None,
        },
        description: Some(format!("CrewAI tool: {tool_class}")),
        parameters: None,
        requires_context: false,
        tool_type: tool_class,
        data_flows: Vec::new(),
    })
}

fn extract_function_name(text: &str) -> Option<String> {
    // Extract function name from definition: def function_name():
    if let Some(idx) = text.find("def ") {
        let rest = &text[idx + 4..];
        if let Some(end_idx) = rest.find('(') {
            return Some(rest[..end_idx].trim().to_string());
        }
    }
    None
}

fn extract_parameter_value(text: &str, param_name: &str) -> Option<String> {
    // Look for param_name="value" or param_name='value'
    let pattern = format!("{param_name}=");
    if let Some(start) = text.find(&pattern) {
        let rest = &text[start + pattern.len()..];

        // Skip whitespace
        let rest = rest.trim_start();

        // Check for string literal
        if rest.starts_with('"') || rest.starts_with('\'') {
            let quote_char = rest.chars().next().unwrap();
            let content = &rest[1..];

            // Handle multi-line strings
            if content.starts_with(quote_char) && content.chars().nth(1) == Some(quote_char) {
                // Triple-quoted string
                let triple_quote = format!("{quote_char}{quote_char}{quote_char}");
                let after_open = &content[2..];
                if let Some(end) = after_open.find(&triple_quote) {
                    return Some(after_open[..end].trim().to_string());
                }
            } else {
                // Regular string
                if let Some(end) = content.find(quote_char) {
                    return Some(content[..end].to_string());
                }
            }
        } else if rest.starts_with("Process.") {
            // Handle Process.sequential, Process.hierarchical, etc.
            if let Some(end) = rest.find(&[',', ')', '\n'][..]) {
                return Some(rest[..end].trim().to_string());
            }
        } else {
            // Handle other values (function references, etc.)
            if let Some(end) = rest.find(&[',', ')', '\n'][..]) {
                let value = rest[..end].trim();
                if !value.is_empty() {
                    return Some(value.to_string());
                }
            }
        }
    }
    None
}

fn extract_tools_list(text: &str) -> Vec<String> {
    // Look for tools=[tool1, tool2, ...]
    if let Some(start) = text.find("tools=") {
        let rest = &text[start + 6..];
        let rest = rest.trim_start();

        if let Some(content) = rest.strip_prefix('[') {
            if let Some(end) = content.find(']') {
                let tools_str = &content[..end];
                return tools_str
                    .split(',')
                    .map(|s| s.trim().to_string())
                    .filter(|s| !s.is_empty())
                    .collect();
            }
        }
    }
    Vec::new()
}

fn extract_list_parameter(text: &str, param_name: &str) -> Vec<String> {
    // Look for param_name=[item1, item2, ...]
    let pattern = format!("{param_name}=");
    if let Some(start) = text.find(&pattern) {
        let rest = &text[start + pattern.len()..];
        let rest = rest.trim_start();

        if let Some(content) = rest.strip_prefix('[') {
            if let Some(end) = content.find(']') {
                let list_str = &content[..end];
                return list_str
                    .split(',')
                    .map(|s| s.trim().to_string())
                    .filter(|s| !s.is_empty())
                    .collect();
            }
        }
    }
    Vec::new()
}

fn build_system_prompt(
    role: &Option<String>,
    goal: &Option<String>,
    backstory: &Option<String>,
) -> Option<String> {
    let mut parts = Vec::new();

    if let Some(r) = role {
        parts.push(format!("Role: {r}"));
    }
    if let Some(g) = goal {
        parts.push(format!("Goal: {g}"));
    }
    if let Some(b) = backstory {
        parts.push(format!("Backstory: {b}"));
    }

    if parts.is_empty() {
        None
    } else {
        Some(parts.join("\n"))
    }
}

fn contains_tool_pattern(text: &str) -> bool {
    // Check if text contains tool instantiation patterns
    text.contains("Tool()")
        || text.contains("SerperDevTool()")
        || text.contains("ScrapeWebsiteTool()")
        || text.contains("WebsiteSearchTool()")
        || text.contains("FileReadTool()")
        || text.contains("DirectoryReadTool()")
        || (text.contains("_tool") && text.contains("()"))
}

fn extract_variable_name(text: &str) -> Option<String> {
    // Extract variable name from assignment: var_name = value
    if let Some(eq_idx) = text.find(" = ") {
        let before = &text[..eq_idx].trim();
        // Take the last word (variable name)
        if let Some(var_name) = before.split_whitespace().last() {
            return Some(var_name.to_string());
        }
    }
    None
}

fn extract_tool_class(text: &str) -> Option<String> {
    // Extract class name from instantiation: ClassName()
    if let Some(eq_idx) = text.find(" = ") {
        let after = &text[eq_idx + 3..].trim();
        if let Some(paren_idx) = after.find('(') {
            let class_name = after[..paren_idx].trim();
            if !class_name.is_empty() {
                return Some(class_name.to_string());
            }
        }
    }
    None
}

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

    #[test]
    fn test_extract_function_name() {
        let text = "def researcher() -> Agent:";
        let name = extract_function_name(text);
        assert_eq!(name, Some("researcher".to_string()));
    }

    #[test]
    fn test_extract_parameter_value_simple() {
        let text = r#"Agent(role="Research Analyst", goal="Find information")"#;

        let role = extract_parameter_value(text, "role");
        assert_eq!(role, Some("Research Analyst".to_string()));

        let goal = extract_parameter_value(text, "goal");
        assert_eq!(goal, Some("Find information".to_string()));
    }

    #[test]
    fn test_extract_tools_list() {
        let text = "Agent(role='test', tools=[search_tool, scrape_tool], verbose=True)";
        let tools = extract_tools_list(text);
        assert_eq!(tools, vec!["search_tool", "scrape_tool"]);
    }

    #[test]
    fn test_extract_list_parameter() {
        let text = "Crew(agents=[researcher, analyst], tasks=[task1, task2])";

        let agents = extract_list_parameter(text, "agents");
        assert_eq!(agents, vec!["researcher", "analyst"]);

        let tasks = extract_list_parameter(text, "tasks");
        assert_eq!(tasks, vec!["task1", "task2"]);
    }

    #[test]
    fn test_build_system_prompt() {
        let role = Some("Research Analyst".to_string());
        let goal = Some("Find information".to_string());
        let backstory = Some("Expert researcher".to_string());

        let prompt = build_system_prompt(&role, &goal, &backstory);
        assert!(prompt.is_some());
        let prompt = prompt.unwrap();
        assert!(prompt.contains("Role: Research Analyst"));
        assert!(prompt.contains("Goal: Find information"));
        assert!(prompt.contains("Backstory: Expert researcher"));
    }

    #[test]
    fn test_extract_variable_name() {
        let text = "search_tool = SerperDevTool()";
        let var = extract_variable_name(text);
        assert_eq!(var, Some("search_tool".to_string()));
    }

    #[test]
    fn test_extract_tool_class() {
        let text = "search_tool = SerperDevTool()";
        let class = extract_tool_class(text);
        assert_eq!(class, Some("SerperDevTool".to_string()));
    }

    #[test]
    fn test_contains_tool_pattern() {
        assert!(contains_tool_pattern("search_tool = SerperDevTool()"));
        assert!(contains_tool_pattern("tool = WebsiteSearchTool()"));
        assert!(!contains_tool_pattern("agent = Agent()"));
    }
}