magi-code 0.96.1

use super::{
    ToolResult, ToolResultDisplay, ToolRuntime,
    args::{FIND_DEFAULT_LIMIT, FindArgs, FindKindArg},
    contract::{metadata_key as meta, tool_name},
    workspace,
};
use crate::cancellation::AgentCancellation;
use serde_json::json;
use std::{
    cmp::Ordering,
    collections::BinaryHeap,
    path::{Path, PathBuf},
    time::{Duration, Instant},
};

pub(crate) const FIND_MAX_VISITED_ENTRIES: usize = 100_000;
pub(crate) const FIND_DEADLINE_SECS: u64 = 10;

#[derive(Debug, Default)]
pub(super) struct FindSearchService;

impl FindSearchService {
    pub(super) fn path_search(
        &self,
        walker: &workspace::WorkspaceWalker,
        request: FindPathSearchRequest,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<FindPathSearchResult> {
        self.path_search_with_bounds(
            walker,
            request,
            cancellation,
            Instant::now() + Duration::from_secs(FIND_DEADLINE_SECS),
            FIND_MAX_VISITED_ENTRIES,
        )
    }

    fn path_search_with_bounds(
        &self,
        walker: &workspace::WorkspaceWalker,
        request: FindPathSearchRequest,
        cancellation: &AgentCancellation,
        deadline: Instant,
        max_entries: usize,
    ) -> anyhow::Result<FindPathSearchResult> {
        cancellation.check()?;
        let cwd = request.cwd.canonicalize()?;
        let (search_root, output_prefix) = match request.path_filter.as_ref() {
            Some(path_filter) => {
                let resolved = cwd.join(path_filter).canonicalize()?;
                if !resolved.is_dir() {
                    anyhow::bail!("find path search path must be a directory");
                }
                let prefix = match resolved.strip_prefix(&cwd) {
                    Ok(relative) => normalize_filter_prefix(relative),
                    Err(_) => resolved.to_string_lossy().replace('\\', "/"),
                };
                (resolved, prefix)
            }
            None => (cwd.clone(), String::new()),
        };

        let scoped_query = if output_prefix.is_empty() {
            request.query.clone()
        } else {
            query_without_filter_prefix(&request.query, &output_prefix)
                .unwrap_or_else(|| request.query.clone())
        };
        let normalized_query = normalize_query(&scoped_query);
        let mut matches = PathMatchAccumulator::new(request.offset.saturating_add(request.limit));
        let search_root_relative = search_root.starts_with(&cwd);
        let status = walker.visit_entries_with_bounds(
            workspace::WorkspaceWalkOptions {
                root: &search_root,
                skip_dirs: &[".git", "target"],
                cancel_interval: 32,
            },
            Some(cancellation),
            deadline,
            max_entries,
            |path, file_type| {
                let item_kind = if file_type.is_file() {
                    FindPathKind::File
                } else if file_type.is_dir() {
                    FindPathKind::Directory
                } else {
                    return Ok(true);
                };
                if !kind_includes(request.kind, item_kind) {
                    return Ok(true);
                }
                let relative = if search_root_relative {
                    normalize_relative_path(path.strip_prefix(&cwd).unwrap_or(path))
                } else {
                    path.to_string_lossy().replace('\\', "/")
                };
                if relative.is_empty()
                    || (search_root_relative && relative.starts_with('/'))
                    || (search_root_relative && relative.contains(":/"))
                {
                    return Ok(true);
                }
                if !output_prefix.is_empty() && !path_matches_filter(&relative, &output_prefix) {
                    return Ok(true);
                }
                let Some(score) = fuzzy_score(&relative, &normalized_query) else {
                    return Ok(true);
                };
                matches.push(ScoredPathMatch {
                    item: FindPathMatch {
                        relative_path: relative,
                        kind: item_kind,
                    },
                    score,
                });
                Ok(true)
            },
        )?;
        let total_matched = matches.total_matched;
        let items = matches
            .into_ranked()
            .into_iter()
            .skip(request.offset)
            .take(request.limit)
            .map(|matched| matched.item)
            .collect::<Vec<_>>();
        let truncated =
            !status.completed || request.offset.saturating_add(items.len()) < total_matched;

        Ok(FindPathSearchResult {
            items,
            total_matched,
            truncated,
            index_ready: status.completed,
            scan_complete: status.completed,
            entries_scanned: status.entries_visited,
            entry_limit_reached: status.entry_limit_reached,
            timed_out: status.timed_out,
            walk_errors: status.walk_errors,
            walk_entries_omitted: status.entries_omitted,
        })
    }

    pub(super) fn index_count(&self) -> usize {
        0
    }
}

#[derive(Debug)]
pub(super) struct FindPathSearchRequest {
    pub(super) cwd: PathBuf,
    pub(super) query: String,
    pub(super) path_filter: Option<PathBuf>,
    pub(super) kind: FindKind,
    pub(super) offset: usize,
    pub(super) limit: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum FindKind {
    Files,
    Directories,
    Mixed,
}

impl FindKind {
    fn as_str(self) -> &'static str {
        match self {
            Self::Files => "files",
            Self::Directories => "directories",
            Self::Mixed => "mixed",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct FindPathSearchResult {
    pub(super) items: Vec<FindPathMatch>,
    /// Matches observed so far; exact only when `scan_complete` is true.
    pub(super) total_matched: usize,
    pub(super) truncated: bool,
    pub(super) index_ready: bool,
    pub(super) scan_complete: bool,
    pub(super) entries_scanned: usize,
    pub(super) entry_limit_reached: bool,
    pub(super) timed_out: bool,
    pub(super) walk_errors: usize,
    pub(super) walk_entries_omitted: usize,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct FindPathMatch {
    pub(super) relative_path: String,
    pub(super) kind: FindPathKind,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum FindPathKind {
    File,
    Directory,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct ScoredPathMatch {
    item: FindPathMatch,
    score: usize,
}

#[derive(Debug)]
struct PathMatchAccumulator {
    matches: BinaryHeap<HeapScoredPathMatch>,
    total_matched: usize,
    capacity: usize,
}

impl PathMatchAccumulator {
    fn new(capacity: usize) -> Self {
        Self {
            matches: BinaryHeap::new(),
            total_matched: 0,
            capacity,
        }
    }

    fn push(&mut self, matched: ScoredPathMatch) {
        self.total_matched = self.total_matched.saturating_add(1);
        if self.capacity == 0 {
            return;
        }
        if self.matches.len() < self.capacity {
            self.matches.push(HeapScoredPathMatch(matched));
            return;
        }

        let should_replace = self
            .matches
            .peek()
            .is_some_and(|worst| scored_path_match_cmp(&matched, &worst.0).is_lt());
        if should_replace {
            self.matches.pop();
            self.matches.push(HeapScoredPathMatch(matched));
        }
    }

    fn into_ranked(self) -> Vec<ScoredPathMatch> {
        let mut matches = self
            .matches
            .into_iter()
            .map(|matched| matched.0)
            .collect::<Vec<_>>();
        matches.sort_by(scored_path_match_cmp);
        matches
    }
}

#[derive(Debug, Eq, PartialEq)]
struct HeapScoredPathMatch(ScoredPathMatch);

// Keep the worst retained match at the heap root so a better match can
// replace it without retaining every candidate from the walk.
impl Ord for HeapScoredPathMatch {
    fn cmp(&self, other: &Self) -> Ordering {
        scored_path_match_cmp(&self.0, &other.0)
    }
}

impl PartialOrd for HeapScoredPathMatch {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl ToolRuntime {
    pub(super) fn find(
        &self,
        args: FindArgs,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<ToolResult> {
        let kind = match args.kind.unwrap_or(FindKindArg::Files) {
            FindKindArg::Files => FindKind::Files,
            FindKindArg::Directories => FindKind::Directories,
            FindKindArg::Mixed => FindKind::Mixed,
        };
        let limit = args.limit.unwrap_or(FIND_DEFAULT_LIMIT);
        let offset = args.offset.unwrap_or(0);
        let path_filter = args
            .path
            .as_deref()
            .map(|path| self.find_path_filter(path))
            .transpose()?;

        let result = self.path_search.path_search(
            &self.workspace_walker,
            FindPathSearchRequest {
                cwd: self.cwd_canonical.clone(),
                query: args.query.clone(),
                path_filter,
                kind,
                offset,
                limit,
            },
            cancellation,
        )?;
        let output = result
            .items
            .iter()
            .map(|item| match item.kind {
                FindPathKind::File => item.relative_path.clone(),
                FindPathKind::Directory => format!("{}/", item.relative_path),
            })
            .collect::<Vec<_>>()
            .join("\n");

        Ok(ToolResult {
            tool_name: tool_name::FIND.to_string(),
            success: true,
            content: output,
            metadata: json!({
                (meta::ENGINE): "rust",
                (meta::QUERY): args.query,
                (meta::PATH): args.path,
                (meta::KIND): kind.as_str(),
                (meta::LIMIT): limit,
                (meta::OFFSET): offset,
                (meta::MATCHES_RETURNED): result.items.len(),
                (meta::TOTAL_MATCHED): result.total_matched,
                (meta::TRUNCATED): result.truncated,
                (meta::INDEX_READY): result.index_ready,
                "scan_complete": result.scan_complete,
                "entries_scanned": result.entries_scanned,
                "entry_limit_reached": result.entry_limit_reached,
                (meta::TIMED_OUT): result.timed_out,
                "walk_errors": result.walk_errors,
                "walk_entries_omitted": result.walk_entries_omitted,
                "total_matched_exact": result.scan_complete,
                "has_more": if result.scan_complete { Some(result.truncated) } else { None },
                "structured_result": {
                    "files": result.items.iter().map(|item| item.relative_path.clone()).collect::<Vec<_>>(),
                    "complete": result.scan_complete && !result.truncated && offset == 0,
                    "scan_complete": result.scan_complete,
                    "has_more": if result.scan_complete { Some(result.truncated) } else { None },
                    "next_offset": result.truncated.then_some(offset.saturating_add(result.items.len())),
                },
            }),
            display: ToolResultDisplay::default(),
        })
    }

    fn find_path_filter(&self, path: &str) -> anyhow::Result<PathBuf> {
        let resolved = self.resolve_existing_path(
            path,
            super::fs::ExistingPathPolicy::find(self.find_absolute_paths),
        )?;
        if !resolved.is_dir() {
            anyhow::bail!("find path must be a directory");
        }
        if !resolved.starts_with(&self.cwd_canonical) && !self.find_absolute_paths {
            self.ensure_inside_with_setting(&resolved, "tools.find.absolute_paths")?;
        }
        Ok(resolved
            .strip_prefix(&self.cwd_canonical)
            .unwrap_or(&resolved)
            .to_path_buf())
    }
}

fn kind_includes(kind: FindKind, item_kind: FindPathKind) -> bool {
    matches!(
        (kind, item_kind),
        (FindKind::Files, FindPathKind::File)
            | (FindKind::Directories, FindPathKind::Directory)
            | (FindKind::Mixed, _)
    )
}

fn path_kind_rank(kind: FindPathKind) -> u8 {
    match kind {
        FindPathKind::File => 0,
        FindPathKind::Directory => 1,
    }
}

fn scored_path_match_cmp(left: &ScoredPathMatch, right: &ScoredPathMatch) -> Ordering {
    right
        .score
        .cmp(&left.score)
        .then_with(|| left.item.relative_path.cmp(&right.item.relative_path))
        .then_with(|| path_kind_rank(left.item.kind).cmp(&path_kind_rank(right.item.kind)))
}

fn normalize_query(query: &str) -> String {
    query.trim().to_lowercase()
}

fn fuzzy_score(path: &str, normalized_query: &str) -> Option<usize> {
    if normalized_query.is_empty() {
        return Some(0);
    }
    let path_lower = path.to_lowercase();
    let mut score = 0usize;
    for token in normalized_query.split_whitespace() {
        if token.is_empty() {
            continue;
        }
        if let Some(index) = path_lower.find(token) {
            score = score.saturating_add(10_000usize.saturating_sub(index));
        } else if fuzzy_subsequence_match(&path_lower, token) {
            score = score.saturating_add(token.len());
        } else {
            return None;
        }
    }
    Some(score)
}

fn fuzzy_subsequence_match(path: &str, token: &str) -> bool {
    let mut chars = token.chars();
    let mut current = chars.next();
    if current.is_none() {
        return true;
    }
    for path_char in path.chars() {
        if Some(path_char) == current {
            current = chars.next();
            if current.is_none() {
                return true;
            }
        }
    }
    false
}

fn normalize_relative_path(path: &Path) -> String {
    path.to_string_lossy()
        .replace('\\', "/")
        .trim_matches('/')
        .to_string()
}

fn normalize_filter_prefix(path: &Path) -> String {
    normalize_relative_path(path)
}

fn path_matches_filter(relative_path: &str, prefix: &str) -> bool {
    prefix.is_empty()
        || relative_path == prefix
        || relative_path
            .strip_prefix(prefix)
            .is_some_and(|suffix| suffix.starts_with('/'))
}

fn query_without_filter_prefix(query: &str, prefix: &str) -> Option<String> {
    let prefix = prefix.trim_matches('/');
    if prefix.is_empty() {
        return None;
    }

    let query = query.trim_start();
    let rest = query.strip_prefix(prefix)?;
    if rest.is_empty() {
        return Some(String::new());
    }
    if let Some(rest) = rest.strip_prefix('/') {
        return Some(rest.trim_start_matches('/').trim_start().to_string());
    }
    if rest.chars().next().is_some_and(char::is_whitespace) {
        return Some(rest.trim_start().to_string());
    }
    None
}