FastSeek.Core/FastSeekEngine.cs

using System.Collections.Concurrent;
using System.Diagnostics;
using FastSeek.Core.Index;
using FastSeek.Core.Mft;
using FastSeek.Core.Search;
using FastSeek.Core.Utils;

namespace FastSeek.Core;

public sealed class StartupMetrics
{
    public bool CacheLoaded { get; init; }
    public int TotalFiles { get; init; }
    public double DriveDiscoveryMs { get; init; }
    public double ScanMs { get; init; }
    public double IndexMs { get; init; }
    public double CacheLoadMs { get; init; }
    public double CacheSaveMs { get; init; }
    public double TotalStartupMs { get; init; }
}

public sealed class FastSeekEngine : IDisposable
{
    private readonly ReaderWriterLockSlim _storeLock = new();
    private readonly BlockingCollection<IndexEvent> _queue = new();
    private readonly CancellationTokenSource _cts = new();
    private readonly object _checkpointLock = new();
    private readonly string _cachePath = Path.Combine(Path.GetTempPath(), "fastseek_cache.bin");

    private readonly List<JournalCheckpoint> _checkpoints = new();
    private Task? _eventTask;
    private readonly List<Task> _watcherTasks = new();

    public IndexStore Store { get; } = new();
    public bool CaseSensitive { get; set; }
    public List<string> ExcludedDirs { get; } = new();
    public StartupMetrics? LastStartupMetrics { get; private set; }
    public string CachePath => _cachePath;

    public async Task InitializeAsync(Action<string>? log = null)
    {
        log?.Invoke("Initializing FastSeek engine...");
        var swTotal = Stopwatch.StartNew();

        var swDrives = Stopwatch.StartNew();
        var drives = Drives.GetNtfsDrives();
        swDrives.Stop();
        log?.Invoke($"Detected NTFS drives: {string.Join(", ", drives.Select(d => $"{d.Letter}:"))}");
        if (drives.Count == 0) throw new InvalidOperationException("No NTFS drives found.");

        var swCacheLoad = Stopwatch.StartNew();
        log?.Invoke($"Loading cache: {_cachePath}");
        var loaded = TryLoadCache(_cachePath);
        swCacheLoad.Stop();
        log?.Invoke(loaded ? "Cache hit. Loaded index from cache." : "Cache miss. Running full MFT scan...");

        double scanMs = 0;
        double indexMs = 0;
        double cacheSaveMs = 0;

        if (!loaded)
        {
            var swScan = Stopwatch.StartNew();
            foreach (var d in drives)
            {
                log?.Invoke($"Scanning drive {d.Letter}:");
                using var reader = MftReader.Open(d);
                var direct = reader.ScanDirect(TimeSpan.FromSeconds(6), m => log?.Invoke($"[{d.Letter}:direct] {m}"));
                var scan = direct ?? reader.Scan(m => log?.Invoke($"[{d.Letter}:ioctl] {m}"));
                var method = direct is null ? "ioctl (fallback)" : "direct";
                _storeLock.EnterWriteLock();
                try { Store.PopulateFromScan(scan, d.Root); } finally { _storeLock.ExitWriteLock(); }
                log?.Invoke($"Finished drive {d.Letter}: {scan.Records.Count:N0} records ({method})");
            }
            swScan.Stop();
            scanMs = swScan.Elapsed.TotalMilliseconds;

            var swIndex = Stopwatch.StartNew();
            _storeLock.EnterWriteLock();
            try { Store.FinalizeStore(); } finally { _storeLock.ExitWriteLock(); }
            swIndex.Stop();
            indexMs = swIndex.Elapsed.TotalMilliseconds;

            var swCacheSave = Stopwatch.StartNew();
            log?.Invoke("Saving cache...");
            SaveCache(_cachePath);
            swCacheSave.Stop();
            cacheSaveMs = swCacheSave.Elapsed.TotalMilliseconds;
            log?.Invoke("Cache saved.");
        }

        _storeLock.EnterReadLock();
        try { _checkpoints.AddRange(Store.Checkpoints); }
        finally { _storeLock.ExitReadLock(); }

        foreach (var d in drives)
        {
            var drive = d;
            _watcherTasks.Add(Task.Run(() =>
            {
                try
                {
                    using var watcher = UsnWatcher.New(drive, _queue);
                    watcher.RunShared(_checkpoints, _checkpointLock, _cts.Token);
                }
                catch { }
            }, _cts.Token));
        }

        _eventTask = Task.Run(() =>
        {
            foreach (var ev in _queue.GetConsumingEnumerable(_cts.Token))
            {
                _storeLock.EnterWriteLock();
                try
                {
                    switch (ev)
                    {
                        case IndexEvent.Created c: Store.Insert(c.Record); break;
                        case IndexEvent.Deleted d: Store.Remove(d.FileRef); break;
                        case IndexEvent.Renamed r: Store.Rename(r.OldRef, r.NewRecord); break;
                        case IndexEvent.Moved m: Store.ApplyMove(m.FileRef, m.NewParentRef, m.Name, m.Kind); break;
                    }
                }
                finally { _storeLock.ExitWriteLock(); }
            }
        }, _cts.Token);

        swTotal.Stop();
        LastStartupMetrics = new StartupMetrics
        {
            CacheLoaded = loaded,
            TotalFiles = Count(),
            DriveDiscoveryMs = swDrives.Elapsed.TotalMilliseconds,
            ScanMs = scanMs,
            IndexMs = indexMs,
            CacheLoadMs = swCacheLoad.Elapsed.TotalMilliseconds,
            CacheSaveMs = cacheSaveMs,
            TotalStartupMs = swTotal.Elapsed.TotalMilliseconds,
        };
        log?.Invoke($"Startup complete in {LastStartupMetrics.TotalStartupMs:F2} ms");

        await Task.CompletedTask;
    }

