Size: 11.0 KB Modified: 6/07/2026 9:12 PM
using System.Text;
using System.Text.RegularExpressions;
using KaiChat.Models;

namespace KaiChat.Services;

public interface IInstructionRewriteService
{
    Task<InstructionRewriteResult> RewriteAsync(string path, string instruction, CancellationToken ct = default);
}

public sealed record InstructionRewriteResult(
    bool Rewritten,
    string Path,
    string Message,
    bool IsError = false,
    string? Error = null,
    GitResult? Commit = null);

public sealed class InstructionRewriteService : IInstructionRewriteService
{
    private static readonly HashSet<string> AllowedInstructionFiles = new(StringComparer.OrdinalIgnoreCase)
    {
        "KaiChat/Kai.md",
        "Story Bible.md",
        "Writing Standards.md",
        "KaiChat/StoryBible/Overview.md",
        "KaiChat/StoryBible/Characters.md",
        "KaiChat/StoryBible/Anatomy.md",
        "KaiChat/StoryBible/Movement.md",
        "KaiChat/StoryBible/Communication.md",
        "KaiChat/StoryBible/Dissolution.md",
        "KaiChat/StoryBible/ScenesDispatchesInterludes.md",
        "KaiChat/StoryBible/IntimacyCanon.md",
        "KaiChat/StoryBible/LanguageAndLore.md",
        "KaiChat/StoryBible/Locations.md",
        "KaiChat/StoryBible/Inventory.md",
        "KaiChat/StoryBible/SupportingCharacters.md",
        "KaiChat/StoryBible/Relationship.md",
        "KaiChat/StoryBible/KnownGaps.md",
        "KaiChat/StoryBible/ArchiveStatus.md",
        "KaiChat/StoryBible/VersionHistory.md"
    };

    private readonly IChatService _chat;
    private readonly IFileService _files;
    private readonly IGitService _git;

    public InstructionRewriteService(IChatService chat, IFileService files, IGitService git)
    {
        _chat = chat;
        _files = files;
        _git = git;
    }

    public async Task<InstructionRewriteResult> RewriteAsync(string path, string instruction, CancellationToken ct = default)
    {
        if (string.IsNullOrWhiteSpace(path))
            return Error(path, "'Path' is required.");

        if (string.IsNullOrWhiteSpace(instruction))
            return Error(path, "'Instruction' is required.");

        if (!AllowedInstructionFiles.Contains(NormalizePath(path))
            && !IsAllowedPersonaFile(path))
            return Error(path, $"Instruction rewrites are limited to: {string.Join(", ", AllowedInstructionFiles.Order())}, and files under KaiChat/Personas/.");

        var currentContent = _files.ReadFile(path);
        if (currentContent is "(invalid path)" or "(file not found)")
            return Error(path, currentContent);

        var prompt = BuildRewritePrompt(path, instruction, currentContent);
        var rewritten = await _chat.ChatAsync(new List<ChatMessage>
        {
            new() { Role = "system", Content = prompt }
        }, ct);

        var candidate = InstructionRewriteGuard.CleanModelOutput(rewritten);
        var validation = InstructionRewriteGuard.Validate(currentContent, candidate, instruction);
        if (!validation.Valid)
        {
            return new InstructionRewriteResult(
                Rewritten: false,
                Path: path,
                Message: $"Rewrite rejected: {validation.Message}",
                IsError: true,
                Error: validation.Message);
        }

        if (!validation.HasMeaningfulChange)
        {
            return new InstructionRewriteResult(
                Rewritten: false,
                Path: path,
                Message: "Rewrite produced no meaningful content change; file was left untouched.");
        }

        await _files.WriteFileAsync(path, candidate);
        var commit = await _git.AutoCommitAsync($"rewrite {path}: {instruction}", path);
        if (!commit.Success)
        {
            return new InstructionRewriteResult(
                Rewritten: true,
                Path: path,
                Message: $"Rewrote {path}, but commit failed: {commit.Error}",
                IsError: true,
                Error: commit.Error,
                Commit: commit);
        }

        return new InstructionRewriteResult(
            Rewritten: true,
            Path: path,
            Message: $"Rewrote {path}: {instruction}\nAuto-commit: {CommitSummary(commit)}\nNo follow-up /api/git/commit call is needed for this rewrite.",
            Commit: commit);
    }

    private static string CommitSummary(GitResult commit)
    {
        return string.IsNullOrWhiteSpace(commit.Output)
            ? "completed"
            : commit.Output;
    }

    private static InstructionRewriteResult Error(string? path, string message)
    {
        return new InstructionRewriteResult(
            Rewritten: false,
            Path: path ?? "",
            Message: $"Error rewriting: {message}",
            IsError: true,
            Error: message);
    }

    private static string BuildRewritePrompt(string path, string instruction, string currentContent)
    {
        return $$"""
        Rewrite one markdown instruction file according to the requested change.

        Rules:
        - Return only the complete updated file content.
        - Preserve every existing heading and section unless the instruction explicitly says to remove it.
        - Preserve unrelated wording and ordering as much as possible.
        - If the requested section does not exist, add it without deleting existing sections.
        - Do not add markdown code fences, commentary, XML tags, or wrapper separators.
        - Do not add leading front matter unless the original file contains real YAML front matter with keys.

        File path: {{path}}

        Current file content:
        <current_file>
        {{currentContent}}
        </current_file>

        Requested change:
        <instruction>
        {{instruction}}
        </instruction>
        """;
    }

