ClaudeDo/src/ClaudeDo.Worker/Lifecycle/TaskMergeService.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Git;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Hub;
using ClaudeDo.Worker.State;
using Microsoft.EntityFrameworkCore;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Lifecycle;

public sealed record MergeResult(
    string Status,
    IReadOnlyList<string> ConflictFiles,
    string? ErrorMessage);

public sealed record MergeTargets(
    string DefaultBranch,
    IReadOnlyList<string> LocalBranches);

public sealed record MergePreviewResult(
    string Status,
    IReadOnlyList<string> ConflictFiles,
    int ChangedFileCount);

public sealed record MergeConflicts(
    string TaskId,
    IReadOnlyList<ConflictFileContent> Files);

public sealed record ConflictFileContent(
    string Path,
    string Ours,
    string Theirs,
    string? Base);

public sealed record ConflictDocuments(
    string TaskId,
    IReadOnlyList<ConflictDocumentContent> Files);

public sealed record ConflictDocumentContent(
    string Path,
    bool IsBinary,
    IReadOnlyList<MergeSegment> Segments);

public sealed class TaskMergeService
{
    public const string StatusMerged   = "merged";
    public const string StatusConflict = "conflict";
    public const string StatusBlocked  = "blocked";
    public const string StatusAborted  = "aborted";

    public const string PreviewClean       = "clean";
    public const string PreviewConflict    = "conflict";
    public const string PreviewUnavailable = "unavailable";

    private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;
    private readonly GitService _git;
    private readonly HubBroadcaster _broadcaster;
    private readonly ITaskStateService _state;
    private readonly ILogger<TaskMergeService> _logger;

    public TaskMergeService(
        IDbContextFactory<ClaudeDoDbContext> dbFactory,
        GitService git,
        HubBroadcaster broadcaster,
        ITaskStateService state,
        ILogger<TaskMergeService> logger)
    {
        _dbFactory = dbFactory;
        _git = git;
        _broadcaster = broadcaster;
        _state = state;
        _logger = logger;
    }

    private async Task<(TaskEntity Task, ListEntity List, WorktreeEntity? Worktree)> LoadMergeContextAsync(
        string taskId, CancellationToken ct)
    {
        using var ctx = _dbFactory.CreateDbContext();
        var task = await new TaskRepository(ctx).GetByIdAsync(taskId, ct)
            ?? throw new KeyNotFoundException($"Task '{taskId}' not found.");
        var list = await new ListRepository(ctx).GetByIdAsync(task.ListId, ct)
            ?? throw new InvalidOperationException("List not found.");
        var wt = await new WorktreeRepository(ctx).GetByTaskIdAsync(taskId, ct);
        return (task, list, wt);
    }

    private async Task MarkWorktreeMergedAsync(string taskId, CancellationToken ct)
    {
        using (var ctx = _dbFactory.CreateDbContext())
        {
            await new WorktreeRepository(ctx).SetStateAsync(taskId, WorktreeState.Merged, ct);
        }
        await _broadcaster.WorktreeUpdated(taskId);
    }

    private async Task ApproveIfWaitingForReviewAsync(TaskEntity task, CancellationToken ct)
    {
        // A merged worktree means the work is integrated, so the task must reach Done.
        // MarkWorktreeMergedAsync only flips the worktree state; transition the task
        // itself when it was still awaiting review (a Done task is already terminal).
        if (task.Status == TaskStatus.WaitingForReview)
            await _state.ApproveReviewAsync(task.Id, ct);
    }

