MemoryService.cs

using System.Text;
using System.Text.RegularExpressions;

namespace KaiChat.Services;

public interface IMemoryService
{
    string MemoryDirectoryRelativePath { get; }
    string UpdateMarkerRelativePath { get; }
    IDisposable EnterPersonaScope(string? persona);
    Task<string> GetPersistentMemoryAsync();
    Task<IReadOnlyList<MemoryFileInfo>> ListMemoryFilesAsync(bool includeTrash = false);
    Task<string> ReadMemoryFileAsync(string relativePath);
    Task SaveMemoryFileAsync(string relativePath, string content);
    Task SaveMemoryAsync(string section, string content);
    Task SaveFullAsync(string content);
    Task SoftDeleteMemoryFileAsync(string relativePath, string reason);
    Task DeleteSectionAsync(string section);
    Task<DateTime> GetLastUpdateAsync();
    Task<bool> EnsureUpdateMarkerAsync(DateTime timestamp);
    Task SetLastUpdateAsync(DateTime timestamp);
    Task DeleteUpdateMarkerAsync();
    Task<MemoryUpdateMarkerInfo> GetUpdateMarkerInfoAsync();
}

public record MemoryFileInfo(
    string Path,
    long Size,
    DateTime UpdatedAt,
    bool IsTrash,
    bool IsLoaded);

public record MemoryUpdateMarkerInfo(
    string Path,
    bool Exists,
    string RawValue,
    DateTime? ParsedUtc,
    DateTime? LastWriteTimeUtc);

public class MemoryService : IMemoryService
{
    private readonly string _baseFolder;
    private readonly string _globalMemoryDir;
    private readonly string _markerFile;
    private readonly ILogger<MemoryService> _logger;
    private readonly SemaphoreSlim _lock = new(1, 1);
    private readonly AsyncLocal<string?> _personaScope = new();
    private static readonly string[] PreferredLoadOrder =
    [
        "core.md",
        "recent.md",
        "history.md",
        "communication.md",
        "raymond.md",
        "kai.md",
        "health.md",
        "story-world.md",
        "projects.md",
        "preferences.md",
        "index.md",
        "manifest.md"
    ];

    public MemoryService(IConfiguration config, ILogger<MemoryService> logger)
    {
        _baseFolder = config.GetValue<string>("KaiChat:BaseFolder") ?? ".";
        var kaiChatDir = Path.Combine(_baseFolder, "KaiChat");
        _globalMemoryDir = Path.Combine(kaiChatDir, "Memory");
        _markerFile = Path.Combine(kaiChatDir, ".memory_update_marker");
        _logger = logger;
    }

    private string _memoryDir => _personaScope.Value is { Length: > 0 } persona
        ? Path.Combine(_baseFolder, "KaiChat", "Personas", persona, "Memory")
        : _globalMemoryDir;

    private string _trashDir => Path.Combine(_memoryDir, "_trash");

    public string MemoryDirectoryRelativePath => Path
        .GetRelativePath(_baseFolder, _memoryDir)
        .Replace('\\', '/');

    public string UpdateMarkerRelativePath => Path
        .GetRelativePath(_baseFolder, _markerFile)
        .Replace('\\', '/');

    /// <summary>
    /// Sets the persona scope for the current async context.
    /// When set, all memory operations read/write from KaiChat/Personas/{persona}/Memory/
    /// instead of the global KaiChat/Memory/.
    /// Returns an IDisposable that restores the previous scope on Dispose.
    /// </summary>
    public IDisposable EnterPersonaScope(string? persona)
    {
        var previous = _personaScope.Value;
        _personaScope.Value = string.IsNullOrWhiteSpace(persona) ? null : persona;
        return new PersonaScopeExit(this, previous);
    }

    private sealed class PersonaScopeExit : IDisposable
    {
        private readonly MemoryService _owner;
        private readonly string? _previous;
        public PersonaScopeExit(MemoryService owner, string? previous)
        {
            _owner = owner;
            _previous = previous;
        }
        public void Dispose() => _owner._personaScope.Value = _previous;
    }