    public List<SearchResult> Search(string input, int limit = 50)
    {
        if (string.IsNullOrWhiteSpace(input)) return [];
        var parsed = ParseQuery(input.Trim());

        _storeLock.EnterReadLock();
        try
        {
            if (parsed.ExtFilter is not null)
            {
                var dot = "." + parsed.ExtFilter;
                return Store.Entries.Where(e => e.NameLower.EndsWith(dot, StringComparison.Ordinal))
                    .Where(e => parsed.Filter == Filter.All || (parsed.Filter == Filter.Dirs && e.IsDir) || (parsed.Filter == Filter.Files && !e.IsDir))
                    .Select(e => new SearchResult { FullPath = SearchEngine.BuildPath(e.FileRef, Store), Name = e.Name, Rank = 0, IsDir = e.IsDir })
                    .Where(r => ExcludedDirs.Count == 0 || !ExcludedDirs.Any(ex => r.FullPath.ToLowerInvariant().StartsWith(ex, StringComparison.Ordinal)))
                    .Take(limit)
                    .ToList();
            }

            var overshootLimit = Math.Max(limit * 4, 400);
            return SearchEngine.Search(Store, parsed.Query, overshootLimit, CaseSensitive, ExcludedDirs)
                .Where(r => parsed.Filter == Filter.All || (parsed.Filter == Filter.Dirs && r.IsDir) || (parsed.Filter == Filter.Files && !r.IsDir))
                .Take(limit)
                .ToList();
        }
        finally { _storeLock.ExitReadLock(); }
    }

    public int Count()
    {
        _storeLock.EnterReadLock();
        try { return Store.Len(); }
        finally { _storeLock.ExitReadLock(); }
    }