    private static string NormalizePath(string path)
    {
        return path.Replace('\\', '/').Trim().TrimStart('/');
    }
    private static bool IsAllowedPersonaFile(string path)
    {
        var normalized = NormalizePath(path);
        // Allow files under KaiChat/Personas/{persona}/ directory
        var segments = normalized.Split('/');
        return segments.Length >= 3
            && string.Equals(segments[0], "KaiChat", StringComparison.OrdinalIgnoreCase)
            && string.Equals(segments[1], "Personas", StringComparison.OrdinalIgnoreCase)
            && segments[2].Length > 0
            && !segments[2].StartsWith('_');
    }}

public sealed record InstructionRewriteValidation(
    bool Valid,
    bool HasMeaningfulChange,
    string Message);

public static class InstructionRewriteGuard
{
    private static readonly Regex WholeFenceRegex = new(
        @"\A\s*```(?:markdown|md|text)?\s*\r?\n(?<content>[\s\S]*?)\r?\n```\s*\z",
        RegexOptions.IgnoreCase | RegexOptions.Compiled);

    private static readonly Regex HeadingRegex = new(
        @"(?m)^(#{1,6})\s+(.+?)\s*$",
        RegexOptions.Compiled);

    private static readonly Regex DestructiveInstructionRegex = new(
        @"\b(delete|remove|drop|discard)\b|\b(replace|rewrite)\s+the\s+entire\s+file\b|\bfrom\s+scratch\b|\bonly\s+keep\b",
        RegexOptions.IgnoreCase | RegexOptions.Compiled);

    public static string CleanModelOutput(string text)
    {
        var cleaned = (text ?? "").Trim('\uFEFF', ' ', '\t', '\r', '\n');
        var fence = WholeFenceRegex.Match(cleaned);
        if (fence.Success)
            cleaned = fence.Groups["content"].Value.Trim('\uFEFF', ' ', '\t', '\r', '\n');

        cleaned = RemoveRedundantLeadingSeparators(cleaned);
        return NormalizeLineEndings(cleaned).TrimEnd() + Environment.NewLine;
    }

    public static InstructionRewriteValidation Validate(string currentContent, string candidateContent, string instruction)
    {
        if (string.IsNullOrWhiteSpace(candidateContent))
            return Invalid("AI returned empty content.");

        if (NormalizeMeaningful(candidateContent).Length <= 20)
            return Invalid("AI returned a too-short file.");

        var currentHeadings = ExtractHeadings(currentContent);
        var candidateHeadings = ExtractHeadings(candidateContent);
        if (!AllowsHeadingRemoval(instruction))
        {
            var missing = currentHeadings
                .Where(h => !candidateHeadings.Contains(h, StringComparer.OrdinalIgnoreCase))
                .ToList();

            if (missing.Count > 0)
                return Invalid($"rewrite removed existing heading(s): {string.Join(", ", missing)}.");
        }

        if (!AllowsHeadingRemoval(instruction) &&
            currentContent.Length >= 200 &&
            candidateContent.Length < currentContent.Length * 0.55)
        {
            return Invalid("rewrite removed too much content for a non-destructive instruction.");
        }

        var hasMeaningfulChange = !string.Equals(
            NormalizeMeaningful(currentContent),
            NormalizeMeaningful(candidateContent),
            StringComparison.Ordinal);

        return new InstructionRewriteValidation(true, hasMeaningfulChange, "OK");
    }

    private static InstructionRewriteValidation Invalid(string message)
    {
        return new InstructionRewriteValidation(false, false, message);
    }

    private static IReadOnlyList<string> ExtractHeadings(string content)
    {
        return HeadingRegex.Matches(NormalizeLineEndings(content ?? ""))
            .Select(m => $"{m.Groups[1].Value} {CollapseWhitespace(m.Groups[2].Value.Trim())}")
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToList();
    }

    private static bool AllowsHeadingRemoval(string instruction)
    {
        return DestructiveInstructionRegex.IsMatch(instruction ?? "");
    }

    private static string NormalizeMeaningful(string content)
    {
        var cleaned = RemoveRedundantLeadingSeparators(content ?? "");
        cleaned = NormalizeLineEndings(cleaned);
        var builder = new StringBuilder();
        foreach (var line in cleaned.Split('\n'))
            builder.AppendLine(line.TrimEnd());
        return builder.ToString().Trim();
    }

    private static string RemoveRedundantLeadingSeparators(string content)
    {
        var normalized = NormalizeLineEndings(content ?? "");
        var lines = normalized.Split('\n').ToList();
        var index = 0;
        var separatorCount = 0;

        while (index < lines.Count)
        {
            var trimmed = lines[index].Trim();
            if (trimmed.Length == 0)
            {
                index++;
                continue;
            }

            if (trimmed == "---")
            {
                separatorCount++;
                index++;
                continue;
            }

            break;
        }

        if (separatorCount == 0 || index >= lines.Count || !lines[index].TrimStart().StartsWith("#", StringComparison.Ordinal))
            return content ?? "";

        return string.Join('\n', lines.Skip(index));
    }

    private static string NormalizeLineEndings(string text)
    {
        return (text ?? "").Replace("\r\n", "\n").Replace('\r', '\n');
    }

    private static string CollapseWhitespace(string text)
    {
        return Regex.Replace(text, @"\s+", " ");
    }
}
Offline