    public async Task<DateTime> GetLastUpdateAsync()
    {
        if (!File.Exists(_markerFile))
            return DateTime.MaxValue;

        var text = await File.ReadAllTextAsync(_markerFile);
        if (DateTimeOffset.TryParse(text.Trim(), out var dto))
            return dto.UtcDateTime;

        _logger.LogWarning("Memory update marker could not be parsed: {Path}", _markerFile);
        return DateTime.MaxValue;
    }

    public async Task<bool> EnsureUpdateMarkerAsync(DateTime timestamp)
    {
        await _lock.WaitAsync();
        try
        {
            if (File.Exists(_markerFile))
                return false;

            Directory.CreateDirectory(Path.GetDirectoryName(_markerFile) ?? ".");
            await File.WriteAllTextAsync(_markerFile, timestamp.ToUniversalTime().ToString("O"));
            _logger.LogInformation("Created memory update marker at {Timestamp}", timestamp.ToUniversalTime());
            return true;
        }
        finally { _lock.Release(); }
    }

    public async Task SetLastUpdateAsync(DateTime timestamp)
    {
        await _lock.WaitAsync();
        try
        {
            Directory.CreateDirectory(Path.GetDirectoryName(_markerFile) ?? ".");
            await File.WriteAllTextAsync(_markerFile, timestamp.ToUniversalTime().ToString("O"));
        }
        finally { _lock.Release(); }
    }

    public async Task DeleteUpdateMarkerAsync()
    {
        await _lock.WaitAsync();
        try
        {
            if (File.Exists(_markerFile))
            {
                File.Delete(_markerFile);
                _logger.LogInformation("Deleted memory update marker; no pending memory update remains");
            }
        }
        finally { _lock.Release(); }
    }

    public async Task<MemoryUpdateMarkerInfo> GetUpdateMarkerInfoAsync()
    {
        if (!File.Exists(_markerFile))
            return new MemoryUpdateMarkerInfo(
                Path: _markerFile,
                Exists: false,
                RawValue: "",
                ParsedUtc: null,
                LastWriteTimeUtc: null);

        var info = new FileInfo(_markerFile);
        var raw = await File.ReadAllTextAsync(_markerFile);
        var parsed = DateTimeOffset.TryParse(raw.Trim(), out var dto)
            ? dto.UtcDateTime
            : (DateTime?)null;

        return new MemoryUpdateMarkerInfo(
            Path: _markerFile,
            Exists: true,
            RawValue: raw.Trim(),
            ParsedUtc: parsed,
            LastWriteTimeUtc: info.LastWriteTimeUtc);
    }

    public async Task<string> GetPersistentMemoryAsync()
    {
        await EnsureDirectoryMemorySeededAsync();

        var files = GetLoadableMemoryFiles().ToList();
        if (files.Count > 0)
            return await ReadMemoryPackAsync(files);

        return "";
    }

    public async Task<IReadOnlyList<MemoryFileInfo>> ListMemoryFilesAsync(bool includeTrash = false)
    {
        await EnsureDirectoryMemorySeededAsync();
        if (!Directory.Exists(_memoryDir))
            return [];

        var files = Directory.GetFiles(_memoryDir, "*.md", SearchOption.AllDirectories)
            .Where(path => includeTrash || !IsTrashPath(path))
            .Select(path =>
            {
                var info = new FileInfo(path);
                return new MemoryFileInfo(
                    ToMemoryRelativePath(path),
                    info.Length,
                    info.LastWriteTimeUtc,
                    IsTrashPath(path),
                    IsLoadableMemoryPath(path));
            })
            .OrderBy(f => f.IsTrash)
            .ThenBy(f => LoadOrder(f.Path))
            .ThenBy(f => f.Path, StringComparer.OrdinalIgnoreCase)
            .ToList();

        return files;
    }

    public async Task<string> ReadMemoryFileAsync(string relativePath)
    {
        await EnsureDirectoryMemorySeededAsync();
        var path = ResolveMemoryPath(relativePath, allowTrash: true);
        if (!File.Exists(path))
            throw new FileNotFoundException("Memory file not found", relativePath);
        return await File.ReadAllTextAsync(path);
    }