    private bool TryLoadCache(string path)
    {
        if (!File.Exists(path)) return false;
        try
        {
            using var fs = File.OpenRead(path);
            using var br = new BinaryReader(fs, System.Text.Encoding.UTF8, leaveOpen: false);
            var magic = br.ReadUInt32();
            if (magic != 0x4B534146) return false; // FASK
            var version = br.ReadInt32();
            if (version != 1) return false;

            var driveRoot = br.ReadString();
            var checkpointCount = br.ReadInt32();
            var checkpoints = new List<JournalCheckpoint>(checkpointCount);
            for (var i = 0; i < checkpointCount; i++)
            {
                checkpoints.Add(new JournalCheckpoint
                {
                    NextUsn = br.ReadInt64(),
                    JournalId = br.ReadUInt64(),
                    DriveLetter = br.ReadChar()
                });
            }

            var entryCount = br.ReadInt32();
            var cacheData = new CacheData
            {
                DriveRoot = driveRoot,
                Checkpoints = checkpoints,
                Entries = new List<CachedEntry>(entryCount)
            };
            for (var i = 0; i < entryCount; i++)
            {
                cacheData.Entries.Add(new CachedEntry
                {
                    FileRef = br.ReadUInt64(),
                    ParentRef = br.ReadUInt64(),
                    Kind = br.ReadByte() == 1 ? FileKind.Directory : FileKind.File,
                    Name = br.ReadString()
                });
            }
            var restored = IndexStore.FromCache(cacheData);

            _storeLock.EnterWriteLock();
            try
            {
                Store.Entries.Clear();
                Store.RefLookup.Clear();
                Store.DriveRoot = restored.DriveRoot;
                Store.Checkpoints = restored.Checkpoints;
                Store.Entries.AddRange(restored.Entries);
                Store.RefLookup.AddRange(restored.RefLookup);
            }
            finally { _storeLock.ExitWriteLock(); }
            return true;
        }
        catch { return false; }
    }

    private void SaveCache(string path)
    {
        try
        {
            _storeLock.EnterReadLock();
            using var fs = File.Create(path);
            using var bw = new BinaryWriter(fs, System.Text.Encoding.UTF8, leaveOpen: false);
            bw.Write(0x4B534146u); // FASK
            bw.Write(1); // version
            bw.Write(Store.DriveRoot ?? string.Empty);
            bw.Write(Store.Checkpoints.Count);
            for (var i = 0; i < Store.Checkpoints.Count; i++)
            {
                var c = Store.Checkpoints[i];
                bw.Write(c.NextUsn);
                bw.Write(c.JournalId);
                bw.Write(c.DriveLetter);
            }

            bw.Write(Store.Entries.Count);
            for (var i = 0; i < Store.Entries.Count; i++)
            {
                var e = Store.Entries[i];
                bw.Write(e.FileRef);
                bw.Write(e.ParentRef);
                bw.Write((byte)(e.IsDir ? 1 : 0));
                bw.Write(e.Name);
            }
        }
        catch { }
    }

    public void Dispose()
    {
        _cts.Cancel();
        _queue.CompleteAdding();
        lock (_checkpointLock)
        {
            _storeLock.EnterWriteLock();
            try { Store.Checkpoints = _checkpoints.ToList(); } finally { _storeLock.ExitWriteLock(); }
        }
        SaveCache(_cachePath);
        _storeLock.Dispose();
        _queue.Dispose();
        _cts.Dispose();
    }

    private enum Filter { All, Dirs, Files }
    private sealed class ParsedQuery { public required string Query; public Filter Filter; public string? ExtFilter; }

    private static ParsedQuery ParseQuery(string input)
    {
        if (input.StartsWith("ext:", StringComparison.OrdinalIgnoreCase)) return new ParsedQuery { Query = string.Empty, Filter = Filter.Files, ExtFilter = input[4..].ToLowerInvariant() };
        if (input.StartsWith("*.", StringComparison.Ordinal)) return new ParsedQuery { Query = string.Empty, Filter = Filter.All, ExtFilter = input[2..].ToLowerInvariant() };
        if (input.StartsWith("folder:", StringComparison.OrdinalIgnoreCase)) return new ParsedQuery { Query = input[7..].Trim(), Filter = Filter.Dirs };
        if (input.StartsWith("file:", StringComparison.OrdinalIgnoreCase)) return new ParsedQuery { Query = input[5..].Trim(), Filter = Filter.Files };
        if (input.StartsWith(':')) return new ParsedQuery { Query = input[1..], Filter = Filter.Dirs };
        if (input.StartsWith('!')) return new ParsedQuery { Query = input[1..], Filter = Filter.Files };
        return new ParsedQuery { Query = input, Filter = Filter.All };
    }
}