    public async Task<MergeResult> MergeAsync(
        string taskId,
        string targetBranch,
        bool removeWorktree,
        string commitMessage,
        bool leaveConflictsInTree,
        CancellationToken ct)
    {
        var (task, list, wt) = await LoadMergeContextAsync(taskId, ct);

        if (task.Status == TaskStatus.Running)
            return Blocked("task is running");
        if (wt is null)
            return Blocked("task has no worktree");
        if (wt.State != WorktreeState.Active)
            return Blocked($"worktree state is {wt.State}");
        if (string.IsNullOrWhiteSpace(list.WorkingDir))
            return Blocked("list has no working directory");
        if (!await _git.IsGitRepoAsync(list.WorkingDir, ct))
            return Blocked("working directory is not a git repository");
        if (await _git.IsMidMergeAsync(list.WorkingDir, ct))
            return Blocked("target working directory is mid-merge");
        if (await _git.HasChangesAsync(list.WorkingDir, ct))
            return Blocked("target working tree has uncommitted changes");

        var currentBranch = await _git.GetCurrentBranchAsync(list.WorkingDir, ct);
        if (!string.Equals(currentBranch, targetBranch, StringComparison.Ordinal))
        {
            try { await _git.CheckoutBranchAsync(list.WorkingDir, targetBranch, ct); }
            catch (Exception ex) { return Blocked($"failed to switch target branch: {ex.Message}"); }
        }

        var (exitCode, stderr) = await _git.MergeNoFfAsync(list.WorkingDir, wt.BranchName, commitMessage, ct);
        if (exitCode != 0)
        {
            List<string> files;
            try { files = await _git.ListConflictedFilesAsync(list.WorkingDir, ct); }
            catch { files = new(); }

            if (leaveConflictsInTree && files.Count > 0)
            {
                return new MergeResult(StatusConflict, files, null);
            }

            // If abort fails the repo is left mid-merge; the caller must resolve manually.
            // Return Blocked (not conflict) so the UI does not offer a stale conflict list.
            try { await _git.MergeAbortAsync(list.WorkingDir, ct); }
            catch (Exception ex)
            {
                _logger.LogError(ex, "git merge --abort failed after conflict — repo is mid-merge");
                return Blocked($"merge conflict and abort failed: {ex.Message} — repo is mid-merge, resolve manually");
            }

            if (files.Count == 0)
            {
                // Non-conflict failure (e.g. unrelated histories).
                return new MergeResult(StatusBlocked, Array.Empty<string>(), $"merge failed: {stderr}");
            }

            return new MergeResult(StatusConflict, files, null);
        }

        string? cleanupWarning = null;
        if (removeWorktree)
        {
            try
            {
                await _git.WorktreeRemoveAsync(list.WorkingDir, wt.Path, force: false, ct);
                try { await _git.BranchDeleteAsync(list.WorkingDir, wt.BranchName, force: false, ct); }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex, "branch delete failed for {Branch}", wt.BranchName);
                    cleanupWarning = $"worktree removed, branch delete failed: {ex.Message}";
                }
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "worktree remove failed for {Path}", wt.Path);
                cleanupWarning = $"worktree remove failed: {ex.Message}";
            }
        }

        await MarkWorktreeMergedAsync(taskId, ct);
        await ApproveIfWaitingForReviewAsync(task, ct);

        _logger.LogInformation(
            "Merged task {TaskId} branch {Branch} into {Target} (remove worktree: {Remove})",
            taskId, wt.BranchName, targetBranch, removeWorktree);
        await _broadcaster.WorkerLog($"Merged \"{task.Title}\" into {targetBranch}", WorkerLogLevel.Success, DateTime.UtcNow);

        return new MergeResult(StatusMerged, Array.Empty<string>(), cleanupWarning);
    }

    public Task<MergeResult> MergeAsync(
        string taskId,
        string targetBranch,
        bool removeWorktree,
        string commitMessage,
        CancellationToken ct)
        => MergeAsync(taskId, targetBranch, removeWorktree, commitMessage, leaveConflictsInTree: false, ct);

    public async Task<MergeResult> ContinueMergeAsync(string taskId, CancellationToken ct)
    {
        var (task, list, wt) = await LoadMergeContextAsync(taskId, ct);

        if (wt is null) return Blocked("task has no worktree");
        if (wt.State != WorktreeState.Active) return Blocked($"worktree state is {wt.State}");
        if (string.IsNullOrWhiteSpace(list.WorkingDir)) return Blocked("list has no working directory");
        if (!await _git.IsMidMergeAsync(list.WorkingDir, ct))
            return Blocked("repo is not mid-merge");

        await _git.AddAllAsync(list.WorkingDir, ct);

        var remaining = await _git.ListConflictedFilesAsync(list.WorkingDir, ct);
        if (remaining.Count > 0)
            return new MergeResult(StatusConflict, remaining, "conflicts not fully resolved");

        try { await _git.CommitAsync(list.WorkingDir, $"Merge branch '{wt.BranchName}'", ct); }
        catch (Exception ex) { return Blocked($"commit failed: {ex.Message}"); }

        await MarkWorktreeMergedAsync(taskId, ct);
        await ApproveIfWaitingForReviewAsync(task, ct);
        _logger.LogInformation("Continued merge of task {TaskId} branch {Branch}", taskId, wt.BranchName);

        return new MergeResult(StatusMerged, Array.Empty<string>(), null);
    }

    public async Task<MergeResult> AbortMergeAsync(string taskId, CancellationToken ct)
    {
        var (_, list, wt) = await LoadMergeContextAsync(taskId, ct);

        if (wt is null) return Blocked("task has no worktree");
        if (wt.State != WorktreeState.Active) return Blocked($"worktree state is {wt.State}");
        if (string.IsNullOrWhiteSpace(list.WorkingDir)) return Blocked("list has no working directory");
        if (!await _git.IsMidMergeAsync(list.WorkingDir, ct))
            return Blocked("repo is not mid-merge");

        try { await _git.MergeAbortAsync(list.WorkingDir, ct); }
        catch (Exception ex) { return Blocked($"abort failed: {ex.Message}"); }
        _logger.LogInformation("Aborted merge of task {TaskId}", taskId);

        return new MergeResult(StatusAborted, Array.Empty<string>(), null);
    }

    public async Task<MergeConflicts> GetConflictsAsync(string taskId, CancellationToken ct)
    {
        var (_, list, _) = await LoadMergeContextAsync(taskId, ct);
        if (string.IsNullOrWhiteSpace(list.WorkingDir))
            throw new InvalidOperationException("list has no working directory");

        var files = await _git.ListConflictedFilesAsync(list.WorkingDir, ct);
        var result = new List<ConflictFileContent>(files.Count);
        foreach (var path in files)
        {
            var ours   = await _git.ShowStageAsync(list.WorkingDir, 2, path, ct) ?? "";
            var theirs = await _git.ShowStageAsync(list.WorkingDir, 3, path, ct) ?? "";
            var @base  = await _git.ShowStageAsync(list.WorkingDir, 1, path, ct);
            result.Add(new ConflictFileContent(path, ours, theirs, @base));
        }
        return new MergeConflicts(taskId, result);
    }

    /// <summary>
    /// Reads each conflicted working-tree file and parses its conflict markers into line-level
    /// segments (with the diff3 merge base when present). Binary files are flagged and skipped.
    /// </summary>
    public async Task<ConflictDocuments> GetConflictDocumentsAsync(string taskId, CancellationToken ct)
    {
        var (_, list, _) = await LoadMergeContextAsync(taskId, ct);
        if (string.IsNullOrWhiteSpace(list.WorkingDir))
            throw new InvalidOperationException("list has no working directory");

        var files = await _git.ListConflictedFilesAsync(list.WorkingDir, ct);
        var result = new List<ConflictDocumentContent>(files.Count);
        foreach (var path in files)
        {
            var full = Path.Combine(list.WorkingDir, path.Replace('/', Path.DirectorySeparatorChar));
            string text;
            try { text = await File.ReadAllTextAsync(full, ct); }
            catch { text = ""; }

            if (LooksBinary(text))
            {
                result.Add(new ConflictDocumentContent(path, true, Array.Empty<MergeSegment>()));
                continue;
            }

            result.Add(new ConflictDocumentContent(path, false, ConflictMarkerParser.Parse(text)));
        }
        return new ConflictDocuments(taskId, result);
    }

    // A NUL byte in the head of the file is the conventional binary sniff.
    private static bool LooksBinary(string text)
    {
        var n = Math.Min(text.Length, 8000);
        for (var i = 0; i < n; i++)
            if (text[i] == '\0') return true;
        return false;
    }

    public async Task WriteResolutionAsync(string taskId, string path, string content, CancellationToken ct)
    {
        var (_, list, _) = await LoadMergeContextAsync(taskId, ct);
        if (string.IsNullOrWhiteSpace(list.WorkingDir))
            throw new InvalidOperationException("list has no working directory");

        var full = Path.Combine(list.WorkingDir, path.Replace('/', Path.DirectorySeparatorChar));
        await File.WriteAllTextAsync(full, content, ct);
        await _git.AddPathAsync(list.WorkingDir, path, ct);
    }

    public async Task<MergeTargets> GetTargetsAsync(string taskId, CancellationToken ct)
    {
        var (_, list, _) = await LoadMergeContextAsync(taskId, ct);

        if (string.IsNullOrWhiteSpace(list.WorkingDir))
            return new MergeTargets("", Array.Empty<string>());

        var current  = await _git.GetCurrentBranchAsync(list.WorkingDir, ct);
        var branches = await _git.ListLocalBranchesAsync(list.WorkingDir, ct);
        return new MergeTargets(current, branches);
    }

    public async Task<MergePreviewResult> PreviewAsync(string taskId, string targetBranch, CancellationToken ct)
    {
        var (_, list, wt) = await LoadMergeContextAsync(taskId, ct);

        if (wt is null || wt.State != WorktreeState.Active)
            return new MergePreviewResult(PreviewUnavailable, Array.Empty<string>(), 0);
        if (string.IsNullOrWhiteSpace(list.WorkingDir) || !await _git.IsGitRepoAsync(list.WorkingDir, ct))
            return new MergePreviewResult(PreviewUnavailable, Array.Empty<string>(), 0);

        var target = string.IsNullOrWhiteSpace(targetBranch)
            ? await _git.GetCurrentBranchAsync(list.WorkingDir, ct)
            : targetBranch;

        var preview = await _git.PreviewMergeAsync(list.WorkingDir, target, wt.BranchName, ct);
        if (!preview.Supported)
            return new MergePreviewResult(PreviewUnavailable, Array.Empty<string>(), 0);
        if (!preview.Clean)
            return new MergePreviewResult(PreviewConflict, preview.ConflictFiles, 0);

        var count = await _git.CountChangedFilesAsync(list.WorkingDir, target, wt.BranchName, ct);
        return new MergePreviewResult(PreviewClean, Array.Empty<string>(), count);
    }

    public async Task<MergeResult> ApproveAndMergeAsync(string taskId, string targetBranch, CancellationToken ct)
    {
        var (task, list, wt) = await LoadMergeContextAsync(taskId, ct);

        if (task.Status != TaskStatus.WaitingForReview)
            return Blocked("task is not waiting for review");

        if (wt is null || wt.State != WorktreeState.Active)
        {
            var done = await _state.ApproveReviewAsync(taskId, ct);
            return done.Ok
                ? new MergeResult(StatusMerged, Array.Empty<string>(), null)
                : Blocked(done.Reason ?? "approve failed");
        }

        if (string.IsNullOrWhiteSpace(list.WorkingDir))
            return Blocked("list has no working directory");

        var target = string.IsNullOrWhiteSpace(targetBranch)
            ? await _git.GetCurrentBranchAsync(list.WorkingDir, ct)
            : targetBranch;

        // MergeAsync transitions the task WaitingForReview -> Done on a successful merge.
        // Remove the worktree on approve (matching the unit-merge path) so merged
        // worktrees don't pile up; the merge commit on the target branch is the record.
        return await MergeAsync(taskId, target, removeWorktree: true, $"Merge {wt.BranchName}", ct);
    }

    private static MergeResult Blocked(string reason) =>
        new(StatusBlocked, Array.Empty<string>(), reason);
}