    public async Task SaveMemoryFileAsync(string relativePath, string content)
    {
        var path = ResolveMemoryPath(relativePath, allowTrash: false);
        var dir = Path.GetDirectoryName(path);
        if (dir != null && !Directory.Exists(dir))
            Directory.CreateDirectory(dir);

        await _lock.WaitAsync();
        try
        {
            var existing = File.Exists(path)
                ? await File.ReadAllTextAsync(path)
                : string.Empty;
            var normalized = NormalizeMemoryFileContent(content);
            if (string.Equals(NormalizeForComparison(existing), NormalizeForComparison(normalized), StringComparison.Ordinal))
            {
                _logger.LogDebug("Memory file save skipped (no effective change): {Path}", ToMemoryRelativePath(path));
                return;
            }

            await WriteAllTextSafeAsync(path, normalized);
            _logger.LogInformation("Memory file saved: {Path}", ToMemoryRelativePath(path));
        }
        finally { _lock.Release(); }
    }

    public async Task SaveMemoryAsync(string section, string content)
    {
        await _lock.WaitAsync();
        try
        {
            await EnsureDirectoryMemorySeededAsync();
            var target = Path.Combine(_memoryDir, "core.md");

            var existing = File.Exists(target)
                ? await File.ReadAllTextAsync(target)
                : "# KaiChat Persistent Memory\n\n";

            var sectionBlock = ComposeSection(section, content);
            existing = UpsertSection(existing, section, sectionBlock);
            existing = RemoveDuplicateSections(existing);

            await WriteAllTextSafeAsync(target, existing);
            _logger.LogInformation("Memory updated: {Section}", section);
        }
        finally { _lock.Release(); }
    }

    public async Task SaveFullAsync(string content)
    {
        await _lock.WaitAsync();
        try
        {
            await EnsureDirectoryMemorySeededAsync();
            var target = Path.Combine(_memoryDir, "core.md");
            await WriteAllTextSafeAsync(target, RemoveDuplicateSections(content));
            _logger.LogInformation("Memory core file replaced from full-save");
        }
        finally { _lock.Release(); }
    }

    public async Task SoftDeleteMemoryFileAsync(string relativePath, string reason)
    {
        var path = ResolveMemoryPath(relativePath, allowTrash: false);
        if (!File.Exists(path))
            throw new FileNotFoundException("Memory file not found", relativePath);

        await _lock.WaitAsync();
        try
        {
            Directory.CreateDirectory(_trashDir);
            var stamp = DateTime.UtcNow.ToString("yyyyMMdd-HHmmss");
            var safeName = ToMemoryRelativePath(path)
                .Replace('/', '_')
                .Replace('\\', '_');
            var deletedPath = Path.Combine(_trashDir, $"{stamp}_{safeName}");
            File.Move(path, deletedPath);

            var tombstone = $"""
            # Deleted Memory File

            - Original path: `{ToMemoryRelativePath(path)}`
            - Deleted at UTC: `{DateTime.UtcNow:O}`
            - Reason: {reason.Trim()}
            - Stored as: `{ToMemoryRelativePath(deletedPath)}`

            This is a soft delete. Restore the stored file if this removal was a mistake.
            """;
            await File.WriteAllTextAsync(deletedPath + ".deleted.md", tombstone);
            _logger.LogInformation("Memory file soft-deleted: {Path}", relativePath);
        }
        finally { _lock.Release(); }
    }

    public async Task DeleteSectionAsync(string section)
    {
        await _lock.WaitAsync();
        try
        {
            await EnsureDirectoryMemorySeededAsync();
            var target = Path.Combine(_memoryDir, "core.md");

            if (!File.Exists(target)) return;
            var existing = await File.ReadAllTextAsync(target);
            var marker = $"## {section}";
            var start = existing.IndexOf(marker);
            if (start == -1) return;

            var end = existing.IndexOf("\n## ", start + marker.Length);
            if (end == -1) end = existing.Length - 1;
            else
            {
                // Find the start of the next header to keep the ## on the right line
                while (end > start && existing[end] != '\n') end--;
            }

            existing = existing[..start] + existing[(end + 1)..];
            await WriteAllTextSafeAsync(target, existing);
            _logger.LogInformation("Memory section deleted: {Section}", section);
        }
        finally { _lock.Release(); }
    }

    private async Task EnsureDirectoryMemorySeededAsync()
    {
        Directory.CreateDirectory(_memoryDir);
        await Task.CompletedTask;
    }

    private async Task<string> ReadMemoryPackAsync(IEnumerable<string> files)
    {
        var builder = new System.Text.StringBuilder();
        builder.AppendLine("# KaiChat Persistent Memory");
        builder.AppendLine();
        builder.AppendLine($"Memory is stored as a directory of markdown files under `{MemoryDirectoryRelativePath}`.");
        builder.AppendLine("Files under `_trash` and directories beginning with `_` are not loaded into the model context.");
        builder.AppendLine();

        foreach (var file in files)
        {
            var rel = ToMemoryRelativePath(file);
            var content = await File.ReadAllTextAsync(file);
            if (string.IsNullOrWhiteSpace(content))
                continue;

            builder.AppendLine($"## Memory File: {rel}");
            builder.AppendLine();
            builder.AppendLine(content.Trim());
            builder.AppendLine();
        }

        return builder.ToString().TrimEnd() + Environment.NewLine;
    }

    private IEnumerable<string> GetLoadableMemoryFiles()
    {
        if (!Directory.Exists(_memoryDir))
            return [];

        return Directory.GetFiles(_memoryDir, "*.md", SearchOption.AllDirectories)
            .Where(IsLoadableMemoryPath)
            .OrderBy(path => LoadOrder(ToMemoryRelativePath(path)))
            .ThenBy(ToMemoryRelativePath, StringComparer.OrdinalIgnoreCase);
    }

    private bool IsLoadableMemoryPath(string fullPath)
    {
        var rel = ToMemoryRelativePath(fullPath);
        if (rel.StartsWith("_", StringComparison.Ordinal))
            return false;

        var parts = rel.Split('/', StringSplitOptions.RemoveEmptyEntries);
        return parts.All(part => !part.StartsWith("_", StringComparison.Ordinal));
    }

    private bool IsTrashPath(string fullPath)
    {
        var rel = ToMemoryRelativePath(fullPath);
        return rel.StartsWith("_trash/", StringComparison.OrdinalIgnoreCase)
            || rel.Equals("_trash", StringComparison.OrdinalIgnoreCase);
    }

    private int LoadOrder(string relativePath)
    {
        var normalized = NormalizeRelativePath(relativePath);
        var idx = Array.FindIndex(PreferredLoadOrder, p => string.Equals(p, normalized, StringComparison.OrdinalIgnoreCase));
        return idx >= 0 ? idx : 1_000;
    }

    private string ResolveMemoryPath(string relativePath, bool allowTrash)
    {
        relativePath = NormalizeRelativePath(relativePath);
        if (string.IsNullOrWhiteSpace(relativePath))
            throw new ArgumentException("Memory path is required", nameof(relativePath));
        if (!relativePath.EndsWith(".md", StringComparison.OrdinalIgnoreCase))
            throw new ArgumentException("Memory files must be markdown files", nameof(relativePath));
        if (!allowTrash && relativePath.StartsWith("_trash/", StringComparison.OrdinalIgnoreCase))
            throw new ArgumentException("Write/delete operations cannot target _trash directly", nameof(relativePath));

        var full = Path.GetFullPath(Path.Combine(_memoryDir, relativePath));
        var root = Path.GetFullPath(_memoryDir);
        if (!full.StartsWith(root + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase)
            && !string.Equals(full, root, StringComparison.OrdinalIgnoreCase))
            throw new UnauthorizedAccessException($"Memory path is outside {MemoryDirectoryRelativePath}");

        return full;
    }

    private static string NormalizeRelativePath(string relativePath)
    {
        relativePath = relativePath.Replace('\\', '/').Trim().TrimStart('/');
        var parts = relativePath.Split('/', StringSplitOptions.RemoveEmptyEntries);
        if (parts.Any(part => part == "." || part == ".."))
            throw new ArgumentException("Memory path cannot contain . or .. segments", nameof(relativePath));
        return string.Join('/', parts);
    }

    private string ToMemoryRelativePath(string fullPath)
    {
        return Path.GetRelativePath(_memoryDir, fullPath).Replace('\\', '/');
    }

    private async Task WriteAllTextSafeAsync(string path, string content)
    {
        var dir = Path.GetDirectoryName(path);
        if (dir != null && !Directory.Exists(dir))
            Directory.CreateDirectory(dir);
        await File.WriteAllTextAsync(path, content.TrimEnd() + Environment.NewLine);
    }

    private static string NormalizeForComparison(string content)
    {
        return NormalizeMemoryFileContent(content).Replace('\r', '\n');
    }

    private static string NormalizeMemoryFileContent(string content)
    {
        var normalized = NormalizeLines(content);
        normalized = RemoveDuplicateSections(normalized);
        return normalized.TrimEnd();
    }

    private static string NormalizeLines(string content)
    {
        return content.Replace("\r\n", "\n").Replace('\r', '\n').TrimEnd();
    }

    private static string ComposeSection(string sectionTitle, string content)
    {
        var normalizedBody = NormalizeLines(content).Trim();
        return $"## {sectionTitle}\n\n{normalizedBody}\n";
    }

    private static string UpsertSection(string markdown, string sectionTitle, string sectionContent)
    {
        var sections = SplitSections(markdown);
        var normalizedTarget = NormalizeHeaderForMatch(sectionTitle);
        for (var i = sections.Count - 1; i >= 0; i--)
        {
            var section = sections[i];
            if (section.Heading is null)
                continue;
            if (!string.Equals(section.HeadingKey, normalizedTarget, StringComparison.OrdinalIgnoreCase))
                continue;

            sections.RemoveAt(i);
        }

        sections.Add(new MemorySection
        {
            Heading = sectionTitle,
            HeadingKey = normalizedTarget,
            Body = sectionContent
        });

        return BuildSections(sections);
    }

    private static string RemoveDuplicateSections(string markdown)
    {
        var sections = SplitSections(markdown);
        if (sections.Count == 0)
            return NormalizeLines(markdown);

        var seenHeadings = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
        var kept = new List<MemorySection>();
        for (var i = sections.Count - 1; i >= 0; i--)
        {
            var section = sections[i];
            var headingKey = section.HeadingKey ?? string.Empty;
            if (section.Heading == null)
            {
                kept.Add(section);
                continue;
            }

            if (!seenHeadings.Add(headingKey))
                continue;

            kept.Add(section);
        }

        kept.Reverse();
        return BuildSections(kept).Trim();
    }

    private static string BuildSections(IReadOnlyList<MemorySection> sections)
    {
        var builder = new StringBuilder();
        foreach (var section in sections)
        {
            if (section.Heading == null)
            {
                if (!string.IsNullOrWhiteSpace(section.Body))
                {
                    builder.AppendLine(section.Body);
                    builder.AppendLine();
                }
                continue;
            }

            builder.AppendLine(section.Body.TrimEnd('\n'));
            builder.AppendLine();
        }

        return builder.ToString().TrimEnd();
    }

    private static string NormalizeHeaderForMatch(string heading)
    {
        return heading.Trim().Replace("\u200b", "").ToLowerInvariant();
    }

    private static List<MemorySection> SplitSections(string markdown)
    {
        var normalized = NormalizeLines(markdown);
        var headingRegex = new Regex(@"(?m)^##\s+.*$");
        var matches = headingRegex.Matches(normalized);
        var sections = new List<MemorySection>();
        if (matches.Count == 0)
            return [new MemorySection { Heading = null, HeadingKey = null, Body = normalized }];

        var start = 0;
        for (var i = 0; i < matches.Count; i++)
        {
            if (start < matches[i].Index)
            {
                var preamble = normalized[start..matches[i].Index];
                if (!string.IsNullOrWhiteSpace(preamble))
                    sections.Add(new MemorySection { Heading = null, HeadingKey = null, Body = preamble.TrimEnd() });
            }

            var sectionStart = matches[i].Index;
            var next = i + 1 < matches.Count ? matches[i + 1].Index : normalized.Length;
            var rawSection = normalized[sectionStart..next].TrimEnd();
            var headingLine = rawSection.Split('\n', 2)[0];
            sections.Add(new MemorySection
            {
                Heading = headingLine.Length >= 3 ? headingLine.Substring(3).Trim() : headingLine.Trim(),
                HeadingKey = NormalizeHeaderForMatch(headingLine.Length >= 3 ? headingLine.Substring(3).Trim() : headingLine.Trim()),
                Body = rawSection
            });
            start = next;
        }

        return sections;
    }

    private sealed class MemorySection
    {
        public string? Heading { get; set; }
        public string? HeadingKey { get; set; }
        public string Body { get; set; } = string.Empty;
    }
}